Skip to content

Commit

Permalink
perf: bv_decide use Lean.RArray
Browse files Browse the repository at this point in the history
  • Loading branch information
hargoniX committed Dec 2, 2024
1 parent 0b8f50f commit 32a8d7e
Show file tree
Hide file tree
Showing 4 changed files with 33 additions and 23 deletions.
24 changes: 15 additions & 9 deletions src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 6 additions & 5 deletions src/Std/Tactic/BVDecide/Bitblast/BVExpr/Basic.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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

/-!
Expand Down Expand Up @@ -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)
Expand All @@ -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]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 32a8d7e

Please sign in to comment.