From f1e485fe245a60324c9156a78b23155dcf8358e6 Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Tue, 8 Oct 2024 23:00:45 -0500
Subject: [PATCH 01/10] feat: Strict Ackermannization for bitvector problems

We implement strict ackermannization,
which is an algorithmic technique to convert QF_BV+UF into QF_BV.

The implementation walks over the goals and the hypotheses.
When it encounters a function application `(f x1 x2 ... xn)`
where the type signature of `f` is `BitVec k1 -> BitVec k2 -> ... -> BitVec kn -> BitVec ko`,
it creates a new variable `fAppX : BitVec k0`, and an equation `fAppX = f x1 ... xn`.
Next, when it encounters another application of the same function `(f y1 y2 ... yn)`,
it creates a new variable `fAppY : BitVec k0`, and another equation `fAppY = f y1 ... yn`.

After the walk, we loop over all pairs of applications of the function `f` that we have abstracted.
In this case, we have `fAppX` and `fAppY`. For these, we build a extentionality equation, which says that if

```
hAppXAppYExt: x1 = x2 -> y1 = y2 -> ... xn = yn -> fAppX = fAppY
```

Intuitively, this is the only fact that they theory UF can propagate,
and we thus encode it directly as a SAT formula,
by abstracting out the actual function application into a new free variable.

This implementation creates the ackermann variables (`fAppX, fAppY`) first,
and then in a subsequent phase, adds all the ackermann equations (`hAppXAppYExt`).
This anticipates the next patches which will implement *incremental* ackermannization,
where we will use the counterexample from the model to selectively add ackerman equations.
---
 src/Lean/Elab/Tactic.lean                 |   1 +
 src/Lean/Elab/Tactic/BVAckermannize.lean  | 438 ++++++++++++++++++++++
 src/Std/Tactic/BVAckermannize.lean        |  12 +
 src/Std/Tactic/BVAckermannize/Syntax.lean |  31 ++
 tests/lean/run/ackermannize.lean          |  43 +++
 5 files changed, 525 insertions(+)
 create mode 100644 src/Lean/Elab/Tactic/BVAckermannize.lean
 create mode 100644 src/Std/Tactic/BVAckermannize.lean
 create mode 100644 src/Std/Tactic/BVAckermannize/Syntax.lean
 create mode 100644 tests/lean/run/ackermannize.lean

diff --git a/src/Lean/Elab/Tactic.lean b/src/Lean/Elab/Tactic.lean
index b8faf9a1749d..e2a75219ff88 100644
--- a/src/Lean/Elab/Tactic.lean
+++ b/src/Lean/Elab/Tactic.lean
@@ -43,3 +43,4 @@ import Lean.Elab.Tactic.Rewrites
 import Lean.Elab.Tactic.DiscrTreeKey
 import Lean.Elab.Tactic.BVDecide
 import Lean.Elab.Tactic.BoolToPropSimps
