fix: Sync len after map rebuilds in iter_mut/iter_keys_mut for HashMap and UHashMap (#9491)

Co-authored-by: Tom French <[email protected]>
Co-authored-by: Jake Fecher <[email protected]>

Author: 3 people

noir_stdlib/src/collections/map.nr

@@ -319,6 +319,7 @@ impl<K, V, let N: u32, B> HashMap<K, V, N, B> {
         }
 
         self._table = new_map._table;
+        self._len = new_map._len;
     }
 
     /// Iterates through the HashMap, mutating each key to the result returned from
@@ -356,6 +357,7 @@ impl<K, V, let N: u32, B> HashMap<K, V, N, B> {
         }
 
         self._table = new_map._table;
+        self._len = new_map._len;
     }
 
     /// Iterates through the HashMap, applying the given function to each value and mutating the

noir_stdlib/src/collections/umap.nr

@@ -210,6 +210,7 @@ impl<K, V, B> UHashMap<K, V, B> {
         }
 
         self._table = new_map._table;
+        self._len = new_map._len;
     }
 
     // For each key applies mutator function.
@@ -229,6 +230,7 @@ impl<K, V, B> UHashMap<K, V, B> {
         }
 
         self._table = new_map._table;
+        self._len = new_map._len;
     }
 
     // For each value applies mutator function.

test_programs/execution_success/hashmap/src/main.nr

@@ -36,6 +36,7 @@ fn main(input: [Entry; HASHMAP_LEN]) {
     test_retain();
     test_iterators();
     test_mut_iterators();
+    test_iter_len_shrink_on_key_collision();
 
     doc_tests();
 }
@@ -199,6 +200,29 @@ fn test_mut_iterators() {
     assert(entries == [(12, 30), (30, 70), (66, 130)], "Got incorrect iteration of entries.");
 }
 
+// Ensure that mutating keys to collide shrinks the length and keeps iterator invariants intact.
+fn test_iter_len_shrink_on_key_collision() {
+    let mut hashmap = ALLOCATE_HASHMAP();
+
+    hashmap.insert(1, 10);
+    hashmap.insert(2, 20);
+    hashmap.insert(3, 30);
+
+    // Mutate all keys to the same value to force collisions.
+    hashmap.iter_keys_mut(|_k: K| -> K { 0 });
+
+    assert_eq(
+        hashmap.len(),
+        1,
+        "HashMap length should shrink to 1 after key collisions during iter_keys_mut.",
+    );
+
+    // These calls will assert internally if _len is not synchronized with the table.
+    let _ = hashmap.entries();
+    let _ = hashmap.keys();
+    let _ = hashmap.values();
+}
+
 // docs:start:type_alias
 type MyMap = HashMap<u8, u32, 10, BuildHasherDefault<Poseidon2Hasher>>;
 // docs:end:type_alias