From 06efcb30f3daa1f789a05d8118f52da6654996d5 Mon Sep 17 00:00:00 2001
From: Josh Jones <joshua.jones@chainsafe.io>
Date: Wed, 8 Nov 2023 09:26:45 +0000
Subject: [PATCH] revert `for_each` code in `amt_tests` and create min test
 case in `iter`

---
 ipld/amt/src/iter.rs        | 52 ++++++++++++++++++++++++++++
 ipld/amt/tests/amt_tests.rs | 68 ++++++++++++++++++++++---------------
 2 files changed, 92 insertions(+), 28 deletions(-)

diff --git a/ipld/amt/src/iter.rs b/ipld/amt/src/iter.rs
index 794c0141c..a07ab0664 100644
--- a/ipld/amt/src/iter.rs
+++ b/ipld/amt/src/iter.rs
@@ -165,7 +165,10 @@ where
 #[cfg(test)]
 mod tests {
     use crate::Amt;
+    use fvm_ipld_blockstore::tracking::BSStats;
+    use fvm_ipld_blockstore::tracking::TrackingBlockstore;
     use fvm_ipld_blockstore::MemoryBlockstore;
+    use fvm_ipld_encoding::BytesDe;
     use quickcheck_macros::quickcheck;
 
     #[test]
@@ -240,6 +243,55 @@ mod tests {
         assert_eq!(expected, restored);
     }
 
+    // Helper function for `for_each` test
+    fn tbytes(bz: &[u8]) -> BytesDe {
+        BytesDe(bz.to_vec())
+    }
+
+    #[test]
+    fn minimal_for_each() {
+        let mem = MemoryBlockstore::default();
+        let db = TrackingBlockstore::new(&mem);
+        let mut a = Amt::new(&db);
+
+        let mut indexes = Vec::new();
+        for i in 0..3 {
+            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();
+        }
+
+        // Flush and regenerate amt
+        let c = a.flush().unwrap();
+        let new_amt = Amt::load(&c, &db).unwrap();
+
+        let mut x = 0;
+        // dbg!(&indexes, &new_amt);
+        // dbg!(new_amt.iter().enumerate());
+        // for (i, v) in new_amt.iter().enumerate() {
+        //     dbg!((i,v));
+        // }
+        #[allow(deprecated)]
+        new_amt
+            .for_each(|i, _: &BytesDe| {
+                if i != indexes[x] {
+                    panic!(
+                        "for each found wrong index: expected {} got {}",
+                        indexes[x], i
+                    );
+                }
+                x += 1;
+                Ok(())
+            })
+            .unwrap();
+        assert_eq!(x, indexes.len());
+    }
+
     #[quickcheck]
     fn vary_bit_width(bit_width: u32) {
         // `bit_width` is only limited due to the test taking too long to run at higher values.
diff --git a/ipld/amt/tests/amt_tests.rs b/ipld/amt/tests/amt_tests.rs
index 23374e900..d6e097d38 100644
--- a/ipld/amt/tests/amt_tests.rs
+++ b/ipld/amt/tests/amt_tests.rs
@@ -332,10 +332,10 @@ fn delete_reduce_height() {
 fn for_each() {
     let mem = MemoryBlockstore::default();
     let db = TrackingBlockstore::new(&mem);
-    let mut a = Amt::new_with_bit_width(&db, 1);
+    let mut a = Amt::new(&db);
 
     let mut indexes = Vec::new();
-    for i in 0..6 {
+    for i in 0..10000 {
         if (i + 1) % 3 == 0 {
             indexes.push(i);
         }
@@ -343,45 +343,57 @@ fn for_each() {
 
     // Set all indices in the Amt
     for i in indexes.iter() {
-        a.set(*i, *i).unwrap();
+        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")));
-    // }
+    // 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);
-    // println!("{:?}", a);
-    // for i in a.iter() {
-    //     println!("{:?}", i);
-    // }
-    assert_eq!(usize::try_from(a.count()).unwrap(), a.iter().count());
+
+    // Iterate over amt with dirty cache
+    let mut x = 0;
+    #[allow(deprecated)]
+    a.for_each(|_, _: &BytesDe| {
+        x += 1;
+        Ok(())
+    })
+    .unwrap();
+
+    assert_eq!(x, indexes.len());
 
     // Flush and regenerate amt
     let c = a.flush().unwrap();
-    let new_amt: AmtImpl<ByteBuf, &TrackingBlockstore<_>, _> = Amt::load(&c, &db).unwrap();
+    let new_amt = Amt::load(&c, &db).unwrap();
     assert_eq!(new_amt.count(), indexes.len() as u64);
 
+    let mut x = 0;
+    #[allow(deprecated)]
     new_amt
-        .iter()
-        .map(|_| Ok::<(), fvm_ipld_amt::Error>(()))
-        .collect::<Result<Vec<_>, _>>()
+        .for_each(|i, _: &BytesDe| {
+            if i != indexes[x] {
+                panic!(
+                    "for each found wrong index: expected {} got {}",
+                    indexes[x], i
+                );
+            }
+            x += 1;
+            Ok(())
+        })
         .unwrap();
-    let x = new_amt.count() as usize;
     assert_eq!(x, indexes.len());
 
-    // new_amt
-    //     .iter()
-    //     .map(|_| Ok::<(), fvm_ipld_amt::Error>(()))
-    //     .collect::<Result<Vec<_>, _>>()
-    //     .unwrap();
-    // assert_eq!(
-    //     c.to_string().as_str(),
-    //     "bafy2bzaceanqxtbsuyhqgxubiq6vshtbhktmzp2if4g6kxzttxmzkdxmtipcm"
-    // );
-
-    // j
+    #[allow(deprecated)]
+    new_amt.for_each(|_, _: &BytesDe| Ok(())).unwrap();
+    assert_eq!(
+        c.to_string().as_str(),
+        "bafy2bzaceanqxtbsuyhqgxubiq6vshtbhktmzp2if4g6kxzttxmzkdxmtipcm"
+    );
+
+    #[rustfmt::skip]
+    assert_eq!(*db.stats.borrow(), BSStats {r: 1431, w: 1431, br: 88649, bw: 88649});
 }
 
 #[test]