perf: fix linearity in (HashSet|HashMap).erase

leanprover · Apr 12, 2024 · 321f8c0 · 321f8c0
1 parent 2e3d523
commit 321f8c0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/Lean/Data/HashMap.lean b/src/Lean/Data/HashMap.lean
@@ -137,8 +137,10 @@ def erase [BEq α] [Hashable α] (m : HashMapImp α β) (a : α) : HashMapImp α
   | ⟨ size, buckets ⟩ =>
     let ⟨i, h⟩ := mkIdx (hash a) buckets.property
     let bkt    := buckets.val[i]
-    if bkt.contains a then ⟨size - 1, buckets.update i (bkt.erase a) h⟩
-    else m
+    if bkt.contains a then
+      ⟨size - 1, buckets.update i (bkt.erase a) h⟩
+    else
+      ⟨size, buckets⟩
 
 inductive WellFormed [BEq α] [Hashable α] : HashMapImp α β → Prop where
   | mkWff          : ∀ n,                    WellFormed (mkHashMapImp n)

diff --git a/src/Lean/Data/HashSet.lean b/src/Lean/Data/HashSet.lean
@@ -112,8 +112,10 @@ def erase [BEq α] [Hashable α] (m : HashSetImp α) (a : α) : HashSetImp α :=
   | ⟨ size, buckets ⟩ =>
     let ⟨i, h⟩ := mkIdx (hash a) buckets.property
     let bkt    := buckets.val[i]
-    if bkt.contains a then ⟨size - 1, buckets.update i (bkt.erase a) h⟩
-    else m
+    if bkt.contains a then
+      ⟨size - 1, buckets.update i (bkt.erase a) h⟩
+    else
+      ⟨size, buckets⟩
 
 inductive WellFormed [BEq α] [Hashable α] : HashSetImp α → Prop where
   | mkWff     : ∀ n,                  WellFormed (mkHashSetImp n)