From 723c340a8b11d5cc083ea8581cd650d025dcc4e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Fri, 12 Apr 2024 10:54:21 +0200
Subject: [PATCH] perf: fix linearity in (HashSet|HashMap).erase (#3887)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes linearity issues in HashSet/HashMap erase functions.

IR before patch:
```
def Lean.HashMapImp.erase._rarg (x_1 : obj) (x_2 : obj) (x_3 : obj) (x_4 : obj) : obj :=
  let x_5 : obj := proj[0] x_3;
  inc x_5;
  let x_6 : obj := proj[1] x_3;
  inc x_6;
  let x_7 : obj := Array.size ◾ x_6;
  inc x_4;
  let x_8 : obj := app x_2 x_4;
  let x_9 : u64 := unbox x_8;
  dec x_8;
  let x_10 : usize := _private.Lean.Data.HashMap.0.Lean.HashMapImp.mkIdx x_7 x_9 ◾;
  let x_11 : obj := Array.uget ◾ x_6 x_10 ◾;
  inc x_11;
  inc x_4;
  inc x_1;
  let x_12 : u8 := Lean.AssocList.contains._rarg x_1 x_4 x_11;
  case x_12 : u8 of
  Bool.false →
    dec x_11;
    dec x_6;
    dec x_5;
    dec x_4;
    dec x_1;
    ret x_3
  Bool.true →
    let x_13 : u8 := isShared x_3;
    case x_13 : u8 of
    Bool.false →
      let x_14 : obj := proj[1] x_3;
      dec x_14;
      let x_15 : obj := proj[0] x_3;
      dec x_15;
      let x_16 : obj := 1;
      let x_17 : obj := Nat.sub x_5 x_16;
      dec x_5;
      let x_18 : obj := Lean.AssocList.erase._rarg x_1 x_4 x_11;
      let x_19 : obj := Array.uset ◾ x_6 x_10 x_18 ◾;
      set x_3[1] := x_19;
      set x_3[0] := x_17;
      ret x_3
    Bool.true →
      dec x_3;
      let x_20 : obj := 1;
      let x_21 : obj := Nat.sub x_5 x_20;
      dec x_5;
      let x_22 : obj := Lean.AssocList.erase._rarg x_1 x_4 x_11;
      let x_23 : obj := Array.uset ◾ x_6 x_10 x_22 ◾;
      let x_24 : obj := ctor_0[Lean.HashMapImp.mk] x_21 x_23;
      ret x_24
```

IR after the patch:
```
def Lean.HashMapImp.erase._rarg (x_1 : obj) (x_2 : obj) (x_3 : obj) (x_4 : obj) : obj :=
  let x_5 : u8 := isShared x_3;
  case x_5 : u8 of
  Bool.false →
    let x_6 : obj := proj[0] x_3;
    let x_7 : obj := proj[1] x_3;
    let x_8 : obj := Array.size ◾ x_7;
    inc x_4;
    let x_9 : obj := app x_2 x_4;
    let x_10 : u64 := unbox x_9;
    dec x_9;
    let x_11 : usize := _private.Lean.Data.HashMap.0.Lean.HashMapImp.mkIdx x_8 x_10 ◾;
    let x_12 : obj := Array.uget ◾ x_7 x_11 ◾;
    inc x_12;
    inc x_4;
    inc x_1;
    let x_13 : u8 := Lean.AssocList.contains._rarg x_1 x_4 x_12;
    case x_13 : u8 of
    Bool.false →
      dec x_12;
      dec x_4;
      dec x_1;
      ret x_3
    Bool.true →
      let x_14 : obj := 1;
      let x_15 : obj := Nat.sub x_6 x_14;
      dec x_6;
      let x_16 : obj := Lean.AssocList.erase._rarg x_1 x_4 x_12;
      let x_17 : obj := Array.uset ◾ x_7 x_11 x_16 ◾;
      set x_3[1] := x_17;
      set x_3[0] := x_15;
      ret x_3
  Bool.true →
    let x_18 : obj := proj[0] x_3;
    let x_19 : obj := proj[1] x_3;
    inc x_19;
    inc x_18;
    dec x_3;
    let x_20 : obj := Array.size ◾ x_19;
    inc x_4;
    let x_21 : obj := app x_2 x_4;
    let x_22 : u64 := unbox x_21;
    dec x_21;
    let x_23 : usize := _private.Lean.Data.HashMap.0.Lean.HashMapImp.mkIdx x_20 x_22 ◾;
    let x_24 : obj := Array.uget ◾ x_19 x_23 ◾;
    inc x_24;
    inc x_4;
    inc x_1;
    let x_25 : u8 := Lean.AssocList.contains._rarg x_1 x_4 x_24;
    case x_25 : u8 of
    Bool.false →
      dec x_24;
      dec x_4;
      dec x_1;
      let x_26 : obj := ctor_0[Lean.HashMapImp.mk] x_18 x_19;
      ret x_26
    Bool.true →
      let x_27 : obj := 1;
      let x_28 : obj := Nat.sub x_18 x_27;
      dec x_18;
      let x_29 : obj := Lean.AssocList.erase._rarg x_1 x_4 x_24;
      let x_30 : obj := Array.uset ◾ x_19 x_23 x_29 ◾;
      let x_31 : obj := ctor_0[Lean.HashMapImp.mk] x_28 x_30;
      ret x_31
```

Previously `x_6` (the buckets array) always gets `inc`remented, now only
if the HashMap itself is shared.
---
 src/Lean/Data/HashMap.lean | 6 ++++--
 src/Lean/Data/HashSet.lean | 6 ++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/Lean/Data/HashMap.lean b/src/Lean/Data/HashMap.lean
index dc30823cb7..52328515f2 100644
--- 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
index 5bf16b14f7..927b6e10d6 100644
--- 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)