From 32a8d7e008539c1d13dbcb762d4574189845f434 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Mon, 2 Dec 2024 16:35:41 +0100
Subject: [PATCH] perf: bv_decide use Lean.RArray

---
 .../BVDecide/Frontend/BVDecide/Reflect.lean   | 24 ++++++++++++-------
 .../Frontend/BVDecide/SatAtBVLogical.lean     | 17 +++++++------
 .../BVDecide/Bitblast/BVExpr/Basic.lean       | 11 +++++----
 .../Bitblast/BVExpr/Circuit/Lemmas/Basic.lean |  4 ++--
 4 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean
index bfc42764f8bc..bb4c3892eb67 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean
@@ -8,6 +8,7 @@ import Std.Data.HashMap
 import Std.Tactic.BVDecide.Bitblast.BVExpr.Basic
 import Lean.Meta.AppBuilder
 import Lean.ToExpr
+import Lean.Data.RArray
 
 /-!
 This module contains the implementation of the reflection monad, used by all other components of this
@@ -138,9 +139,11 @@ structure State where
   -/
   atoms : Std.HashMap Expr Atom := {}
   /--
-  A cache for `atomsAssignment`.
+  A cache for `atomsAssignment`. We maintain the invariant that this value is only used if
+  `atoms` is non empty. The reason for not using an `Option` is that it would pollute a lot of code
+  with error handling that is never hit as this invariant is enforced before all of this code.
   -/
