From 85c80e1bd384e21135faf40bbfcb4558bb252e00 Mon Sep 17 00:00:00 2001
From: skoupidi <skoupidi@systemli.org>
Date: Mon, 1 Apr 2024 21:48:14 +0300
Subject: [PATCH] blockchain: store txs locations using the new tree

---
 src/blockchain/mod.rs      | 26 ++++++++++++++++++--------
 src/blockchain/tx_store.rs | 18 ++++++++++++++++--
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/src/blockchain/mod.rs b/src/blockchain/mod.rs
index c2672193c..00c96d7dd 100644
--- a/src/blockchain/mod.rs
+++ b/src/blockchain/mod.rs
@@ -76,11 +76,6 @@ impl Blockchain {
         let mut trees = vec![];
         let mut batches = vec![];
 
-        // Store transactions
-        let (txs_batch, _) = self.transactions.insert_batch(&block.txs)?;
-        trees.push(self.transactions.main.clone());
-        batches.push(txs_batch);
-
         // Store header
         let (headers_batch, _) = self.headers.insert_batch(&[block.header.clone()])?;
         trees.push(self.headers.0.clone());
@@ -100,6 +95,17 @@ impl Blockchain {
         trees.push(self.blocks.order.clone());
         batches.push(blocks_order_batch);
 
+        // Store transactions
+        let (txs_batch, txs_hashes) = self.transactions.insert_batch(&block.txs)?;
+        trees.push(self.transactions.main.clone());
+        batches.push(txs_batch);
+
+        // Store transactions_locations
+        let txs_locations_batch =
+            self.transactions.insert_batch_location(&txs_hashes, block.header.height)?;
+        trees.push(self.transactions.location.clone());
+        batches.push(txs_locations_batch);
+
         // Perform an atomic transaction over the trees and apply the batches.
         self.atomic_write(&trees, &batches)?;
 
@@ -381,20 +387,24 @@ impl BlockchainOverlay {
     /// Since we are adding to the overlay, we don't need to exeucte
     /// the writes atomically.
     pub fn add_block(&self, block: &BlockInfo) -> Result<blake3::Hash> {
-        // Store transactions
-        self.transactions.insert(&block.txs)?;
-
         // Store header
         self.headers.insert(&[block.header.clone()])?;
 
         // Store block
         let blk: Block = Block::from_block_info(block)?;
+        let txs_hashes = blk.txs.clone();
         let block_hash = self.blocks.insert(&[blk])?[0];
         let block_hash_vec = [block_hash];
 
         // Store block order
         self.blocks.insert_order(&[block.header.height], &block_hash_vec)?;
 
+        // Store transactions
+        self.transactions.insert(&block.txs)?;
+
+        // Store transactions locations
+        self.transactions.insert_location(&txs_hashes, block.header.height)?;
+
         Ok(block_hash)
     }
 
diff --git a/src/blockchain/tx_store.rs b/src/blockchain/tx_store.rs
index 966d63b18..653375add 100644
--- a/src/blockchain/tx_store.rs
+++ b/src/blockchain/tx_store.rs
@@ -116,8 +116,8 @@ impl TxStore {
 
     /// Generate the sled batch corresponding to an insert to the location tree,
     /// so caller can handle the write operation.
-    /// The tuple is built using the index of each location in the slice,
-    /// along with the provided block height
+    /// The location tuple is built using the index of each transaction has in
+    /// the slice, along with the provided block height
     pub fn insert_batch_location(
         &self,
         txs_hashes: &[blake3::Hash],
@@ -405,6 +405,20 @@ impl TxStoreOverlay {
         Ok(ret)
     }
 
+    /// Insert a slice of [`blake3::Hash`] into the overlay's location tree.
+    /// The location tuple is built using the index of each transaction hash
+    /// in the slice, along with the provided block height
+    pub fn insert_location(&self, txs_hashes: &[blake3::Hash], block_height: u64) -> Result<()> {
+        let mut lock = self.0.lock().unwrap();
+
+        for (index, tx_hash) in txs_hashes.iter().enumerate() {
+            let serialized = serialize(&(block_height, index as u64));
+            lock.insert(SLED_TX_LOCATION_TREE, tx_hash.as_bytes(), &serialized)?;
+        }
+
+        Ok(())
+    }
+
     /// Fetch given tx hashes from the overlay's main tree.
     /// The resulting vector contains `Option`, which is `Some` if the tx
     /// was found in the overlay, and otherwise it is `None`, if it has not.