+import Lean.Elab.Tactic.BVAckermannize
diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
new file mode 100644
index 000000000000..5459c9055211
--- /dev/null
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -0,0 +1,438 @@
+/-
+Copyright (c) 2024 Siddharth Bhat. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Siddharth Bhat
+
+This file implements lazy ackermannization [1, 2]
+
+[1] https://lara.epfl.ch/w/_media/model-based.pdf
+[2]  https://leodemoura.github.io/files/oregon08.pdf
+[3] https://github.com/Z3Prover/z3/blob/d047b86439ec209446d211f0f6b251ebfba070d8/src/ackermannization/lackr.cpp#L206
+[4]https://github.com/Z3Prover/z3/blob/d047b86439ec209446d211f0f6b251ebfba070d8/src/ackermannization/lackr_model_constructor.cpp#L344
+-/
+prelude
+import Lean.Expr
+import Lean.Message
+import Std.Tactic.BVDecide.Bitblast
+import Std.Tactic.BVAckermannize.Syntax
+import Lean.Meta.Basic
+import Lean.Elab.Tactic.Basic
+import Lean.Meta.LitValues
+import Lean.Meta.InferType
+import Lean.Elab.Tactic.FalseOrByContra
+import Lean.Meta.Tactic.Assert
+import Lean.Meta.Tactic.Util
+import Lean.Meta.Tactic.Congr
+import Lean.Meta.Tactic.Replace
+
+open Lean Elab Meta Tactic
+
+initialize Lean.registerTraceClass `bv_ack
+
+namespace Ack
+
+structure Config where
+
+structure Context extends Config where
+ config : Config
+
+ /-- Types that can be bitblasted by bv_decide -/
+ inductive BVTy
+ /-- booleans -/
+ | Bool
+ /-- Bitvectors of a fixed width `w` -/
+ | BitVec (w : Nat)
+ deriving Hashable, DecidableEq, Inhabited
+
+instance : ToMessageData BVTy where
+  toMessageData 
+  | .Bool => m!"bool"
+  | .BitVec w => m!"BitVec {w}"
+
+namespace BVTy
+
+ /-- info: _root_.BitVec (w : Nat) : Type -/
+ #guard_msgs in #check _root_.BitVec
+
+/-- Reify a raw expression into the types of bitvectors we can bitblast -/
+def ofExpr? (e : Expr) : OptionT MetaM BVTy :=
+  match_expr e.consumeMData with
+  | _root_.Bool => return Bool
+  | _root_.BitVec w => do
+     let w ← getNatValue? w
+     return .BitVec w
+  | _ => OptionT.fail
+
+def toExpr : BVTy → Expr
+| .Bool => mkConst ``_root_.Bool
+| .BitVec w => mkApp (mkConst ``_root_.BitVec) (mkNatLit w)
+
+end BVTy
+
+structure Argument where
+  /-- The expression corresponding to the argument -/
+  x : Expr
+  /-- The cached type of the expression x -/
+  xTy : BVTy
+deriving Hashable, BEq, Inhabited
+
+namespace Argument
+
+instance : ToMessageData Argument where
+  toMessageData arg := m!"{arg.x} : {arg.xTy}"
+
+/-- Build an `Argument` from a raw expression -/
+def ofExpr? (e : Expr) : OptionT MetaM Argument := do
+  let t ← BVTy.ofExpr? (← inferType e)
+  return { x := e, xTy := t}
+
+end Argument
+
+
+structure Function where
+  /-- The function -/
+  f : Expr
+  codTy : BVTy
+ deriving Hashable, BEq, Inhabited
+
+
+namespace Function
+
+instance : ToMessageData Function where
+  toMessageData fn := m!"{fn.f} (cod: {fn.codTy})"
+
+/--
+Reify an expression `e` of the kind 
+`f x₁ ... xₙ`, where all the arguments and the return type are a `BVTy` into an Ap
+-/
+def reifyAp (f : Expr) : OptionT MetaM (Function × Array Argument) := do
+  let xs := f.getAppArgs
+  /- We need at least one argument for this to be a function call we can ackermannize -/
+  guard <| xs.size > 0
+  let codTy ← BVTy.ofExpr? (← inferType f)
+  let args ← xs.mapM Argument.ofExpr?
+  let fn : Function := { f, codTy }
+  return (fn, args)
+end Function
+
+
+/--
+TODO: is it sensible to hash an array of arguments?
+We may want to use something like a trie to index these.
+Consider swiching to something like `Trie`.
+-/
+abbrev ArgumentList := Array Argument
+
+
+/--
+The data stored for an ackermannized call to allow us to build proofs.
+-/
+structure CallVal where
+  /-- the free variable `ack_fx₁...xₙ := (f x₁ x₂ ... xₙ)`. -/
+  fvar : FVarId 
+  /-- heqProof : The proof that `ack_fx₁...fxₙ = f x₁ x₂ ... xₙ` (name/fvar = value/expr). -/
+  heqProof : Expr  
+deriving Hashable, BEq, Inhabited
+
+namespace CallVal
+instance : ToMessageData CallVal where
+  toMessageData cv := m!"{Expr.fvar cv.fvar} ({cv.heqProof})"
+end CallVal
+
+
+structure State where
+  /--
+  A maping from a function `f`, to a map of arguments `x₁ ... xₙ`, to the information stored about the call.
+  This is used to generate equations of the form `x₁ = y₁ → x₂ = y₂ → ... → xₙ = yₙ → ack_fx₁...xₙ = ack_fy₁...yₙ on-demand.
+  -/
+  fn2args2call : Std.HashMap Function (Std.HashMap ArgumentList CallVal) := {}
+  /-- A counter for generating fresh names. -/
+  gensymCounter : Nat := 0
+
+
+def State.init (_cfg : Config) : State where
+
+abbrev AckM := StateRefT State (ReaderT Context MetaM)
+
+namespace AckM
+
+def run (m : AckM α) (ctx : Context) : MetaM α :=
+  m.run' (State.init ctx.config) |>.run ctx
+
+/-- Generate a fresh name. -/
+def gensym : AckM Name := do
+  modify fun s => { s with gensymCounter := s.gensymCounter + 1 }
+  return Name.mkNum (Name.mkSimple "ackConst") (← get).gensymCounter
+
+def withContext (g : MVarId) (ma : AckM α) : AckM α := g.withContext ma
+
+/-- Get the calls to a function `fn`. -/
+def getCallMap (fn : Function) : AckM (Std.HashMap ArgumentList CallVal) := do
+  return (← get).fn2args2call.getD fn {}
+
+/-- Get the calls to a function `fn`. -/
+def getCallVal? (fn : Function) (args : Array Argument) : AckM (Option CallVal) := do
+  let calls ← getCallMap fn
+  if let .some val := calls.get? args then
+    return some val
+  return none
+
+structure IntroDefResult where
+  -- the new name/fvar of the defn.
+  defn : FVarId
+  -- a proof 'hdefn : name = value
+  eqProof : FVarId
+  
+/-
+Introduce a new definition with name `name : hdefTy` into the local context,
+and return the FVarId of the new definition in the new goal (the MVarId) returned.
+-/
+def introDefExt (g : MVarId) (name : Name) (hdefTy : Expr) (hdefVal : Expr) : AckM (IntroDefResult × MVarId) := do
+  withContext g do
+    let g ← g.assertExt name (hName := Name.str name "value") hdefTy hdefVal
+    let (defn, g) ← g.intro1P
+    let (eqProof, g) ← g.intro1P
+    return ({ defn, eqProof}, g)
+
+/-- Insert the CallVal `cv` at `(fn, args)` into the state. -/
+private def _insertCallVal (fn : Function) (args : ArgumentList) (cv : CallVal) : AckM Unit := do
+  let calls ← getCallMap fn
+  modify fun s => { s with fn2args2call := s.fn2args2call.insert fn (calls.insert args cv) }
+      
+
+/--
+Replace a call to the function `f` with an `fvar`. Since the `fvar` is defeq to the call,
+we can just replace occurrences of the call with the fvar `f`.
+
+We will later need to add another hypothesis with the equality that the `fvar = f x₁ ... xₙ`
+-/
+
+def replaceCallWithFVar (g : MVarId) (fn : Function) (args : ArgumentList) : AckM (CallVal × MVarId) := do
+  g.withContext do
+    if let some val ← getCallVal? fn args then 
+      trace[bv_ack] "using cached call {val} for {fn} {args}"
+      return (val, g)
+    let e := mkAppN fn.f (args.map Argument.x)
+    let name ← gensym
+    let (introDef, g) ← introDefExt g name fn.codTy.toExpr e
+    let cv := { fvar := introDef.defn, heqProof := Expr.fvar introDef.eqProof : CallVal }
+    _insertCallVal fn args cv
+    return (cv, g)
+
+/-- create a trace node in trace class (i.e. `set_option traceClass true`),
+with header `header`, whose default collapsed state is `collapsed`. -/
+def withTraceNode (header : MessageData) (k : AckM α)
+    (collapsed : Bool := true)
+    (traceClass : Name := `bv_ack) : AckM α :=
+  Lean.withTraceNode traceClass (fun _ => return header) k (collapsed := collapsed)
+
+/-- An emoji used to report intemediate states where the tactic is processing hypotheses. -/
+def processingEmoji : String := "⚙️"
+
+/--
+Create a trace note that folds `header` with `(NOTE: can be large)`,
+and prints `msg` under such a trace node.
+Used to print goal states, which can be quite noisy in the trace.
+-/
+def traceLargeMsg (header : MessageData) (msg : MessageData) : AckM Unit :=
+    withTraceNode m!"{header} (NOTE: can be large)" do
+      trace[ack] msg
+
+
+/-- Returns `True` if the type is a function type that is understood by the bitblaster. -/
+def isBitblastTy (e : Expr) : Bool :=
+  match_expr e with
+  | BitVec _ => true
+  | Bool => true
+  | _ => false
+
+mutual
+partial def ackAppChildren (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := do
+  g.withContext do
+    trace[bv_ack] "{crossEmoji} bailing out on illegal application {e}"
+    let f := e.getAppFn
+    let (f, g) ← introAckForExpr g f
+    -- NOTE: Rewrite as a fold? It's not too much cleaner, sadly.
+    let mut args := #[]
+    let mut g := g
+    for arg in e.getAppArgs do
+      let gArg ← introAckForExpr g arg
+      g := gArg.2
+      args := args.push gArg.1
+    return (mkAppN f args, g) -- NOTE: is there some way to use update to update this?
+    
+
+partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := do
+    withTraceNode m!"🪵 {e}" do
+    match e with
+    | .mdata _ e => introAckForExpr g e
+    | .bvar .. | .fvar .. | .mvar .. | .sort .. | .const .. | .proj .. | .lit .. => return (e, g)
+    | .lam _binderName binderTy body _binderInfo =>
+        let (binderTy, g) ← introAckForExpr g binderTy
+        let (body, g) ← introAckForExpr g body
+        return (e.updateLambdaE! binderTy body, g)
+    | .letE _declName type value body _nonDep =>      
+        let (type, g) ← introAckForExpr g type
+        let (value, g) ← introAckForExpr g value
+        let (body, g) ← introAckForExpr g body
+        return (e.updateLet! type value body, g)
+    | .forallE _binderName binderTy body _binderInfo => 
+        let (binderTy, g) ← introAckForExpr g binderTy
+        let (body, g) ← introAckForExpr g body
+        return (e.updateForallE! binderTy body, g)
+    | .app .. => do 
+      withTraceNode m!"@ Expr.app '{e}'" do
+        let f := e.getAppFn
+        let te ← inferType e
+        let .some codTy ← BVTy.ofExpr? te |>.run
+          | do
+            trace[bv_ack] "{crossEmoji} '{te}' not BitVec/Bool."
+            return (← ackAppChildren g e)
+        trace[bv_ack] "{checkEmoji} {e}'s codomain '{te}'"
+
+        let fn := { f, codTy : Function }
+
+        let args := e.getAppArgs
+        assert! args.size > 0 -- since we are an application, we must have at least one argument.
+         -- run ackermannization on all the arguments.
+        -- This ensures that we process bottom-up, and at this stage, our argument
+        -- has been ackermannized, if such an opportunity exists.
+        let mut ackArgs := #[]
+        for arg in args do
+          trace[bv_ack] "@ arg {arg}"
+          let (arg, g) ← introAckForExpr g arg
+          -- do I need a `withContext` here? :(
+          if let .some ackArg ← Argument.ofExpr? arg |>.run then
+            trace[bv_ack] "{checkEmoji} arg {arg}"
+            ackArgs := ackArgs.push ackArg
+          else
+            -- bail out, and recurse into the `app` since we can't ackermannize it.
+            trace[bv_ack] "{crossEmoji} {arg}"
+            return (← ackAppChildren g e)
+           
+        let (call, g) ← replaceCallWithFVar g fn ackArgs
+        trace[bv_ack] "{e} → {call}."
+        return (Expr.fvar call.fvar, g)
+end
+
+
+/--
+Return true if the argument lists are trivially different.
+This is an optimization that we do not yet implement.
+-/
+def areArgListsTriviallyDifferent (_arg₁ _arg₂ : Array Argument) : AckM Bool := return false
+
+/-
+Return true if the argument lists are trivially the same.
+This is an optimization that we do not yet implement.
+If possible, return the simplified hypothesis of the equality of these expressions.
+TODO: -- def areArgListsTriviallySame (arg₁ arg₂ : Array Argument) : AckM (Option Expr) := return none
+-/
+
+
+/-- info: congr.{u, v} {α : Sort u} {β : Sort v} {f₁ f₂ : α → β} {a₁ a₂ : α} (h₁ : f₁ = f₂) (h₂ : a₁ = a₂) : f₁ a₁ = f₂ a₂ -/
+#guard_msgs in #check congr
+
+/-- This is how we build the congruence lemma for `n` arguments -/
+private example (f : X1 → X2 → O) 
+    -- we have these already.
+    (x1 x1' : X1) (x2 x2' : X2)
+    (ackEqApp : fx1x2 = f x1 x2)
+    (ackEqApp' : fx1'x2' = f x1' x2') 
+    -- to be intros'd
+    (h1 : x1 = x1')
+    (h2 : x2 = x2') :
+  fx1x2 = fx1'x2' := 
+  let appEqApp : f x1 x2 = f x1' x2' := congr (congr (Eq.refl f) h1) h2
+  Eq.trans (Eq.trans ackEqApp appEqApp) (Eq.symm ackEqApp')
+
+
+/-
+Make the ackermannization theorem, which states that: `(∀ i, arg₁[i] = arg₂[i]) -> call₁ = call₂`.
+Formally, we build an expr such as `arg₁ = arg'₁ -> arg₂ = arg'₂ -> ... argₙ = arg'ₙ -> call₁ = call₂`,
+where the proof is by congruence over the equalities.
+-/
+def mkAckThm (g : MVarId) (fn : Function) (args args' : Array Argument) (call call' : CallVal): AckM MVarId := do
+    trace[bv_ack] "making ack congr thm for '{fn}' '{args}' ~ '{args'}',  calls '{call}', '{call'}'"
+    if args.size = 0 then
+      throwTacticEx `bv_ack g
+        m!"expected {args} to have more than zero arguments when building congr theorem for {fn}."
+    if args'.size = 0 then
+      throwTacticEx `bv_ack g
+        m!"expected {args'} to have more than zero arguments when building congr theorem for {fn}."
+
+    if args.size ≠ args'.size then
+      throwTacticEx `bv_ack g 
+        m!"internal error: expected {args} to have the same size as {args'} when building congr thm for {fn}."
+    let mut eqHyps := #[]
+    for (arg, arg') in args.zip args' do
+      eqHyps := eqHyps.push (← mkEq arg.x arg'.x)
+    -- we build the equality according to the example above.
+    let mut localDecls : Array (Name × BinderInfo × (Array Expr → AckM Expr)) := #[]
+    let mut i := 0
+    for (arg, arg') in args.zip args' do
+      let name := Name.num (Name.mkSimple "ack_arg") i
+      localDecls := localDecls.push (name, BinderInfo.default, fun _ => mkEq arg.x arg'.x)
+    let ackEqn ← withLocalDecls localDecls fun argsEq => do 
+      let mut fEq ← mkEqRefl fn.f
+      for argEq in argsEq do
+        fEq ← mkCongr fEq argEq
+      -- now fEq  ~ appEqApp
+      let finalEq ←  mkEqTrans (← mkEqTrans call.heqProof fEq) (← mkEqSymm call'.heqProof)
+      mkLambdaFVars argsEq  finalEq
+    trace[bv_ack] "made ackermann equation: {ackEqn}"
+    let (_ackEqn, g) ← g.note (Name.mkSimple s!"ackEqn{fn.f}") ackEqn
+    return g
+
+/-
+For every bitvector (x : BitVec w), for every function `(f : BitVec w → BitVec w')`,
+walk every function application (f x), and add a new fvar (fx : BitVec w').
+- Keep an equality that says `fx = f x`.
+- For function application of f, for each pair of bitvectors x, y,
+  add a hypothesis that says `x = y => fx = fy, with proof.
+-/
+def ack (g : MVarId) : AckM MVarId := do
+  g.withContext do
+    let (target', g) ← introAckForExpr g (← inferType (Expr.mvar g))
+    let g ← g.replaceTargetDefEq target'
+
+    let hyps ← g.getNondepPropHyps
+    let mut g := g
+    for hyp in hyps do
+      g ← g.withContext do
+        withTraceNode m!"@ hyp '{← hyp.getType}'" do
+          let hypG ← introAckForExpr g (← hyp.getType)
+          pure hypG.2
+
+    -- trace[bv_ack] "done with ackermannization collection, now adding new theorems..."
+    for (fn, arg2call) in (← get).fn2args2call do
+      let argCallsArr := arg2call.toArray
+      for i in [0:argCallsArr.size] do
+        let (arg₁, call₁) := argCallsArr[i]!
+        for j in [i+1:argCallsArr.size] do
+          let (arg₂, call₂) := argCallsArr[j]!
+          if ← areArgListsTriviallyDifferent arg₁ arg₂ then
+            continue
+          g ← withContext g do
+                mkAckThm g fn arg₁ arg₂ call₁ call₂
+    trace[bv_ack] "{checkEmoji} ack.{indentD g}"
+    return g
+
+end AckM
+
+/-- Entry point for programmatic usage of `bv_ackermannize` -/
+def ackTac (ctx : Context) : TacticM Unit := do
+  withoutRecover do
+    liftMetaTactic fun g => do
+      let g ← (AckM.ack g).run ctx
+      return [g]
+
+end Ack
+
+@[builtin_tactic Lean.Parser.Tactic.bvAckEager]
+def evalBvAckEager : Tactic := fun
+  | `(tactic| bv_ack_eager) => 
+    let config : Ack.Config := {}
+    let ctx : Ack.Context := { config := config }
+    Ack.ackTac ctx
+  | _ => throwUnsupportedSyntax
diff --git a/src/Std/Tactic/BVAckermannize.lean b/src/Std/Tactic/BVAckermannize.lean
new file mode 100644
index 000000000000..5e555a161137
--- /dev/null
+++ b/src/Std/Tactic/BVAckermannize.lean
@@ -0,0 +1,12 @@
+/-
+Copyright (c) 2024 Siddharth Bhat. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Siddharth Bhat
+-/
+prelude
+import Std.Tactic.BVAckermannize.Syntax
+
+/-!
+This directory contains the lazy ackermannization tactic.
+This uses lean's inbuilt bitblaster to uninterpreted functions + bitvectors to SAT.
+-/
diff --git a/src/Std/Tactic/BVAckermannize/Syntax.lean b/src/Std/Tactic/BVAckermannize/Syntax.lean
new file mode 100644
index 000000000000..f91e0a794b27
--- /dev/null
+++ b/src/Std/Tactic/BVAckermannize/Syntax.lean
@@ -0,0 +1,31 @@
+/-
+Copyright (c) 2024 Siddharth Bhat. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Siddharth Bhat
+-/
+prelude
+import Init.Notation
+import Init.Simproc
+
+set_option linter.missingDocs true -- keep it documented
+
+namespace Lean.Parser
+
+namespace Tactic
+
+/--
+Close fixed-width `BitVec`, `Bool`, and uninterpreted function goals by obtaining a proof from an external SAT solver and
+verifying it inside Lean. The solvable goals are currently limited to the Lean equivalent of
+[`QF_BV`](https://smt-lib.org/logics-all.shtml#QF_BV) + [`QF_UF`](https://smt-lib.org/logics-all.shtml#QF_UF)
+
+```lean
+example : ∀ (a b : BitVec 64), (a &&& b) + (a ^^^ b) = a ||| b := by
+  intros
+  bv_decide
+```
+-/
+syntax (name := bvAckEager) "bv_ack_eager" : tactic
+
+end Tactic
+
+end Lean.Parser
diff --git a/tests/lean/run/ackermannize.lean b/tests/lean/run/ackermannize.lean
new file mode 100644
index 000000000000..df8b2c913ae1
--- /dev/null
+++ b/tests/lean/run/ackermannize.lean
@@ -0,0 +1,43 @@
+import Std.Tactic.BVAckermannize
+import Std.Tactic.BVDecide
+import Lean.Elab.Tactic.BVAckermannize
+
+set_option trace.bv_ack true
+
+theorem foo (f : BitVec 1 → BitVec 1) (x : BitVec 1) : 
+    ((1#1 ^^^ f x ^^^ (f (x + 1))) = 0#1) → 
+    ((f 0#1 = 1#1) ∨ (f 1#1 = 1#1)) := by
+  try bv_decide
+  bv_ack_eager
+  bv_decide
+
+/-- info: 'foo' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound] -/
+#guard_msgs in  #print axioms foo
+
+theorem bar (f : BitVec 1 -> BitVec 1) (x y : BitVec 1)
+   (hfxy : f x = 1#1 ∨ f y = 1#1)
+   (hxy : x ^^^ y = 0#1) : 
+   (f 0#1 = 1#1 ∨ f 1#1 = 1#1) := by
+  try bv_decide
+  bv_ack_eager
+  bv_decide
+ 
+/-- info: 'bar' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound] -/
+#guard_msgs in  #print axioms bar
+
+-- set_option trace.ack true in
+theorem less_6E (m : BitVec 64) (P : BitVec 64 → Bool)
+    (h : m ≤ 0x00000006 ∧
+         P 0x00000000 ∧
+         P 0x00000001 ∧
+         P 0x00000002 ∧
+         P 0x00000003 ∧
+         P 0x00000004 ∧
+         P 0x00000005 ∧
+         P 0x00000006) :
+    P m := by
+  bv_ack_eager
+  bv_decide
+
+/-- info: 'less_6E' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound] -/
+#guard_msgs in  #print axioms less_6E

From 9bc7719b5e70b1e67381d6cadf3f9f85a176c0ef Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Tue, 8 Oct 2024 23:17:12 -0500
Subject: [PATCH 02/10] feat: add more tests

---
 tests/lean/run/ackermannize.lean | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tests/lean/run/ackermannize.lean b/tests/lean/run/ackermannize.lean
index df8b2c913ae1..38188e2e589c 100644
--- a/tests/lean/run/ackermannize.lean
+++ b/tests/lean/run/ackermannize.lean
@@ -4,6 +4,7 @@ import Lean.Elab.Tactic.BVAckermannize
 
 set_option trace.bv_ack true
 
+/- Test that we correctly handle arrow / forall -/
 theorem foo (f : BitVec 1 → BitVec 1) (x : BitVec 1) : 
     ((1#1 ^^^ f x ^^^ (f (x + 1))) = 0#1) → 
     ((f 0#1 = 1#1) ∨ (f 1#1 = 1#1)) := by
@@ -14,6 +15,7 @@ theorem foo (f : BitVec 1 → BitVec 1) (x : BitVec 1) :
 /-- info: 'foo' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound] -/
 #guard_msgs in  #print axioms foo
 
+/- Test that we correctly ackermannize hyps and goal -/
 theorem bar (f : BitVec 1 -> BitVec 1) (x y : BitVec 1)
    (hfxy : f x = 1#1 ∨ f y = 1#1)
    (hxy : x ^^^ y = 0#1) : 
@@ -25,6 +27,7 @@ theorem bar (f : BitVec 1 -> BitVec 1) (x y : BitVec 1)
 /-- info: 'bar' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound] -/
 #guard_msgs in  #print axioms bar
 
+/- Test that the example that motivated ackermannization works -/
 -- set_option trace.ack true in
 theorem less_6E (m : BitVec 64) (P : BitVec 64 → Bool)
     (h : m ≤ 0x00000006 ∧
@@ -41,3 +44,13 @@ theorem less_6E (m : BitVec 64) (P : BitVec 64 → Bool)
 
 /-- info: 'less_6E' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound] -/
 #guard_msgs in  #print axioms less_6E
+
+/- Test that nested applications work -/
+theorem true_on_one_input_of_neg_comm (g : BitVec 1 → BitVec 1) (x : BitVec 1) (h : g (~~~ x) = ~~~ (g x)) : 
+    (g 0#1) ||| (g 1#1) = 1#1 := by
+  try bv_decide
+  bv_ack_eager
+  bv_decide
+
+/-- info: 'true_on_one_input_of_neg_comm' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound] -/
+#guard_msgs in  #print axioms true_on_one_input_of_neg_comm

From 9b0e32794aaa431191130d6435b3d527eb563f99 Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Tue, 8 Oct 2024 23:39:39 -0500
Subject: [PATCH 03/10] chore: cleanup

---
 src/Lean/Elab/Tactic/BVAckermannize.lean | 25 ++++++------------------
 1 file changed, 6 insertions(+), 19 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
index 5459c9055211..d165830d581a 100644
--- a/src/Lean/Elab/Tactic/BVAckermannize.lean
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -226,19 +226,6 @@ def withTraceNode (header : MessageData) (k : AckM α)
     (traceClass : Name := `bv_ack) : AckM α :=
   Lean.withTraceNode traceClass (fun _ => return header) k (collapsed := collapsed)
 
-/-- An emoji used to report intemediate states where the tactic is processing hypotheses. -/
-def processingEmoji : String := "⚙️"
-
-/--
-Create a trace note that folds `header` with `(NOTE: can be large)`,
-and prints `msg` under such a trace node.
-Used to print goal states, which can be quite noisy in the trace.
--/
-def traceLargeMsg (header : MessageData) (msg : MessageData) : AckM Unit :=
-    withTraceNode m!"{header} (NOTE: can be large)" do
-      trace[ack] msg
-
-
 /-- Returns `True` if the type is a function type that is understood by the bitblaster. -/
 def isBitblastTy (e : Expr) : Bool :=
   match_expr e with
@@ -263,7 +250,7 @@ partial def ackAppChildren (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := do
     
 
 partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := do
-    withTraceNode m!"🪵 {e}" do
+  Lean.withTraceNode `bv_ack (fun _ => g.withContext do pure m!"🎯 {e}") (collapsed := false) do
     match e with
     | .mdata _ e => introAckForExpr g e
     | .bvar .. | .fvar .. | .mvar .. | .sort .. | .const .. | .proj .. | .lit .. => return (e, g)
@@ -281,7 +268,7 @@ partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := d
         let (body, g) ← introAckForExpr g body
         return (e.updateForallE! binderTy body, g)
     | .app .. => do 
-      withTraceNode m!"@ Expr.app '{e}'" do
+      withTraceNode m!"🎯 Expr.app '{e}'" (collapsed := false) do
         let f := e.getAppFn
         let te ← inferType e
         let .some codTy ← BVTy.ofExpr? te |>.run
@@ -299,7 +286,7 @@ partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := d
         -- has been ackermannized, if such an opportunity exists.
         let mut ackArgs := #[]
         for arg in args do
-          trace[bv_ack] "@ arg {arg}"
+          trace[bv_ack] "🎯 arg {arg}"
           let (arg, g) ← introAckForExpr g arg
           -- do I need a `withContext` here? :(
           if let .some ackArg ← Argument.ofExpr? arg |>.run then
@@ -311,7 +298,7 @@ partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := d
             return (← ackAppChildren g e)
            
         let (call, g) ← replaceCallWithFVar g fn ackArgs
-        trace[bv_ack] "{e} → {call}."
+        trace[bv_ack] "{checkEmoji} {e} → {call}."
         return (Expr.fvar call.fvar, g)
 end
 
@@ -352,7 +339,7 @@ Make the ackermannization theorem, which states that: `(∀ i, arg₁[i] = arg
 Formally, we build an expr such as `arg₁ = arg'₁ -> arg₂ = arg'₂ -> ... argₙ = arg'ₙ -> call₁ = call₂`,
 where the proof is by congruence over the equalities.
 -/
-def mkAckThm (g : MVarId) (fn : Function) (args args' : Array Argument) (call call' : CallVal): AckM MVarId := do
+def mkAckThm (g : MVarId) (fn : Function) (args args' : Array Argument) (call call' : CallVal) : AckM MVarId := do
     trace[bv_ack] "making ack congr thm for '{fn}' '{args}' ~ '{args'}',  calls '{call}', '{call'}'"
     if args.size = 0 then
       throwTacticEx `bv_ack g
@@ -400,7 +387,7 @@ def ack (g : MVarId) : AckM MVarId := do
     let mut g := g
     for hyp in hyps do
       g ← g.withContext do
-        withTraceNode m!"@ hyp '{← hyp.getType}'" do
+        withTraceNode m!"🎯 hyp '{← hyp.getType}'" (collapsed := false) do
           let hypG ← introAckForExpr g (← hyp.getType)
           pure hypG.2
 

From 4337495ef3e02fa912805109eb4b8906f44b3165 Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Tue, 8 Oct 2024 23:46:09 -0500
Subject: [PATCH 04/10] chore: localize withContext

---
 src/Lean/Elab/Tactic/BVAckermannize.lean | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
index d165830d581a..7432130865e8 100644
--- a/src/Lean/Elab/Tactic/BVAckermannize.lean
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -6,9 +6,9 @@ Authors: Siddharth Bhat
 This file implements lazy ackermannization [1, 2]
 
 [1] https://lara.epfl.ch/w/_media/model-based.pdf
-[2]  https://leodemoura.github.io/files/oregon08.pdf
+[2] https://leodemoura.github.io/files/oregon08.pdf
 [3] https://github.com/Z3Prover/z3/blob/d047b86439ec209446d211f0f6b251ebfba070d8/src/ackermannization/lackr.cpp#L206
-[4]https://github.com/Z3Prover/z3/blob/d047b86439ec209446d211f0f6b251ebfba070d8/src/ackermannization/lackr_model_constructor.cpp#L344
+[4] https://github.com/Z3Prover/z3/blob/d047b86439ec209446d211f0f6b251ebfba070d8/src/ackermannization/lackr_model_constructor.cpp#L344
 -/
 prelude
 import Lean.Expr
@@ -117,7 +117,7 @@ end Function
 
 
 /--
-TODO: is it sensible to hash an array of arguments?
+NOTE: is it sensible to hash an array of arguments?
 We may want to use something like a trie to index these.
 Consider swiching to something like `Trie`.
 -/
@@ -135,8 +135,10 @@ structure CallVal where
 deriving Hashable, BEq, Inhabited
 
 namespace CallVal
+
 instance : ToMessageData CallVal where
   toMessageData cv := m!"{Expr.fvar cv.fvar} ({cv.heqProof})"
+
 end CallVal
 
 
@@ -351,6 +353,7 @@ def mkAckThm (g : MVarId) (fn : Function) (args args' : Array Argument) (call ca
     if args.size ≠ args'.size then
       throwTacticEx `bv_ack g 
         m!"internal error: expected {args} to have the same size as {args'} when building congr thm for {fn}."
+
     let mut eqHyps := #[]
     for (arg, arg') in args.zip args' do
       eqHyps := eqHyps.push (← mkEq arg.x arg'.x)
@@ -360,12 +363,12 @@ def mkAckThm (g : MVarId) (fn : Function) (args args' : Array Argument) (call ca
     for (arg, arg') in args.zip args' do
       let name := Name.num (Name.mkSimple "ack_arg") i
       localDecls := localDecls.push (name, BinderInfo.default, fun _ => mkEq arg.x arg'.x)
-    let ackEqn ← withLocalDecls localDecls fun argsEq => do 
+    let ackEqn ← g.withContext <| withLocalDecls localDecls fun argsEq => do 
       let mut fEq ← mkEqRefl fn.f
       for argEq in argsEq do
         fEq ← mkCongr fEq argEq
-      -- now fEq  ~ appEqApp
-      let finalEq ←  mkEqTrans (← mkEqTrans call.heqProof fEq) (← mkEqSymm call'.heqProof)
+
+      let finalEq ← mkEqTrans (← mkEqTrans call.heqProof fEq) (← mkEqSymm call'.heqProof)
       mkLambdaFVars argsEq  finalEq
     trace[bv_ack] "made ackermann equation: {ackEqn}"
     let (_ackEqn, g) ← g.note (Name.mkSimple s!"ackEqn{fn.f}") ackEqn
@@ -391,7 +394,6 @@ def ack (g : MVarId) : AckM MVarId := do
           let hypG ← introAckForExpr g (← hyp.getType)
           pure hypG.2
 
-    -- trace[bv_ack] "done with ackermannization collection, now adding new theorems..."
     for (fn, arg2call) in (← get).fn2args2call do
       let argCallsArr := arg2call.toArray
       for i in [0:argCallsArr.size] do
@@ -400,8 +402,7 @@ def ack (g : MVarId) : AckM MVarId := do
           let (arg₂, call₂) := argCallsArr[j]!
           if ← areArgListsTriviallyDifferent arg₁ arg₂ then
             continue
-          g ← withContext g do
-                mkAckThm g fn arg₁ arg₂ call₁ call₂
+          g ← mkAckThm g fn arg₁ arg₂ call₁ call₂
     trace[bv_ack] "{checkEmoji} ack.{indentD g}"
     return g
 

From 07fee4245d5a65bb01be67c42e703ce797b0e340 Mon Sep 17 00:00:00 2001
From: Siddharth <siddu.druid@gmail.com>
Date: Wed, 9 Oct 2024 00:27:51 -0500
Subject: [PATCH 05/10] Apply suggestions from code review

Co-authored-by: Tobias Grosser <github@grosser.es>
---
 src/Lean/Elab/Tactic/BVAckermannize.lean | 35 ++++++------------------
 1 file changed, 9 insertions(+), 26 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
index 7432130865e8..c5ea4b56c2c8 100644
--- a/src/Lean/Elab/Tactic/BVAckermannize.lean
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -38,7 +38,7 @@ structure Context extends Config where
 
  /-- Types that can be bitblasted by bv_decide -/
  inductive BVTy
- /-- booleans -/
+ /-- Booleans -/
  | Bool
  /-- Bitvectors of a fixed width `w` -/
  | BitVec (w : Nat)
@@ -81,54 +81,46 @@ namespace Argument
 instance : ToMessageData Argument where
   toMessageData arg := m!"{arg.x} : {arg.xTy}"
 
-/-- Build an `Argument` from a raw expression -/
+/-- Build an `Argument` from a raw expression. -/
 def ofExpr? (e : Expr) : OptionT MetaM Argument := do
   let t ← BVTy.ofExpr? (← inferType e)
   return { x := e, xTy := t}
 
 end Argument
-
-
 structure Function where
   /-- The function -/
   f : Expr
   codTy : BVTy
  deriving Hashable, BEq, Inhabited
-
-
 namespace Function
 
 instance : ToMessageData Function where
   toMessageData fn := m!"{fn.f} (cod: {fn.codTy})"
 
 /--
-Reify an expression `e` of the kind 
-`f x₁ ... xₙ`, where all the arguments and the return type are a `BVTy` into an Ap
+Reify an expression `e` of the kind `f x₁ ... xₙ`, where all the arguments and the return type are
+a `BVTy` into an App.
 -/
 def reifyAp (f : Expr) : OptionT MetaM (Function × Array Argument) := do
   let xs := f.getAppArgs
-  /- We need at least one argument for this to be a function call we can ackermannize -/
+  /- We need at least one argument for this to be a function call we can ackermannize. -/
   guard <| xs.size > 0
   let codTy ← BVTy.ofExpr? (← inferType f)
   let args ← xs.mapM Argument.ofExpr?
   let fn : Function := { f, codTy }
   return (fn, args)
 end Function
-
-
 /--
 NOTE: is it sensible to hash an array of arguments?
 We may want to use something like a trie to index these.
 Consider swiching to something like `Trie`.
 -/
 abbrev ArgumentList := Array Argument
-
-
 /--
 The data stored for an ackermannized call to allow us to build proofs.
 -/
 structure CallVal where
-  /-- the free variable `ack_fx₁...xₙ := (f x₁ x₂ ... xₙ)`. -/
+  /-- The free variable `ack_fx₁...xₙ := (f x₁ x₂ ... xₙ)`. -/
   fvar : FVarId 
   /-- heqProof : The proof that `ack_fx₁...fxₙ = f x₁ x₂ ... xₙ` (name/fvar = value/expr). -/
   heqProof : Expr  
@@ -140,18 +132,14 @@ instance : ToMessageData CallVal where
   toMessageData cv := m!"{Expr.fvar cv.fvar} ({cv.heqProof})"
 
 end CallVal
-
-
 structure State where
   /--
-  A maping from a function `f`, to a map of arguments `x₁ ... xₙ`, to the information stored about the call.
+  A mapping from a function `f`, to a map of arguments `x₁ ... xₙ`, to the information stored about the call.
   This is used to generate equations of the form `x₁ = y₁ → x₂ = y₂ → ... → xₙ = yₙ → ack_fx₁...xₙ = ack_fy₁...yₙ on-demand.
   -/
   fn2args2call : Std.HashMap Function (Std.HashMap ArgumentList CallVal) := {}
   /-- A counter for generating fresh names. -/
   gensymCounter : Nat := 0
-
-
 def State.init (_cfg : Config) : State where
 
 abbrev AckM := StateRefT State (ReaderT Context MetaM)
@@ -206,9 +194,8 @@ private def _insertCallVal (fn : Function) (args : ArgumentList) (cv : CallVal)
 Replace a call to the function `f` with an `fvar`. Since the `fvar` is defeq to the call,
 we can just replace occurrences of the call with the fvar `f`.
 
-We will later need to add another hypothesis with the equality that the `fvar = f x₁ ... xₙ`
+We will later need to add another hypothesis with the equality that `fvar = f x₁ ... xₙ`
 -/
-
 def replaceCallWithFVar (g : MVarId) (fn : Function) (args : ArgumentList) : AckM (CallVal × MVarId) := do
   g.withContext do
     if let some val ← getCallVal? fn args then 
@@ -221,7 +208,7 @@ def replaceCallWithFVar (g : MVarId) (fn : Function) (args : ArgumentList) : Ack
     _insertCallVal fn args cv
     return (cv, g)
 
-/-- create a trace node in trace class (i.e. `set_option traceClass true`),
+/-- Create a trace node in trace class (i.e. `set_option traceClass true`),
 with header `header`, whose default collapsed state is `collapsed`. -/
 def withTraceNode (header : MessageData) (k : AckM α)
     (collapsed : Bool := true)
@@ -303,8 +290,6 @@ partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := d
         trace[bv_ack] "{checkEmoji} {e} → {call}."
         return (Expr.fvar call.fvar, g)
 end
-
-
 /--
 Return true if the argument lists are trivially different.
 This is an optimization that we do not yet implement.
@@ -334,8 +319,6 @@ private example (f : X1 → X2 → O)
   fx1x2 = fx1'x2' := 
   let appEqApp : f x1 x2 = f x1' x2' := congr (congr (Eq.refl f) h1) h2
   Eq.trans (Eq.trans ackEqApp appEqApp) (Eq.symm ackEqApp')
-
-
 /-
 Make the ackermannization theorem, which states that: `(∀ i, arg₁[i] = arg₂[i]) -> call₁ = call₂`.
 Formally, we build an expr such as `arg₁ = arg'₁ -> arg₂ = arg'₂ -> ... argₙ = arg'ₙ -> call₁ = call₂`,

From a2ba6ca9a2037da600d36a882fbe158247721aee Mon Sep 17 00:00:00 2001
From: Siddharth <siddu.druid@gmail.com>
Date: Wed, 9 Oct 2024 00:29:16 -0500
Subject: [PATCH 06/10] Apply suggestions from code review

Co-authored-by: Tobias Grosser <github@grosser.es>
---
 src/Lean/Elab/Tactic/BVAckermannize.lean | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
index c5ea4b56c2c8..a87252401297 100644
--- a/src/Lean/Elab/Tactic/BVAckermannize.lean
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -188,7 +188,6 @@ def introDefExt (g : MVarId) (name : Name) (hdefTy : Expr) (hdefVal : Expr) : Ac
 private def _insertCallVal (fn : Function) (args : ArgumentList) (cv : CallVal) : AckM Unit := do
   let calls ← getCallMap fn
   modify fun s => { s with fn2args2call := s.fn2args2call.insert fn (calls.insert args cv) }
-      
 
 /--
 Replace a call to the function `f` with an `fvar`. Since the `fvar` is defeq to the call,

From cbab54334a04b1dfc62c99eda3d3eee2e18c6c84 Mon Sep 17 00:00:00 2001
From: Siddharth <siddu.druid@gmail.com>
Date: Thu, 17 Oct 2024 02:07:53 -0500
Subject: [PATCH 07/10] Apply suggestions from code review

Co-authored-by: Alex Keizer <alex@keizer.dev>
---
 src/Lean/Elab/Tactic/BVAckermannize.lean | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
index a87252401297..2ffd8f0245fd 100644
--- a/src/Lean/Elab/Tactic/BVAckermannize.lean
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -54,7 +54,9 @@ namespace BVTy
  /-- info: _root_.BitVec (w : Nat) : Type -/
  #guard_msgs in #check _root_.BitVec
 
-/-- Reify a raw expression into the types of bitvectors we can bitblast -/
+/-- Reify a raw expression of type `Type` into the types of bitvectors we can bitblast,
+returning `none` if `e` was not recognized as either `Bool` or `BitVec ?w`, 
+with `?w` a literal `Nat`  -/
 def ofExpr? (e : Expr) : OptionT MetaM BVTy :=
   match_expr e.consumeMData with
   | _root_.Bool => return Bool
@@ -81,7 +83,7 @@ namespace Argument
 instance : ToMessageData Argument where
   toMessageData arg := m!"{arg.x} : {arg.xTy}"
 
-/-- Build an `Argument` from a raw expression. -/
+/-- Build an `Argument` from a raw expression of type `Bool` or `BitVec _`. -/
 def ofExpr? (e : Expr) : OptionT MetaM Argument := do
   let t ← BVTy.ofExpr? (← inferType e)
   return { x := e, xTy := t}

From e50a6f8eb684da8024fe77187ad7508008aff19f Mon Sep 17 00:00:00 2001
From: Siddharth <siddu.druid@gmail.com>
Date: Thu, 17 Oct 2024 02:18:40 -0500
Subject: [PATCH 08/10] Apply suggestions from code review

Co-authored-by: Alex Keizer <alex@keizer.dev>
---
 src/Std/Tactic/BVAckermannize.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Std/Tactic/BVAckermannize.lean b/src/Std/Tactic/BVAckermannize.lean
index 5e555a161137..95c7a576e070 100644
--- a/src/Std/Tactic/BVAckermannize.lean
+++ b/src/Std/Tactic/BVAckermannize.lean
@@ -8,5 +8,5 @@ import Std.Tactic.BVAckermannize.Syntax
 
 /-!
 This directory contains the lazy ackermannization tactic.
-This uses lean's inbuilt bitblaster to uninterpreted functions + bitvectors to SAT.
+This uses lean's builtin bitblaster to reduce uninterpreted functions + bitvectors to SAT.
 -/

From 7c1ee08a770552fc07af063499c3d9911d040574 Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Thu, 17 Oct 2024 02:18:54 -0500
Subject: [PATCH 09/10] chore: address alex's comments

---
 src/Lean/Elab/Tactic/BVAckermannize.lean | 49 +++++++++++-------------
 1 file changed, 23 insertions(+), 26 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
index 2ffd8f0245fd..de5541d4243d 100644
--- a/src/Lean/Elab/Tactic/BVAckermannize.lean
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -3,7 +3,7 @@ Copyright (c) 2024 Siddharth Bhat. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Siddharth Bhat
 
-This file implements lazy ackermannization [1, 2]
+This file implements strict ackermannization [1, 2]
 
 [1] https://lara.epfl.ch/w/_media/model-based.pdf
 [2] https://leodemoura.github.io/files/oregon08.pdf
@@ -31,11 +31,6 @@ initialize Lean.registerTraceClass `bv_ack
 
 namespace Ack
 
-structure Config where
-
-structure Context extends Config where
- config : Config
-
  /-- Types that can be bitblasted by bv_decide -/
  inductive BVTy
  /-- Booleans -/
@@ -51,9 +46,6 @@ instance : ToMessageData BVTy where
 
 namespace BVTy
 
- /-- info: _root_.BitVec (w : Nat) : Type -/
- #guard_msgs in #check _root_.BitVec
-
 /-- Reify a raw expression of type `Type` into the types of bitvectors we can bitblast,
 returning `none` if `e` was not recognized as either `Bool` or `BitVec ?w`, 
 with `?w` a literal `Nat`  -/
@@ -65,12 +57,14 @@ def ofExpr? (e : Expr) : OptionT MetaM BVTy :=
      return .BitVec w
   | _ => OptionT.fail
 
+/-- Convert a `BVTy` back into an `Expr` -/
 def toExpr : BVTy → Expr
 | .Bool => mkConst ``_root_.Bool
 | .BitVec w => mkApp (mkConst ``_root_.BitVec) (mkNatLit w)
 
 end BVTy
 
+/-- An argument to an uninterpreted function, which we track for ackermannization. -/
 structure Argument where
   /-- The expression corresponding to the argument -/
   x : Expr
@@ -89,15 +83,21 @@ def ofExpr? (e : Expr) : OptionT MetaM Argument := do
   return { x := e, xTy := t}
 
 end Argument
+
+/--
+A function, which packs the expression and the type of the codomain of the function.
+We use the type of the codomain to build an abstracted value for every call of this function.
+-/
 structure Function where
   /-- The function -/
   f : Expr
-  codTy : BVTy
+  /-- The type of the function's codomain -/
+  codomain : BVTy
  deriving Hashable, BEq, Inhabited
 namespace Function
 
 instance : ToMessageData Function where
-  toMessageData fn := m!"{fn.f} (cod: {fn.codTy})"
+  toMessageData fn := m!"{fn.f} (cod: {fn.codomain})"
 
 /--
 Reify an expression `e` of the kind `f x₁ ... xₙ`, where all the arguments and the return type are
@@ -107,9 +107,9 @@ def reifyAp (f : Expr) : OptionT MetaM (Function × Array Argument) := do
   let xs := f.getAppArgs
   /- We need at least one argument for this to be a function call we can ackermannize. -/
   guard <| xs.size > 0
-  let codTy ← BVTy.ofExpr? (← inferType f)
+  let codomain ← BVTy.ofExpr? (← inferType f)
   let args ← xs.mapM Argument.ofExpr?
-  let fn : Function := { f, codTy }
+  let fn : Function := { f, codomain }
   return (fn, args)
 end Function
 /--
@@ -142,14 +142,14 @@ structure State where
   fn2args2call : Std.HashMap Function (Std.HashMap ArgumentList CallVal) := {}
   /-- A counter for generating fresh names. -/
   gensymCounter : Nat := 0
-def State.init (_cfg : Config) : State where
 
-abbrev AckM := StateRefT State (ReaderT Context MetaM)
+def State.init : State where
+
+abbrev AckM := StateRefT State MetaM
 
 namespace AckM
 
-def run (m : AckM α) (ctx : Context) : MetaM α :=
-  m.run' (State.init ctx.config) |>.run ctx
+def run (m : AckM α) : MetaM α := m.run' State.init 
 
 /-- Generate a fresh name. -/
 def gensym : AckM Name := do
@@ -204,7 +204,7 @@ def replaceCallWithFVar (g : MVarId) (fn : Function) (args : ArgumentList) : Ack
       return (val, g)
     let e := mkAppN fn.f (args.map Argument.x)
     let name ← gensym
-    let (introDef, g) ← introDefExt g name fn.codTy.toExpr e
+    let (introDef, g) ← introDefExt g name fn.codomain.toExpr e
     let cv := { fvar := introDef.defn, heqProof := Expr.fvar introDef.eqProof : CallVal }
     _insertCallVal fn args cv
     return (cv, g)
@@ -261,13 +261,13 @@ partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := d
       withTraceNode m!"🎯 Expr.app '{e}'" (collapsed := false) do
         let f := e.getAppFn
         let te ← inferType e
-        let .some codTy ← BVTy.ofExpr? te |>.run
+        let .some codomain ← BVTy.ofExpr? te |>.run
           | do
             trace[bv_ack] "{crossEmoji} '{te}' not BitVec/Bool."
             return (← ackAppChildren g e)
         trace[bv_ack] "{checkEmoji} {e}'s codomain '{te}'"
 
-        let fn := { f, codTy : Function }
+        let fn := { f, codomain : Function }
 
         let args := e.getAppArgs
         assert! args.size > 0 -- since we are an application, we must have at least one argument.
@@ -278,7 +278,6 @@ partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := d
         for arg in args do
           trace[bv_ack] "🎯 arg {arg}"
           let (arg, g) ← introAckForExpr g arg
-          -- do I need a `withContext` here? :(
           if let .some ackArg ← Argument.ofExpr? arg |>.run then
             trace[bv_ack] "{checkEmoji} arg {arg}"
             ackArgs := ackArgs.push ackArg
@@ -393,10 +392,10 @@ def ack (g : MVarId) : AckM MVarId := do
 end AckM
 
 /-- Entry point for programmatic usage of `bv_ackermannize` -/
-def ackTac (ctx : Context) : TacticM Unit := do
+def ackTac : TacticM Unit := do
   withoutRecover do
     liftMetaTactic fun g => do
-      let g ← (AckM.ack g).run ctx
+      let g ← (AckM.ack g).run
       return [g]
 
 end Ack
@@ -404,7 +403,5 @@ end Ack
 @[builtin_tactic Lean.Parser.Tactic.bvAckEager]
 def evalBvAckEager : Tactic := fun
   | `(tactic| bv_ack_eager) => 
-    let config : Ack.Config := {}
-    let ctx : Ack.Context := { config := config }
-    Ack.ackTac ctx
+    Ack.ackTac
   | _ => throwUnsupportedSyntax

From b386209d1a0fa511fb7138671e660971df0d4d6d Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Sat, 19 Oct 2024 09:07:02 -0500
Subject: [PATCH 10/10] chore: more lazy ack

---
 src/Lean/Elab/Tactic/BVAckermannize.lean     |   5 +-
 src/Lean/Elab/Tactic/BVAckermannizeLazy.lean | 110 +++++++++++++++++++
 src/Std/Tactic/BVAckermannize/Syntax.lean    |  11 ++
 3 files changed, 125 insertions(+), 1 deletion(-)
 create mode 100644 src/Lean/Elab/Tactic/BVAckermannizeLazy.lean

diff --git a/src/Lean/Elab/Tactic/BVAckermannize.lean b/src/Lean/Elab/Tactic/BVAckermannize.lean
index de5541d4243d..58f19c64e9bb 100644
--- a/src/Lean/Elab/Tactic/BVAckermannize.lean
+++ b/src/Lean/Elab/Tactic/BVAckermannize.lean
@@ -145,6 +145,7 @@ structure State where
 
 def State.init : State where
 
+
 abbrev AckM := StateRefT State MetaM
 
 namespace AckM
@@ -211,7 +212,7 @@ def replaceCallWithFVar (g : MVarId) (fn : Function) (args : ArgumentList) : Ack
 
 /-- Create a trace node in trace class (i.e. `set_option traceClass true`),
 with header `header`, whose default collapsed state is `collapsed`. -/
-def withTraceNode (header : MessageData) (k : AckM α)
+private def withTraceNode (header : MessageData) (k : AckM α)
     (collapsed : Bool := true)
     (traceClass : Name := `bv_ack) : AckM α :=
   Lean.withTraceNode traceClass (fun _ => return header) k (collapsed := collapsed)
@@ -290,6 +291,8 @@ partial def introAckForExpr (g : MVarId) (e : Expr) : AckM (Expr × MVarId) := d
         trace[bv_ack] "{checkEmoji} {e} → {call}."
         return (Expr.fvar call.fvar, g)
 end
+
+#check ackAppChildren
 /--
 Return true if the argument lists are trivially different.
 This is an optimization that we do not yet implement.
diff --git a/src/Lean/Elab/Tactic/BVAckermannizeLazy.lean b/src/Lean/Elab/Tactic/BVAckermannizeLazy.lean
new file mode 100644
index 000000000000..c4d13da9d30e
--- /dev/null
+++ b/src/Lean/Elab/Tactic/BVAckermannizeLazy.lean
@@ -0,0 +1,110 @@
+/-
+Copyright (c) 2024 Siddharth Bhat. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Siddharth Bhat
+
+This file implements lazy ackermannization [1, 2]
+
+[1] https://lara.epfl.ch/w/_media/model-based.pdf
+[2] https://leodemoura.github.io/files/oregon08.pdf
+[3] https://github.com/Z3Prover/z3/blob/d047b86439ec209446d211f0f6b251ebfba070d8/src/ackermannization/lackr.cpp#L206
+[4] https://github.com/Z3Prover/z3/blob/d047b86439ec209446d211f0f6b251ebfba070d8/src/ackermannization/lackr_model_constructor.cpp#L344
+-/
+prelude
+import Lean.Expr
+import Lean.Message
+import Std.Tactic.BVDecide.Bitblast
+import Std.Tactic.BVAckermannize.Syntax
+import Lean.Meta.Basic
+import Lean.Elab.Tactic.Basic
+import Lean.Meta.LitValues
+import Lean.Meta.InferType
+import Lean.Elab.Tactic.FalseOrByContra
+import Lean.Meta.Tactic.Assert
+import Lean.Meta.Tactic.Util
+import Lean.Meta.Tactic.Congr
+import Lean.Meta.Tactic.Replace
+import Lean.Elab.Tactic.BVAckermannize
+
+open Lean Elab Meta Tactic
+
+
+namespace Ack
+
+abbrev Fn2Args2Call := Std.HashMap Function (Std.HashMap ArgumentList CallVal) 
+/--
+A lazily built model,
+which knows whch assignments need the funext axiom to be added.
+-/
+structure LazyState where
+  /-- The set of calls for which we should add the extensionality axiom -/
+  fn2args2callUnfolded : Fn2Args2Call
+
+/-- Uninterpreted functions that created violations. -/
+def findViolations : Fn2Args2Call := sorry
+
+/--
+Lazy ackermannization maintains the ackermannization state,
+and enriches it with information necessary for lazy unfolding.
+-/
+abbrev AckLazyM := StateRefT LazyState AckM
+namespace AckLazyM
+
+/-- Create a trace node in trace class (i.e. `set_option traceClass true`),
+with header `header`, whose default collapsed state is `collapsed`. -/
+private def withTraceNode (header : MessageData) (k : AckM α)
+    (collapsed : Bool := true)
+    (traceClass : Name := `bv_ack) : AckM α :=
+  Lean.withTraceNode traceClass (fun _ => return header) k (collapsed := collapsed)
+
+def ackLazyOnce (g : MVarId) : AckLazyM MVarId := do
+  g.withContext do
+    let (target', g) ← Ack.AckM.introAckForExpr g (← inferType (Expr.mvar g))
+    let g ← g.replaceTargetDefEq target'
+
+    let hyps ← g.getNondepPropHyps
+    let mut g := g
+    for hyp in hyps do
+      g ← g.withContext do
+        withTraceNode m!"🎯 hyp '{← hyp.getType}'" (collapsed := false) do
+          let hypG ← Ack.AckM.introAckForExpr g (← hyp.getType)
+          pure hypG.2
+
+    for (fn, arg2call) in (← getThe Ack.State).fn2args2call do
+      let argCallsArr := arg2call.toArray
+      for i in [0:argCallsArr.size] do
+        let (arg₁, call₁) := argCallsArr[i]!
+        for j in [i+1:argCallsArr.size] do
+          let (arg₂, call₂) := argCallsArr[j]!
+          if ← Ack.AckM.areArgListsTriviallyDifferent arg₁ arg₂ then
+            continue
+          g ← Ack.AckM.mkAckThm g fn arg₁ arg₂ call₁ call₂
+    trace[bv_ack] "{checkEmoji} ack.{indentD g}"
+    return g
+
+
+def ack (g : MVarId) : AckM MVarId := do
+  let mut madeProgress := true
+  let mut g := g
+  while true do
+    if !madeProgress then
+      throwError "{crossEmoji} unable to prove goal with lazy ackermannization"
+  return g
+
+end AckLazyM
+
+/-- Entry point for programmatic usage of `bv_ackermannize` -/
+def ackLazyTac : TacticM Unit := do
+  withoutRecover do
+    liftMetaTactic fun g => do
+      let g ← (AckLazyM.ack g).run
+      return [g]
+
+end Ack
+
+@[builtin_tactic Lean.Parser.Tactic.bvAckLazy]
+def evalBvLazy : Tactic := fun
+  | `(tactic| bv_ack_eager) => 
+    Ack.ackLazyTac
+  | _ => throwUnsupportedSyntax
+
diff --git a/src/Std/Tactic/BVAckermannize/Syntax.lean b/src/Std/Tactic/BVAckermannize/Syntax.lean
index f91e0a794b27..8fc843b544bd 100644
--- a/src/Std/Tactic/BVAckermannize/Syntax.lean
+++ b/src/Std/Tactic/BVAckermannize/Syntax.lean
@@ -26,6 +26,17 @@ example : ∀ (a b : BitVec 64), (a &&& b) + (a ^^^ b) = a ||| b := by
 -/
 syntax (name := bvAckEager) "bv_ack_eager" : tactic
 
+
+/--
+Close fixed-width `BitVec`, `Bool`, and uninterpreted function goals by obtaining a proof from an external SAT solver and
+verifying it inside Lean. The solvable goals are currently limited to the Lean equivalent of
+[`QF_BV`](https://smt-lib.org/logics-all.shtml#QF_BV) + [`QF_UF`](https://smt-lib.org/logics-all.shtml#QF_UF)
+
+Note that this uses *lazy* ackermannization,
+and will therefore lazily add hypotheses of functional extensionality based on the SAT solver's counterexample model.
+-/
+syntax (name := bvAckLazy) "bv_ack_lazy" : tactic
+
 end Tactic
 
 end Lean.Parser