feat: cacheless amt iteration (#2189)

* feat: cacheless amt iteration

* docs, comments

* doc: remove flush hint

* revert doc fix

Author: LesnyRumcajs

ipld/amt/src/amt.rs

@@ -590,4 +590,39 @@ where
 
         Ok(())
     }
+
+    /// Iterates over each value in the Amt and runs a function on the values. This is a
+    /// non-caching version of [`Self::for_each`]. It can potentially be more efficient, especially memory-wise,
+    /// for large AMTs or when the iteration occurs only once.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use fvm_ipld_amt::Amt;
+    ///
+    /// let store = fvm_ipld_blockstore::MemoryBlockstore::default();
+    ///
+    /// let mut map: Amt<String, _> = Amt::new(&store);
+    /// map.set(1, "One".to_owned()).unwrap();
+    /// map.set(4, "Four".to_owned()).unwrap();
+    ///
+    /// let mut values: Vec<(u64, String)> = Vec::new();
+    /// map.for_each_cacheless(|i, v| {
+    ///    values.push((i, v.clone()));
+    ///    Ok(())
+    /// }).unwrap();
+    /// assert_eq!(&values, &[(1, "One".to_owned()), (4, "Four".to_owned())]);
+    /// ```
+    pub fn for_each_cacheless<F>(&self, mut f: F) -> Result<(), Error>
+    where
+        F: FnMut(u64, &V) -> anyhow::Result<()>,
+    {
+        self.root.node.for_each_cacheless(
+            &self.block_store,
+            self.height(),
+            self.bit_width(),
+            0,
+            &mut f,
+        )
+    }
 }

ipld/amt/src/node.rs

@@ -420,6 +420,52 @@ where
         }
     }
 
+    /// Non-caching iteration over the values in the node.
+    pub(super) fn for_each_cacheless<S, F>(
+        &self,
+        bs: &S,
+        height: u32,
+        bit_width: u32,
+        offset: u64,
+        f: &mut F,
+    ) -> Result<(), Error>
+    where
+        F: FnMut(u64, &V) -> anyhow::Result<()>,
+        S: Blockstore,
+    {
+        match self {
+            Node::Leaf { vals } => {
+                for (i, v) in (0..).zip(vals.iter()) {
+                    if let Some(v) = v {
+                        let _ = f(offset + i, v);
+                    }
+                }
+            }
+            Node::Link { links } => {
+                for (i, l) in (0..).zip(links.iter()) {
+                    if let Some(link) = l {
+                        let offs = offset + (i * nodes_for_height(bit_width, height));
+                        match link {
+                            Link::Dirty(sub) => {
+                                sub.for_each_cacheless(bs, height - 1, bit_width, offs, f)?;
+                            }
+                            Link::Cid { cid, cache: _ } => {
+                                let node = bs
+                                    .get_cbor::<CollapsedNode<V>>(cid)?
+                                    .ok_or_else(|| Error::CidNotFound(cid.to_string()))?
+                                    .expand(bit_width)?;
+
+                                node.for_each_cacheless(bs, height - 1, bit_width, offs, f)?;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     /// Returns a `(keep_going, did_mutate)` pair. `keep_going` will be `false` iff
     /// a closure call returned `Ok(false)`, indicating that a `break` has happened.
     /// `did_mutate` will be `true` iff any of the values in the node was actually

ipld/amt/tests/amt_tests.rs

@@ -390,6 +390,77 @@ fn for_each() {
     assert_eq!(*db.stats.borrow(), BSStats {r: 1431, w: 1431, br: 88649, bw: 88649});
 }
 
+#[test]
+fn for_each_cacheless() {
+    let mem = MemoryBlockstore::default();
+    let db = TrackingBlockstore::new(&mem);
+    let mut a = Amt::new(&db);
+
+    let mut indexes = Vec::new();
+    for i in 0..10000 {
+        if (i + 1) % 3 == 0 {
+            indexes.push(i);
+        }
+    }
+
+    // Set all indices in the Amt
+    for i in indexes.iter() {
+        a.set(*i, tbytes(b"value")).unwrap();
+    }
+
+    // Ensure all values were added into the amt
+    for i in indexes.iter() {
+        assert_eq!(a.get(*i).unwrap(), Some(&tbytes(b"value")));
+    }
+
+    assert_eq!(a.count(), indexes.len() as u64);
+
+    // Iterate over amt with dirty cache. This should not touch the database at all.
+    let mut x = 0;
+    a.for_each_cacheless(|_, _: &BytesDe| {
+        x += 1;
+        Ok(())
+    })
+    .unwrap();
+    #[rustfmt::skip]
+    assert_eq!(*db.stats.borrow(), BSStats {r: 0, w: 0, br: 0, bw: 0});
+    assert_eq!(x, indexes.len());
+
+    // Flush and regenerate amt
+    let c = a.flush().unwrap();
+    let new_amt = Amt::load(&c, &db).unwrap();
+    assert_eq!(new_amt.count(), indexes.len() as u64);
+
+    let mut x = 0;
+    new_amt
+        .for_each_cacheless(|i, _: &BytesDe| {
+            if i != indexes[x] {
+                panic!(
+                    "for each cacheless found wrong index: expected {} got {}",
+                    indexes[x], i
+                );
+            }
+            x += 1;
+            Ok(())
+        })
+        .unwrap();
+    assert_eq!(x, indexes.len());
+
+    // After 1st pass, the amount of block reads should be the same as in the caching iteration.
+    #[rustfmt::skip]
+    assert_eq!(*db.stats.borrow(), BSStats {r: 1431, w: 1431, br: 88649, bw: 88649});
+
+    new_amt.for_each_cacheless(|_, _: &BytesDe| Ok(())).unwrap();
+    assert_eq!(
+        c.to_string().as_str(),
+        "bafy2bzaceanqxtbsuyhqgxubiq6vshtbhktmzp2if4g6kxzttxmzkdxmtipcm"
+    );
+
+    // 2nd pass, no caching so block reads roughly doubled.
+    #[rustfmt::skip]
+    assert_eq!(*db.stats.borrow(), BSStats {r: 2861, w: 1431, br: 177158, bw: 88649});
+}
+
 #[test]
 fn for_each_ranged() {
     let mem = MemoryBlockstore::default();