feat: cacheless iterator for hamt

hanabi1224 · hanabi1224 · commit bbdb34f6ee06 · 2025-09-01T23:44:51.000+08:00
diff --git a/ipld/hamt/src/hamt.rs b/ipld/hamt/src/hamt.rs
@@ -13,7 +13,7 @@ use multihash_codetable::Code;
 use serde::de::DeserializeOwned;
 use serde::{Serialize, Serializer};
 
-use crate::iter::IterImpl;
+use crate::iter::{IterImpl, IterItem};
 use crate::node::Node;
 use crate::pointer::version::Version;
 use crate::{Config, Error, Hash, HashAlgorithm, Sha256, pointer::version};
@@ -372,12 +372,13 @@ where
     #[inline]
     pub fn for_each<F>(&self, mut f: F) -> Result<(), Error>
     where
-        V: DeserializeOwned,
+        K: Clone,
+        V: DeserializeOwned + Clone,
         F: FnMut(&K, &V) -> anyhow::Result<()>,
     {
         for res in self {
             let (k, v) = res?;
-            (f)(k, v)?;
+            (f)(k.as_ref(), v.as_ref())?;
         }
         Ok(())
     }
@@ -430,7 +431,7 @@ where
     where
         K: Borrow<Q> + Clone,
         Q: Eq + Hash + ?Sized,
