Fix archived B-tree map iteration

Fixes rkyv#590

Author: djkoloski

rkyv/src/collections/btree_map/iter.rs

@@ -196,32 +196,35 @@ impl<K, V, const E: usize> Iterator for RawIter<K, V, E> {
         match kind {
             NodeKind::Inner => {
                 let inner = current.cast::<InnerNode<K, V, E>>();
-                if next_i < E {
+                let next = if next_i < E {
                     // More values in the current node
                     self.stack.push((current, next_i));
 
-                    // Recurse to a lesser if valid
-                    let next_lesser =
-                        unsafe { addr_of_mut!((*inner).lesser_nodes[next_i]) };
-                    let next_lesser_is_invalid =
-                        unsafe { RelPtr::is_invalid_raw(next_lesser) };
-                    if !next_lesser_is_invalid {
-                        self.stack.push((
-                            unsafe { RelPtr::as_ptr_raw(next_lesser).cast() },
-                            0,
-                        ));
-                    }
+                    // Next is a lesser node
+                    unsafe { addr_of_mut!((*inner).lesser_nodes[next_i]) }
                 } else {
-                    // Recurse to a greater if valid
-                    let next_greater =
-                        unsafe { addr_of_mut!((*inner).greater_node) };
-                    let next_greater_is_invalid =
-                        unsafe { RelPtr::is_invalid_raw(next_greater) };
-                    if !next_greater_is_invalid {
-                        self.stack.push((
-                            unsafe { RelPtr::as_ptr_raw(next_greater).cast() },
-                            0,
-                        ));
+                    // Next is a greater node
+                    unsafe { addr_of_mut!((*inner).greater_node) }
+                };
+
+                let next_is_invalid = unsafe { RelPtr::is_invalid_raw(next) };
+                if !next_is_invalid {
+                    // Recurse left on next node
+                    let mut current = unsafe { RelPtr::as_ptr_raw(next) };
+                    loop {
+                        self.stack.push((current, 0));
+                        let kind = unsafe { (*current).kind };
+                        match kind {
+                            NodeKind::Inner => {
+                                let inner =
+                                    current.cast::<InnerNode<K, V, E>>();
+                                let lesser = unsafe {
+                                    addr_of_mut!((*inner).lesser_nodes[0])
+                                };
+                                current = unsafe { RelPtr::as_ptr_raw(lesser) };
+                            }
+                            NodeKind::Leaf => break,
+                        }
                     }
                 }
             }

rkyv/src/impls/alloc/collections/btree_map.rs

@@ -272,4 +272,41 @@ mod tests {
             assert_eq!(i.next(), None);
         });
     }
+
+    // MIRI can't run this test quickly enough
+    #[cfg(not(miri))]
+    #[test]
+    fn size_matches_iter_len() {
+        #[derive(Archive, Deserialize, Serialize)]
+        #[rkyv(crate)]
+        struct Container {
+            transforms: BTreeMap<String, String>,
+        }
+
+        impl Container {
+            pub fn fill(count: usize) -> Self {
+                let mut transforms = BTreeMap::new();
+                for i in 0..count {
+                    transforms.insert(i.to_string(), i.to_string());
+                }
+                Container { transforms }
+            }
+
+            pub fn check(&self, expected: usize) {
+                to_archived(self, |archived| {
+                    assert_eq!(archived.transforms.len(), expected);
+                    let mut count = 0;
+                    for (..) in archived.transforms.iter() {
+                        count += 1;
+                    }
+                    assert_eq!(count, expected);
+                });
+            }
+        }
+
+        for expected in 0..=200 {
+            let container = Container::fill(expected);
+            container.check(expected);
+        }
+    }
 }

rkyv/src/impls/std/collections/hash_map.rs

@@ -293,8 +293,6 @@ mod tests {
 
         let encoded = to_bytes::<Error>(&map).unwrap();
 
-        eprintln!("{encoded:?}");
-
         // This .unwrap() fails!
         let _decoded = access::<Archived<MyHashMap>, Panic>(&encoded).unwrap();
     }