fix(collections): visit wrapped collision entries once in retain

hashmap/utils.mbt

@@ -409,6 +409,7 @@ pub fn[K, V] HashMap::values(self : HashMap[K, V]) -> Iter[V] {
 /// Retains only the key-value pairs that satisfy the given predicate function.
 /// This method modifies the hash map in-place, removing all entries for which
 /// the predicate returns `false`.
+/// Calls the predicate once per entry, in unspecified order.
 ///
 /// Parameters:
 ///
@@ -432,8 +433,15 @@ pub fn[K, V] HashMap::values(self : HashMap[K, V]) -> Iter[V] {
 #locals(f)
 pub fn[K, V] HashMap::retain(self : HashMap[K, V], f : (K, V) -> Bool) -> Unit {
   let size = self.size
+  guard size > 0 else { return }
+  // Begin at an empty slot so a wrapped collision chain is visited together.
+  // Starting at zero could revisit retained entries shifted across that boundary.
+  let mut start = 0
+  while self.entries[start] is Some(_) {
+    start += 1
+  }
   let mut j = 0
-  for i = 0; j < size; i = i + 1 {
+  for i = start; j < size; i = (i + 1) & self.capacity_mask {
     while self.entries[i] is Some(entry) {
       j += 1
       if f(entry.key, entry.value) {

hashmap/utils_test.mbt

@@ -176,3 +176,40 @@ test "T::retain - preserve map integrity after retain" {
   inspect(map.length(), content="0")
   inspect(map.is_empty(), content="true")
 }
+
+///|
+priv struct RetainCollisionKey(Int) derive(Eq)
+
+///|
+impl Hash for RetainCollisionKey with fn hash(_) {
+  7
+}
+
+///|
+impl Hash for RetainCollisionKey with fn hash_combine(_, hasher) {
+  hasher.combine_int(7)
+}
+
+///|
+test "retain visits wrapped collision entries once" {
+  let map = @hashmap.HashMap(
+    [
+      (RetainCollisionKey(0), 0),
+      (RetainCollisionKey(1), 1),
+      (RetainCollisionKey(2), 2),
+      (RetainCollisionKey(3), 3),
+    ],
+    capacity=8,
+  )
+  let visits = []
+  map.retain((_, value) => {
+    visits.push(value)
+    value != 0
+  })
+  visits.sort()
+  debug_inspect(visits, content="[0, 1, 2, 3]")
+  inspect(map.length(), content="3")
+  for id in 1..<4 {
+    assert_eq(map.get(RetainCollisionKey(id)), Some(id))
+  }
+}

hashset/hashset.mbt

@@ -503,11 +503,19 @@ pub impl[K : Hash + Eq] Sub for HashSet[K] with fn sub(self, other) {
 
 ///|
 /// Removes all elements for which the predicate returns `false`.
+/// Calls the predicate once per element, in unspecified order.
 #locals(f)
 pub fn[K] HashSet::retain(self : HashSet[K], f : (K) -> Bool) -> Unit {
   let size = self.size
+  guard size > 0 else { return }
+  // Begin at an empty slot so a wrapped collision chain is visited together.
+  // Starting at zero could revisit retained entries shifted across that boundary.
+  let mut start = 0
+  while self.entries[start] is Some(_) {
+    start += 1
+  }
   let mut j = 0
-  for i = 0; j < size; i = i + 1 {
+  for i = start; j < size; i = (i + 1) & self.capacity_mask {
     while self.entries[i] is Some(entry) {
       j += 1
       if f(entry.key) {

hashset/hashset_test.mbt

@@ -477,3 +477,23 @@ test "retain/empty_set" {
   set.retain(_x => true)
   @test.assert_eq(set.length(), 0)
 }
+
+///|
+test "retain visits wrapped collision entries once" {
+  let set = @hashset.HashSet(
+    [HashSetKey(0, 7), HashSetKey(1, 7), HashSetKey(2, 7), HashSetKey(3, 7)],
+    capacity=8,
+  )
+  let visits = []
+  set.retain(key => {
+    let HashSetKey(id, _) = key
+    visits.push(id)
+    id != 0
+  })
+  visits.sort()
+  debug_inspect(visits, content="[0, 1, 2, 3]")
+  inspect(set.length(), content="3")
+  for id in 1..<4 {
+    assert_true(set.contains(HashSetKey(id, 7)))
+  }
+}