merkle: skip the leaf fetch when the read value identifies the leaf

nomt/src/lib.rs

@@ -130,14 +130,17 @@ impl KeyReadWrite {
         }
     }
 
+    /// Hashes the values. The read value hash lets the merkle update recognize the key's own leaf
+    /// without fetching it.
     fn to_compact<T: HashAlgorithm>(&self) -> crate::merkle::KeyReadWrite {
         let hash = |v: &Value| T::hash_value(v);
         match self {
-            KeyReadWrite::Read(_) => crate::merkle::KeyReadWrite::Read,
+            KeyReadWrite::Read(read) => crate::merkle::KeyReadWrite::Read(read.as_ref().map(hash)),
             KeyReadWrite::Write(val) => crate::merkle::KeyReadWrite::Write(val.as_ref().map(hash)),
-            KeyReadWrite::ReadThenWrite(_, val) => {
-                crate::merkle::KeyReadWrite::ReadThenWrite(val.as_ref().map(hash))
-            }
+            KeyReadWrite::ReadThenWrite(read, val) => crate::merkle::KeyReadWrite::ReadThenWrite(
+                read.as_ref().map(hash),
+                val.as_ref().map(hash),
+            ),
         }
     }
 }
@@ -684,15 +687,13 @@ fn compact_actuals<T: HashAlgorithm>(
     if max_threads == 1 || actuals.len() < 2 * MIN_HASHES_PER_THREAD {
         return compact(actuals);
     }
-    let hashes = actuals
+    let hashes: usize = actuals
         .iter()
-        .filter(|(_, read_write)| {
-            matches!(
-                read_write,
-                KeyReadWrite::Write(Some(_)) | KeyReadWrite::ReadThenWrite(_, Some(_))
-            )
+        .map(|(_, read_write)| match read_write {
+            KeyReadWrite::Read(v) | KeyReadWrite::Write(v) => v.is_some() as usize,
+            KeyReadWrite::ReadThenWrite(r, w) => r.is_some() as usize + w.is_some() as usize,
         })
-        .count();
+        .sum();
     let threads = max_threads.min(hashes / MIN_HASHES_PER_THREAD);
     if threads < 2 {
         return compact(actuals);
@@ -1060,19 +1061,24 @@ mod tests {
             })
             .collect();
 
-        for len in [0, 1, 16_383, 16_384, 49_151, 49_152, 49_153, 131_073] {
+        for len in [0, 1, 16_383, 16_384, 24_575, 24_576, 24_577, 131_073] {
             let sequential = compact_actuals::<Blake3Hasher>(&actuals[..len], 1);
             for threads in [2, 16] {
                 let parallel = compact_actuals::<Blake3Hasher>(&actuals[..len], threads);
                 assert_eq!(parallel.len(), sequential.len());
                 for ((path, op), (expected_path, expected_op)) in parallel.iter().zip(&sequential) {
                     assert_eq!(path, expected_path);
                     match (op, expected_op) {
-                        (Compact::Read, Compact::Read) => {}
-                        (Compact::Write(value), Compact::Write(expected))
-                        | (Compact::ReadThenWrite(value), Compact::ReadThenWrite(expected)) => {
+                        (Compact::Read(value), Compact::Read(expected))
+                        | (Compact::Write(value), Compact::Write(expected)) => {
                             assert_eq!(value, expected);
                         }
+                        (
+                            Compact::ReadThenWrite(read, write),
+                            Compact::ReadThenWrite(expected_read, expected_write),
+                        ) => {
+                            assert_eq!((read, write), (expected_read, expected_write));
+                        }
                         _ => panic!("operation changed: {op:?} != {expected_op:?}"),
                     }
                 }

nomt/src/merkle/mod.rs

@@ -142,33 +142,42 @@ pub enum BucketInfo {
 /// Whether a key was read or written.
 #[derive(Debug, Clone)]
 pub enum KeyReadWrite {
-    /// The key was read.
-    Read,
+    /// The key was read. Contains the read value hash or `None` if absent.
+    Read(Option<ValueHash>),
     /// The key was written. Contains the written value or `None` if deleted.
     Write(Option<ValueHash>),
-    /// The key was both read and written. Contains the written value of `None` if deleted.
-    ReadThenWrite(Option<ValueHash>),
+    /// The key was both read and written. Contains the read value hash (or `None` if absent) and
+    /// the written value hash (or `None` if deleted).
+    ReadThenWrite(Option<ValueHash>, Option<ValueHash>),
 }
 
 impl KeyReadWrite {
     fn is_write(&self) -> bool {
         match self {
-            KeyReadWrite::Read => false,
-            KeyReadWrite::Write(_) | KeyReadWrite::ReadThenWrite(_) => true,
+            KeyReadWrite::Read(_) => false,
+            KeyReadWrite::Write(_) | KeyReadWrite::ReadThenWrite(_, _) => true,
         }
     }
 
     fn is_read(&self) -> bool {
         match self {
-            KeyReadWrite::Read | KeyReadWrite::ReadThenWrite(_) => true,
+            KeyReadWrite::Read(_) | KeyReadWrite::ReadThenWrite(_, _) => true,
             KeyReadWrite::Write(_) => false,
         }
     }
 
+    /// The hash of the value the key was read with, if it was read and present.
+    fn read_value_hash(&self) -> Option<ValueHash> {
+        match self {
+            KeyReadWrite::Read(v) | KeyReadWrite::ReadThenWrite(v, _) => *v,
+            KeyReadWrite::Write(_) => None,
+        }
+    }
+
     fn written_value(&self) -> Option<Option<ValueHash>> {
         match self {
-            KeyReadWrite::Read => None,
-            KeyReadWrite::Write(v) | KeyReadWrite::ReadThenWrite(v) => {
+            KeyReadWrite::Read(_) => None,
+            KeyReadWrite::Write(v) | KeyReadWrite::ReadThenWrite(_, v) => {
                 Some(v.as_ref().map(|vh| *vh))
             }
         }
@@ -339,7 +348,7 @@ impl Updater {
             true, // collect witness
         );
 
-        seeker.push(key_path);
+        seeker.push(key_path, None);
         let mut page_set = PageSet::new(self.page_pool.clone(), None);
 
         // Blocking I/O loop.

nomt/src/merkle/seek.rs

@@ -38,6 +38,9 @@ const MAX_INFLIGHT: usize = 1024;
 
 struct SeekRequest {
     key: KeyPath,
+    /// Hash of the value the caller read for `key`, if it exists. A terminal leaf whose hash
+    /// matches `(key, known_value_hash)` is this key's leaf, so it needs no beatree fetch.
+    known_value_hash: Option<ValueHash>,
     position: TriePosition,
     page_id: Option<PageId>,
     siblings: Vec<Node>,
@@ -50,18 +53,27 @@ impl SeekRequest {
         read_transaction: &BeatreeReadTx,
         overlay: &LiveOverlay,
         key: KeyPath,
+        known_value_hash: Option<ValueHash>,
         root: Node,
     ) -> SeekRequest {
         let state = if trie::is_terminator::<H>(&root) {
             RequestState::Completed(None)
         } else if trie::is_leaf::<H>(&root) {
-            RequestState::begin_leaf_fetch::<H>(read_transaction, overlay, &TriePosition::new())
+            RequestState::resolve_leaf::<H>(
+                read_transaction,
+                overlay,
+                &TriePosition::new(),
+                &root,
+                key,
+                known_value_hash,
+            )
         } else {
             RequestState::Seeking
         };
 
         let mut request = SeekRequest {
             key,
+            known_value_hash,
             position: TriePosition::new(),
             page_id: None,
             siblings: Vec::new(),
@@ -144,8 +156,14 @@ impl SeekRequest {
             }
 
             if trie::is_leaf::<H>(&cur_node) {
-                self.state =
-                    RequestState::begin_leaf_fetch::<H>(read_transaction, overlay, &self.position);
+                self.state = RequestState::resolve_leaf::<H>(
+                    read_transaction,
+                    overlay,
+                    &self.position,
+                    &cur_node,
+                    self.key,
+                    self.known_value_hash,
+                );
                 if let RequestState::FetchingLeaf { .. } = self.state {
                     self.continue_leaf_fetch::<H>(None);
                 }
@@ -412,6 +430,28 @@ enum RequestState {
 }
 
 impl RequestState {
+    /// Resolves the leaf node `leaf_node` at `pos`: without I/O if it is the leaf of `key` with the
+    /// known value hash, otherwise by fetching the leaf data from the beatree.
+    fn resolve_leaf<H: HashAlgorithm>(
+        read_transaction: &BeatreeReadTx,
+        overlay: &LiveOverlay,
+        pos: &TriePosition,
+        leaf_node: &Node,
+        key: KeyPath,
+        known_value_hash: Option<ValueHash>,
+    ) -> Self {
+        if let Some(value_hash) = known_value_hash {
+            let leaf = trie::LeafData {
+                key_path: key,
+                value_hash,
+            };
+            if H::hash_leaf(&leaf) == *leaf_node {
+                return RequestState::Completed(Some(leaf));
+            }
+        }
+        Self::begin_leaf_fetch::<H>(read_transaction, overlay, pos)
+    }
+
     fn begin_leaf_fetch<H: HashAlgorithm>(
         read_transaction: &BeatreeReadTx,
         overlay: &LiveOverlay,
@@ -634,12 +674,15 @@ impl<H: HashAlgorithm> Seeker<H> {
     }
 
     /// Push a request for key path.
-    pub fn push(&mut self, key: KeyPath) {
+    ///
+    /// `known_value_hash` is the hash of the value the caller read for `key`, if any.
+    pub fn push(&mut self, key: KeyPath, known_value_hash: Option<ValueHash>) {
         let request_index = self.processed + self.requests.len();
         self.requests.push_back(SeekRequest::new::<H>(
             &self.beatree_read_transaction,
             &self.overlay,
             key,
+            known_value_hash,
             self.root,
         ));
         self.idle_requests.push_back(request_index);

nomt/src/merkle/worker.rs

@@ -171,7 +171,7 @@ fn warm_up_phase<H: HashAlgorithm>(
                     Err(_) => panic!("Warm-Up worker, unexpected failure of the warmup channel"),
                 };
 
-                seeker.push(warm_up_command.key_path);
+                seeker.push(warm_up_command.key_path, None);
             } else if index == page_idx {
                 seeker.try_recv_page(&mut page_set)?;
             } else {
@@ -510,7 +510,8 @@ impl<H: HashAlgorithm> RangeUpdater<H> {
                         break;
                     }
                 } else {
-                    seeker.push(self.shared.read_write[next_push].0);
+                    let (key, read_write) = &self.shared.read_write[next_push];
+                    seeker.push(*key, read_write.read_value_hash());
                     seeker.submit_all(page_set);
                 }
             }
@@ -538,10 +539,10 @@ fn subtrie_ops(read_write: &[(KeyPath, KeyReadWrite)]) -> Vec<(KeyPath, Option<V
     read_write
         .iter()
         .filter_map(|(key, read_write)| match read_write {
-            KeyReadWrite::Write(val) | KeyReadWrite::ReadThenWrite(val) => {
+            KeyReadWrite::Write(val) | KeyReadWrite::ReadThenWrite(_, val) => {
                 Some((key.clone(), val.clone()))
             }
-            KeyReadWrite::Read => None,
+            KeyReadWrite::Read(_) => None,
         })
         .collect::<Vec<_>>()
 }