-        V: DeserializeOwned,
+        V: DeserializeOwned + Clone,
         F: FnMut(&K, &V) -> anyhow::Result<()>,
     {
         let mut iter = match &starting_key {
@@ -441,10 +442,10 @@ where
         let mut traversed = 0usize;
         for res in iter.by_ref().take(max.unwrap_or(usize::MAX)) {
             let (k, v) = res?;
-            (f)(k, v)?;
+            (f)(k.as_ref(), v.as_ref())?;
             traversed += 1;
         }
-        let next = iter.next().transpose()?.map(|kv| kv.0).cloned();
+        let next = iter.next().transpose()?.map(|kv| kv.0.as_ref().clone());
         Ok((traversed, next))
     }
 
@@ -456,8 +457,8 @@ where
 
 impl<BS, V, K, H, Ver> HamtImpl<BS, V, K, H, Ver>
 where
-    K: DeserializeOwned + PartialOrd,
-    V: DeserializeOwned,
+    K: DeserializeOwned + PartialOrd + Clone,
+    V: DeserializeOwned + Clone,
     Ver: Version,
     BS: Blockstore,
 {
@@ -513,10 +514,10 @@ where
     /// for res in hamt.iter_from(results.last().unwrap().0)?.skip(1) {
     ///     results.push((res?));
     /// }
-    /// results.sort_by_key(|kv| kv.1);
+    /// results.sort_by_key(|kv| kv.1.clone());
     ///
     /// // Assert that we got out what we put in.
-    /// let results: Vec<_> = results.into_iter().map(|(k, v)|(k.clone(), v.clone())).collect();
+    /// let results: Vec<_> = results.into_iter().map(|(k, v)|(k.clone(), v.as_ref().clone())).collect();
     /// assert_eq!(kvs, results);
     ///
     /// # anyhow::Ok(())
@@ -533,12 +534,12 @@ where
 
 impl<'a, BS, V, K, H, Ver> IntoIterator for &'a HamtImpl<BS, V, K, H, Ver>
 where
-    K: DeserializeOwned + PartialOrd,
-    V: DeserializeOwned,
+    K: DeserializeOwned + PartialOrd + Clone,
+    V: DeserializeOwned + Clone,
     Ver: Version,
     BS: Blockstore,
 {
-    type Item = Result<(&'a K, &'a V), Error>;
+    type Item = Result<(IterItem<'a, K>, IterItem<'a, V>), Error>;
     type IntoIter = IterImpl<'a, BS, V, K, H, Ver>;
 
     fn into_iter(self) -> Self::IntoIter {
diff --git a/ipld/hamt/src/iter.rs b/ipld/hamt/src/iter.rs
@@ -1,6 +1,7 @@
 // Copyright 2021-2023 Protocol Labs
 // SPDX-License-Identifier: Apache-2.0, MIT
 use std::borrow::Borrow;
+use std::fmt::Display;
 use std::iter::FusedIterator;
 
 use forest_hash_utils::BytesKey;
@@ -17,8 +18,8 @@ use crate::{Config, Error, Hash, HashAlgorithm, KeyValuePair, Sha256};
 pub struct IterImpl<'a, BS, V, K = BytesKey, H = Sha256, Ver = version::V3> {
     store: &'a BS,
     conf: &'a Config,
-    stack: Vec<std::slice::Iter<'a, Pointer<K, V, H, Ver>>>,
-    current: std::slice::Iter<'a, KeyValuePair<K, V>>,
+    stack: Vec<StackItem<'a, Pointer<K, V, H, Ver>>>,
+    current: StackItem<'a, KeyValuePair<K, V>>,
 }
 
 /// Iterator over HAMT Key/Value tuples (hamt v0).
@@ -27,19 +28,90 @@ pub type Iterv0<'a, BS, V, K = BytesKey, H = Sha256> = IterImpl<'a, BS, V, K, H,
 /// Iterator over HAMT Key/Value tuples.
 pub type Iter<'a, BS, V, K = BytesKey, H = Sha256> = IterImpl<'a, BS, V, K, H, version::V3>;
 
+enum StackItem<'a, V> {
+    Iter(std::slice::Iter<'a, V>),
+    IntoIter(std::vec::IntoIter<V>),
+}
+
+impl<'a, V> From<std::slice::Iter<'a, V>> for StackItem<'a, V> {
+    fn from(value: std::slice::Iter<'a, V>) -> Self {
+        Self::Iter(value)
+    }
+}
+
+impl<'a, V> From<std::vec::IntoIter<V>> for StackItem<'a, V> {
+    fn from(value: std::vec::IntoIter<V>) -> Self {
+        Self::IntoIter(value)
+    }
+}
+
+impl<'a, V> Iterator for StackItem<'a, V> {
+    type Item = IterItem<'a, V>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            Self::Iter(it) => it.next().map(|i| i.into()),
+            Self::IntoIter(it) => it.next().map(|i| i.into()),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub enum IterItem<'a, V> {
+    Borrowed(&'a V),
+    Owned(V),
+}
+
+impl<V: Display> Display for IterItem<'_, V> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.as_ref().fmt(f)
+    }
+}
+
+impl<V> AsRef<V> for IterItem<'_, V> {
+    fn as_ref(&self) -> &V {
+        match self {
+            Self::Borrowed(v) => v,
+            Self::Owned(v) => v,
+        }
+    }
+}
+
+impl<V: PartialEq> PartialEq<V> for IterItem<'_, V> {
+    fn eq(&self, other: &V) -> bool {
+        self.as_ref().eq(other)
+    }
+
+    fn ne(&self, other: &V) -> bool {
+        self.as_ref().ne(other)
+    }
+}
+
+impl<'a, V> From<V> for IterItem<'a, V> {
+    fn from(value: V) -> Self {
+        Self::Owned(value)
+    }
+}
+
+impl<'a, V> From<&'a V> for IterItem<'a, V> {
+    fn from(value: &'a V) -> Self {
+        Self::Borrowed(value)
+    }
+}
+
 impl<'a, K, V, BS, H, Ver> IterImpl<'a, BS, V, K, H, Ver>
 where