-  atomsAssignmentCache : Expr := mkConst ``List.nil [.zero]
+  atomsAssignmentCache : Expr := mkConst `illegal
 
 /--
 The reflection monad, used to track `BitVec` variables that we see as we traverse the context.
@@ -228,9 +231,9 @@ def run (m : M α) : MetaM α :=
 /--
 Retrieve the atoms as pairs of their width and expression.
 -/
-def atoms : M (List (Nat × Expr)) := do
+def atoms : M (Array (Nat × Expr)) := do
   let sortedAtoms := (← getThe State).atoms.toArray.qsort (·.2.atomNumber < ·.2.atomNumber)
-  return sortedAtoms.map (fun (expr, {width, ..}) => (width, expr)) |>.toList
+  return sortedAtoms.map (fun (expr, {width, ..}) => (width, expr))
 
 /--
 Retrieve a `BitVec.Assignment` representing the atoms we found so far.
@@ -257,11 +260,14 @@ def lookup (e : Expr) (width : Nat) (synthetic : Bool) : M Nat := do
 where
   updateAtomsAssignment : M Unit := do
     let as ← atoms
-    let packed :=
-      as.map (fun (width, expr) => mkApp2 (mkConst ``BVExpr.PackedBitVec.mk) (toExpr width) expr)
-    let packedType := mkConst ``BVExpr.PackedBitVec
-    let newAtomsAssignment ← mkListLit packedType packed
-    modify fun s => { s with atomsAssignmentCache := newAtomsAssignment }
+    if h : 0 < as.size then
+      let ras := Lean.RArray.ofArray as h
+      let packedType := mkConst ``BVExpr.PackedBitVec
+      let pack := fun (width, expr) => mkApp2 (mkConst ``BVExpr.PackedBitVec.mk) (toExpr width) expr
+      let newAtomsAssignment := ras.toExpr packedType pack
+      modify fun s => { s with atomsAssignmentCache := newAtomsAssignment }
+    else
+      throwError "updateAtomsAssignment should only be called when there is an atom"
 
 @[specialize]
 def simplifyBinaryProof' (mkFRefl : Expr → Expr) (fst : Expr) (fproof : Option Expr)
diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/SatAtBVLogical.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/SatAtBVLogical.lean
index 9e570abeb5d4..bceb1a0580c2 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/SatAtBVLogical.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/SatAtBVLogical.lean
@@ -61,13 +61,16 @@ def and (x y : SatAtBVLogical) : SatAtBVLogical where
 
 /-- Given a proof that `x.expr.Unsat`, produce a proof of `False`. -/
 def proveFalse (x : SatAtBVLogical) (h : Expr) : M Expr := do
-  let atomsList ← M.atomsAssignment
-  let evalExpr := mkApp2 (mkConst ``BVLogicalExpr.eval) atomsList x.expr
-  return mkApp3
-    (mkConst ``Std.Tactic.BVDecide.Reflect.Bool.false_of_eq_true_of_eq_false)
-    evalExpr
-    (← x.satAtAtoms)
-    (.app h atomsList)
+  if (← get).atoms.isEmpty then
+    throwError "Unable to identify any relevant atoms."
+  else
+    let atomsList ← M.atomsAssignment
+    let evalExpr := mkApp2 (mkConst ``BVLogicalExpr.eval) atomsList x.expr
+    return mkApp3
+      (mkConst ``Std.Tactic.BVDecide.Reflect.Bool.false_of_eq_true_of_eq_false)
+      evalExpr
+      (← x.satAtAtoms)
+      (.app h atomsList)
 
 
 end SatAtBVLogical
diff --git a/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Basic.lean b/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Basic.lean
index 22af1d1cb39f..f1ab2f404f62 100644
--- a/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Basic.lean
+++ b/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Basic.lean
@@ -6,6 +6,7 @@ Authors: Henrik Böving
 prelude
 import Init.Data.Hashable
 import Init.Data.BitVec
+import Init.Data.RArray
 import Std.Tactic.BVDecide.Bitblast.BoolExpr.Basic
 
 /-!
@@ -279,20 +280,20 @@ structure PackedBitVec where
 /--
 The notion of variable assignments for `BVExpr`.
 -/
-abbrev Assignment := List PackedBitVec
+abbrev Assignment := Lean.RArray PackedBitVec
 
 /--
 Get the value of a `BVExpr.var` from an `Assignment`.
 -/
-def Assignment.getD (assign : Assignment) (idx : Nat) : PackedBitVec :=
-  List.getD assign idx ⟨BitVec.zero 0⟩
+def Assignment.get (assign : Assignment) (idx : Nat) : PackedBitVec :=
+  Lean.RArray.get assign idx
 
 /--
 The semantics for `BVExpr`.
 -/
 def eval (assign : Assignment) : BVExpr w → BitVec w
   | .var idx =>
-    let ⟨bv⟩ := assign.getD idx
+    let ⟨bv⟩ := assign.get idx
     bv.truncate w
   | .const val => val
   | .zeroExtend v expr => BitVec.zeroExtend v (eval assign expr)
@@ -307,7 +308,7 @@ def eval (assign : Assignment) : BVExpr w → BitVec w
   | .arithShiftRight lhs rhs => BitVec.sshiftRight' (eval assign lhs) (eval assign rhs)
 
 @[simp]
-theorem eval_var : eval assign ((.var idx) : BVExpr w) = (assign.getD idx).bv.truncate _ := by
+theorem eval_var : eval assign ((.var idx) : BVExpr w) = (assign.get idx).bv.truncate _ := by
   rfl
 
 @[simp]
diff --git a/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Lemmas/Basic.lean b/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Lemmas/Basic.lean
index b57266bfc5d0..ecb877a42563 100644
--- a/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Lemmas/Basic.lean
+++ b/src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Lemmas/Basic.lean
@@ -18,11 +18,11 @@ namespace Std.Tactic.BVDecide
 namespace BVExpr
 
 def Assignment.toAIGAssignment (assign : Assignment) : BVBit → Bool :=
-  fun bit => (assign.getD bit.var).bv.getLsbD bit.idx
+  fun bit => (assign.get bit.var).bv.getLsbD bit.idx
 
 @[simp]
 theorem Assignment.toAIGAssignment_apply (assign : Assignment) (bit : BVBit) :
-    assign.toAIGAssignment bit = (assign.getD bit.var).bv.getLsbD bit.idx := by
+    assign.toAIGAssignment bit = (assign.get bit.var).bv.getLsbD bit.idx := by
   rfl
 
 end BVExpr