fix: implement jmt (#7036)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added support for reading versioned values, allowing retrieval of the
latest value for a given key.
- **Bug Fixes**
- Improved value lookup logic to ensure the most recent data is returned
when querying by key.
- **Chores**
  - Enhanced logging for output fetching and spend status determination.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: SW van Heerden <swvheerden@gmail.com>

Author: stringhandler

base_layer/core/src/chain_storage/blockchain_database.rs

@@ -477,6 +477,12 @@ where B: BlockchainBackend
         for hash in hashes {
             let output = db.fetch_output(&hash)?;
 
+            trace!(
+                target: LOG_TARGET,
+                "fetch_outputs_with_spend_status_at_tip: hash: {}, output: {:?}",
+                hash.to_hex(),
+                output
+            );
             if let Some(mined_info) = output {
                 let smt_key = KeyHash(
                     mined_info
@@ -490,7 +496,13 @@ where B: BlockchainBackend
                 let spent = smt
                     .get(smt_key, tip)
                     .map_err(ChainStorageError::JellyfishMerkleTreeError)?
-                    .is_some();
+                    .is_none();
+                trace!(
+                    target: LOG_TARGET,
+                    "fetch_outputs_with_spend_status_at_tip: smt_key: {:?}, spent: {}",
+                    smt_key,
+                    spent
+                );
                 result.push(Some((mined_info.output, spent)));
             } else {
                 result.push(None);

base_layer/core/src/chain_storage/lmdb_db/lmdb_db.rs

@@ -989,7 +989,7 @@ impl LMDBDatabase {
         self.delete_block_inputs_outputs(write_txn, block_hash.as_slice())?;
 
         let new_tip_header = self.fetch_chain_header_by_height(prev_height)?;
-        let reader = LmdbTreeReader::new(write_txn, self.jmt_node_data.clone());
+        let reader = LmdbTreeReader::new(write_txn, self.jmt_node_data.clone(), self.jmt_unique_key_data.clone());
         let jmt = JellyfishMerkleTree::<_, SmtHasher>::new(&reader);
 
         let root = jmt
@@ -1252,7 +1252,7 @@ impl LMDBDatabase {
         // smt_writer: &mut LmdbTreeWriter,
     ) -> Result<(), ChainStorageError> {
         // let smt_reader = LmdbTreeReader::new();
-        let smt_reader = LmdbTreeReader::new(txn, self.jmt_node_data.clone());
+        let smt_reader = LmdbTreeReader::new(txn, self.jmt_node_data.clone(), self.jmt_unique_key_data.clone());
         let output_smt = JellyfishMerkleTree::<_, SmtHasher>::new(&smt_reader);
         if self.fetch_block_accumulated_data(txn, header.height + 1)?.is_some() {
             return Err(ChainStorageError::InvalidOperation(format!(
@@ -1889,7 +1889,8 @@ fn acquire_exclusive_file_lock(db_path: &Path) -> Result<File, ChainStorageError
 impl BlockchainBackend for LMDBDatabase {
     fn create_smt_reader(&self) -> Result<OwnedLmdbTreeReader<'_>, ChainStorageError> {
         let read_tx = self.read_transaction()?;
-        let smt_reader = OwnedLmdbTreeReader::new(read_tx, self.jmt_node_data.clone());
+        let smt_reader =
+            OwnedLmdbTreeReader::new(read_tx, self.jmt_node_data.clone(), self.jmt_unique_key_data.clone());
 
         Ok(smt_reader)
     }

base_layer/core/src/chain_storage/lmdb_db/lmdb_tree_reader.rs

@@ -25,20 +25,29 @@ use std::ops::Deref;
 use borsh::BorshSerialize;
 use jmt::storage::TreeReader;
 use lmdb_zero::{ConstTransaction, ReadTransaction};
+use log::warn;
 use tari_storage::lmdb_store::DatabaseRef;
 
-use crate::chain_storage::lmdb_db::lmdb::lmdb_get;
+use crate::chain_storage::lmdb_db::lmdb::{lmdb_fetch_matching_after, lmdb_get};
+
+pub const LOG_TARGET: &str = "c::cs::lmdb_db::lmdb_db";
 
 pub struct LmdbTreeReader<'a> {
     txn: &'a ConstTransaction<'a>,
     node_db: DatabaseRef,
+    unique_key_db: DatabaseRef,
 }
 
 impl<'a> LmdbTreeReader<'a> {
-    pub fn new<T: Deref<Target = ConstTransaction<'a>>>(txn: &'a T, node_db: DatabaseRef) -> Self {
+    pub fn new<T: Deref<Target = ConstTransaction<'a>>>(
+        txn: &'a T,
+        node_db: DatabaseRef,
+        unique_key_db: DatabaseRef,
+    ) -> Self {
         Self {
             txn: txn.deref(),
             node_db,
+            unique_key_db,
         }
     }
 }
@@ -54,12 +63,24 @@ impl TreeReader for LmdbTreeReader<'_> {
 
     fn get_value_option(
         &self,
-        _max_version: jmt::Version,
-        _key_hash: jmt::KeyHash,
+        max_version: jmt::Version,
+        key_hash: jmt::KeyHash,
     ) -> anyhow::Result<Option<jmt::OwnedValue>> {
-        todo!()
-        // TODO: implement after saving
-        // Ok(None)
+        // see if there are any values already.
+        let existing_values: Vec<(Vec<u8>, Option<Vec<u8>>)> =
+            lmdb_fetch_matching_after(self.txn, &self.unique_key_db, &key_hash.0)?;
+        let mut existing_history = vec![];
+        for (key, x) in existing_values {
+            let version = u64::from_be_bytes(key[32..].try_into().unwrap());
+            existing_history.push((version, x));
+            warn!(target: LOG_TARGET, "found version {} for key {:?}", version, key);
+        }
+        // sort by version
+        existing_history.sort_by(|a, b| a.0.cmp(&b.0));
+
+        let latest_value = existing_history.last().and_then(|x| x.1.clone());
+
+        Ok(latest_value)
     }
 
     fn get_rightmost_leaf(&self) -> anyhow::Result<Option<(jmt::storage::NodeKey, jmt::storage::LeafNode)>> {
@@ -71,26 +92,32 @@ impl TreeReader for LmdbTreeReader<'_> {
 pub struct OwnedLmdbTreeReader<'a> {
     txn: ReadTransaction<'a>,
     node_db: DatabaseRef,
+    unique_key_db: DatabaseRef,
 }
 
 impl<'a> OwnedLmdbTreeReader<'a> {
-    pub fn new(txn: ReadTransaction<'a>, node_db: DatabaseRef) -> Self {
-        Self { txn, node_db }
+    pub fn new(txn: ReadTransaction<'a>, node_db: DatabaseRef, unique_key_db: DatabaseRef) -> Self {
+        Self {
+            txn,
+            node_db,
+            unique_key_db,
+        }
     }
 }
 
 impl TreeReader for OwnedLmdbTreeReader<'_> {
     fn get_node_option(&self, node_key: &jmt::storage::NodeKey) -> anyhow::Result<Option<jmt::storage::Node>> {
-        let inner = LmdbTreeReader::new(&self.txn, self.node_db.clone());
+        let inner = LmdbTreeReader::new(&self.txn, self.node_db.clone(), self.unique_key_db.clone());
         inner.get_node_option(node_key)
     }
 
     fn get_value_option(
         &self,
-        _max_version: jmt::Version,
-        _key_hash: jmt::KeyHash,
+        max_version: jmt::Version,
+        key_hash: jmt::KeyHash,
     ) -> anyhow::Result<Option<jmt::OwnedValue>> {
-        todo!()
+        let inner = LmdbTreeReader::new(&self.txn, self.node_db.clone(), self.unique_key_db.clone());
+        inner.get_value_option(max_version, key_hash)
     }
 
     fn get_rightmost_leaf(&self) -> anyhow::Result<Option<(jmt::storage::NodeKey, jmt::storage::LeafNode)>> {