-    K: DeserializeOwned,
-    V: DeserializeOwned,
+    K: DeserializeOwned + Clone,
+    V: DeserializeOwned + Clone,
     Ver: Version,
     BS: Blockstore,
 {
     pub(crate) fn new(store: &'a BS, root: &'a Node<K, V, H, Ver>, conf: &'a Config) -> Self {
         Self {
             conf,
             store,
-            stack: vec![root.pointers.iter()],
-            current: [].iter(),
+            stack: vec![root.pointers.iter().into()],
+            current: [].iter().into(),
         }
     }
 
@@ -56,24 +128,52 @@ where
     {
         let hashed_key = H::hash(key);
         let mut hash = HashBits::new(&hashed_key);
-        let mut node = root;
+        let mut node = IterItem::Borrowed(root);
         let mut stack = Vec::new();
         loop {
             let idx = hash.next(conf.bit_width)?;
-            stack.push(node.pointers[node.index_for_bit_pos(idx)..].iter());
+            match node.clone() {
+                IterItem::Borrowed(node) => {
+                    stack.push(StackItem::from(
+                        node.pointers[node.index_for_bit_pos(idx)..].iter(),
+                    ));
+                }
+                IterItem::Owned(node) => {
+                    stack.push(StackItem::from(
+                        node.pointers[node.index_for_bit_pos(idx)..]
+                            .to_vec()
+                            .into_iter(),
+                    ));
+                }
+            }
             node = match stack.last_mut().unwrap().next() {
                 Some(p) => match p {
-                    Pointer::Link { cid, cache } => cache.get_or_try_init(|| {
-                        Node::load(conf, store, cid, stack.len() as u32).map(Box::new)
-                    })?,
-                    Pointer::Dirty(node) => node,
-                    Pointer::Values(values) => {
+                    IterItem::Borrowed(Pointer::Link { cid, cache: _ }) => {
+                        Node::load(conf, store, cid, stack.len() as u32)?.into()
+                    }
+                    IterItem::Owned(Pointer::Link { cid, cache: _ }) => {
+                        Node::load(conf, store, &cid, stack.len() as u32)?.into()
+                    }
+                    IterItem::Borrowed(Pointer::Dirty(node)) => node.as_ref().into(),
+                    IterItem::Owned(Pointer::Dirty(node)) => (*node).into(),
+                    IterItem::Borrowed(Pointer::Values(values)) => {
+                        return match values.iter().position(|kv| kv.key().borrow() == key) {
+                            Some(offset) => Ok(Self {
+                                conf,
+                                store,
+                                stack,
+                                current: values[offset..].iter().into(),
+                            }),
+                            None => Err(Error::StartKeyNotFound),
+                        };
+                    }
+                    IterItem::Owned(Pointer::Values(values)) => {
                         return match values.iter().position(|kv| kv.key().borrow() == key) {
                             Some(offset) => Ok(Self {
                                 conf,
                                 store,
                                 stack,
-                                current: values[offset..].iter(),
+                                current: values[offset..].to_vec().into_iter().into(),
                             }),
                             None => Err(Error::StartKeyNotFound),
                         };
@@ -92,33 +192,64 @@ where
     K: DeserializeOwned + PartialOrd,
     V: DeserializeOwned,
 {
-    type Item = Result<(&'a K, &'a V), Error>;
+    type Item = Result<(IterItem<'a, K>, IterItem<'a, V>), Error>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if let Some(v) = self.current.next() {
-            return Some(Ok((v.key(), v.value())));
+        match self.current.next() {
+            Some(IterItem::Borrowed(v)) => return Some(Ok((v.key().into(), v.value().into()))),
+            Some(IterItem::Owned(KeyValuePair(k, v))) => return Some(Ok((k.into(), v.into()))),
+            _ => {}
         }
         loop {
             let Some(next) = self.stack.last_mut()?.next() else {
                 self.stack.pop();
                 continue;
             };
             match next {
-                Pointer::Link { cid, cache } => {
-                    let node = match cache.get_or_try_init(|| {
-                        Node::load(self.conf, &self.store, cid, self.stack.len() as u32)
-                            .map(Box::new)
-                    }) {
-                        Ok(node) => node,
-                        Err(e) => return Some(Err(e)),
-                    };
-                    self.stack.push(node.pointers.iter())
+                IterItem::Borrowed(Pointer::Link { cid, cache: _ }) => {
+                    let node =
+                        match Node::load(self.conf, &self.store, cid, self.stack.len() as u32) {
+                            Ok(node) => node,
+                            Err(e) => return Some(Err(e)),
+                        };
+                    self.stack.push(node.pointers.into_iter().into())
+                }
+                IterItem::Owned(Pointer::Link { cid, cache: _ }) => {
+                    let node =
+                        match Node::load(self.conf, &self.store, &cid, self.stack.len() as u32) {
+                            Ok(node) => node,
+                            Err(e) => return Some(Err(e)),
+                        };
+                    self.stack.push(node.pointers.into_iter().into())
+                }
+                IterItem::Borrowed(Pointer::Dirty(node)) => {
+                    self.stack.push(node.pointers.iter().into())
+                }
+                IterItem::Owned(Pointer::Dirty(node)) => {
+                    self.stack.push(node.pointers.into_iter().into())
+                }
+                IterItem::Borrowed(Pointer::Values(kvs)) => {
+                    self.current = kvs.iter().into();
+                    match self.current.next() {
+                        Some(IterItem::Borrowed(v)) => {
+                            return Some(Ok((v.key().into(), v.value().into())));
+                        }
+                        Some(IterItem::Owned(KeyValuePair(k, v))) => {
+                            return Some(Ok((k.into(), v.into())));
+                        }
+                        _ => {}
+                    }
                 }
-                Pointer::Dirty(node) => self.stack.push(node.pointers.iter()),
-                Pointer::Values(kvs) => {
-                    self.current = kvs.iter();
-                    if let Some(v) = self.current.next() {
-                        return Some(Ok((v.key(), v.value())));
+                IterItem::Owned(Pointer::Values(kvs)) => {
+                    self.current = kvs.into_iter().into();
+                    match self.current.next() {
+                        Some(IterItem::Borrowed(v)) => {
+                            return Some(Ok((v.key().into(), v.value().into())));
+                        }
+                        Some(IterItem::Owned(KeyValuePair(k, v))) => {
+                            return Some(Ok((k.into(), v.into())));
+                        }
+                        _ => {}
                     }
                 }
             }
diff --git a/ipld/hamt/src/lib.rs b/ipld/hamt/src/lib.rs
@@ -75,7 +75,7 @@ impl Default for Config {
 
 type HashedKey = [u8; 32];
 
-#[derive(Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 struct KeyValuePair<K, V>(K, V);
 
 impl<K, V> KeyValuePair<K, V> {
diff --git a/ipld/hamt/src/node.rs b/ipld/hamt/src/node.rs
@@ -29,6 +29,20 @@ pub(crate) struct Node<K, V, H, Ver = version::V3> {
     hash: PhantomData<H>,
 }
 
+impl<K, V, H, Ver> Clone for Node<K, V, H, Ver>
+where
+    K: Clone,
+    V: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            bitfield: self.bitfield.clone(),
+            pointers: self.pointers.clone(),
+            hash: Default::default(),
+        }
+    }
+}
+
 impl<K: PartialEq, V: PartialEq, H, Ver> PartialEq for Node<K, V, H, Ver> {
     fn eq(&self, other: &Self) -> bool {
         (self.bitfield == other.bitfield) && (self.pointers == other.pointers)
diff --git a/ipld/hamt/src/pointer.rs b/ipld/hamt/src/pointer.rs
@@ -46,6 +46,23 @@ pub(crate) enum Pointer<K, V, H, Ver = version::V3> {
     Dirty(Box<Node<K, V, H, Ver>>),
 }
 
+impl<K, V, H, Ver> Clone for Pointer<K, V, H, Ver>
+where
+    K: Clone,
+    V: Clone,
+{
+    fn clone(&self) -> Self {
+        match self {
+            Self::Values(v) => Self::Values(v.clone()),
+            Self::Link { cid, cache: _ } => Self::Link {
+                cid: *cid,
+                cache: Default::default(),
+            },
+            Self::Dirty(n) => Self::Dirty(n.clone()),
+        }
+    }
+}
+
 impl<K: PartialEq, V: PartialEq, H, Ver> PartialEq for Pointer<K, V, H, Ver> {
     fn eq(&self, other: &Self) -> bool {
         match (self, other) {
