Add a tactic to genralize the widths of bitvectors

Author: ineol

Blase/Blase/MultiWidth/Tactic.lean

@@ -1093,6 +1093,123 @@ def evalBvMultiWidth : Tactic := fun
     solveEntrypoint g cfg
 | _ => throwUnsupportedSyntax
 
+/-
+A tactic to generalize the width of BitVectors
+-/
+
+structure State where
+  mapping : Std.HashMap Expr Expr
+  deriving Inhabited
+
+abbrev GenM := StateT State TermElabM 
+
+def State.get? (e : Expr) : GenM (Option Expr) := do
+  let s ← get
+  for (e', x) in s.mapping do
+    if ← isDefEq e e' then
+      return x
+  pure none
+
+def State.setMapping (e x : Expr) : GenM Unit := do
+  modify fun s =>
+    {s with mapping := s.mapping.insert e x}
+
+def State.add? (e : Expr) : GenM Expr := do
+  match ← get? e with
+  | some x => pure x
+  | none =>
+    if e.isFVar || e.isBVar then pure e else
+    let x ← mkFreshExprMVar (some (.const ``Nat [])) (userName := `w)
+    setMapping e x
+    pure x
+
+/--
+This table determines which arguments of important functions are bitwidths and
+should be generalized and which ones are normal parameters which should be
+recursively visited.
+-/
+def genTable : Std.HashMap Name (Array Bool) := Id.run do
+  let mut table := .emptyWithCapacity 16
+  table := table.insert ``BitVec #[true]
+  table := table.insert ``BitVec.zeroExtend #[true, true, false]
+  table := table.insert ``BitVec.signExtend #[true, true, false]
+  table := table.insert ``BitVec.instAdd #[true]
+  table := table.insert ``BitVec.instSub #[true]
+  table := table.insert ``BitVec.instMul #[true]
+  table := table.insert ``BitVec.instDiv #[true]
+  table
+
+partial def visit (t : Expr) : GenM Expr := do
+  let t ← instantiateMVars t
+  match t with
+  | .app _ _ =>
+    let f := t.getAppFn
+    let args := t.getAppArgs
+    let table := 
+      if let some (f, _) := f.const? then
+        genTable[f]?
+      else
+        none
+    let bv? (n : Nat) :=
+      match table with
+      | .some xs => xs.getD n false
+      | .none => false
+    /- let f ← visit f -/
+    args.zipIdx.foldlM (init := f) fun res (arg, i) => do
+      let arg ← if bv? i then State.add? arg else visit arg
+      pure <| .app res arg
+  | .forallE n e₁ e₂ info =>
+    pure <| .forallE n (← visit e₁) (← visit e₂) info
+  | e => 
+    pure e
+
+def doBvGeneralize (g : MVarId) : GenM (Expr × MVarId) := do
+  let lctx ← getLCtx
+  let mut allFVars := #[]
+  for h in lctx do
+    if not h.isImplementationDetail then
+      allFVars := allFVars.push h.fvarId
+  let (_, g) ← g.revert allFVars
+  let e ← visit (← g.getType)
+  let mut newVars := #[]
+  for (_, x) in (←get).mapping do
+    newVars := newVars.push x
+
+  let e ← mkForallFVars newVars e (binderInfoForMVars := .default)
+  let e ← instantiateMVars e
+  pure (e, g)
+
+/--
+This tactic tries to generalize the bitvector widths, and only the bitvector
+widths. See `genTable` if the tactic fails to generalize the right parameters
+of a function over bitvectors.
+-/
+syntax (name := bvGeneralize) "bv_generalize" Lean.Parser.Tactic.optConfig : tactic
+@[tactic bvGeneralize]
+def evalBvGeneralize : Tactic := fun
+| `(tactic| bv_generalize) => do
+  let g₀ ← getMainGoal
+  g₀.withContext do
+    let ((e, g), s) ← (doBvGeneralize g₀).run default
+    g.withContext do
+      let g' ← mkFreshExprMVar (some e)
+      -- TODO: instantiate the old goal
+      let mut newVals := #[]
+      for (e, x) in s.mapping do
+        newVals := newVals.push e
+      g.assign <| mkAppN g' newVals 
+      replaceMainGoal [g'.mvarId!]
+| _ => throwUnsupportedSyntax
+
+theorem test_bv_generalize_simple (x y : BitVec 32) (zs : List (BitVec 44)) : 
+    x = x := by
+  bv_generalize
+  bv_multi_width
+
+theorem test_bv_generalize (x y : BitVec 32) (zs : List (BitVec 44)) (z : BitVec 10) (h : 52 + 10 = 42) (heq : x = y) : 
+    x.zeroExtend 10 = y.zeroExtend 10 + 0 := by
+  bv_generalize
+  bv_multi_width
 
 end Tactic
 end MultiWidth