Improve the performance of condSimpTac

Makefile

@@ -178,7 +178,7 @@ test-%: build-dev
 # Replay the Lean tests and time them
 .PHONY: timed-lean
 timed-lean:
-	cd tests/lean && find . -type f -iname "*.lean" -not -path "./.lake/*" -exec printf "\n{}\n" \; -exec lake env time lean {} \; >& timing.out
+	cd tests/lean && find . -type f -iname "*.lean" -not -path "./.lake/*" -exec printf "\n{}\n" \; -exec lake env time lean {} \;
 
 # =============================================================================
 # Nix

backends/lean/Aeneas.lean

@@ -19,6 +19,7 @@ import Aeneas.Saturate
 import Aeneas.ScalarDecrTac
 import Aeneas.ScalarNF
 import Aeneas.ScalarTac
+import Aeneas.Simp
 import Aeneas.SimpIfs
 import Aeneas.SimpLemmas
 import Aeneas.SimpLists

backends/lean/Aeneas/Array.lean

@@ -5,6 +5,8 @@ namespace Array
 
 attribute [-simp] List.getElem!_eq_getElem?_getD
 
+attribute [scalar_tac_simps, simp_lists_simps] Array.size
+
 def setSlice! {α} (a : Array α) (i : ℕ) (s : List α) : Array α :=
   ⟨ a.toList.setSlice! i s⟩
 

backends/lean/Aeneas/BitVec.lean

@@ -17,6 +17,9 @@ open Lean
 
 attribute [-simp] List.getElem!_eq_getElem?_getD
 
+attribute [bvify_simps, simp_scalar_simps] BitVec.zero_eq
+attribute [bvify_simps, simp_scalar_simps] BitVec.instInhabited
+
 def BitVec.toArray {n} (bv: BitVec n) : Array Bool := Array.finRange n |>.map (bv[·])
 def BitVec.ofFn {n} (f: Fin n → Bool) : BitVec n := (BitVec.ofBoolListLE <| List.ofFn f).cast (by simp)
 def BitVec.set {n} (i: Fin n) (b: Bool) (bv: BitVec n) : BitVec n := bv ^^^ (((bv[i] ^^ b).toNat : BitVec n) <<< i.val)

backends/lean/Aeneas/BvTac/BvTac.lean

@@ -55,26 +55,26 @@ partial def bvTacPreprocess (config : Config) (n : Option Expr): TacticM Unit :=
   /- First try simplifying the goal - if it is an (in-)equality between scalars, it may get
      the bitwidth to use for the bit-vectors might be obvious from the goal: we marked some
      theorems wiht `bvify_simps` for this reason. -/
-  Bvify.bvifyTacSimp (Utils.Location.targets #[] true)
+  let r ← Bvify.bvifyTacSimp (Utils.Location.targets #[] true)
+  if r.isNone then return
   /- The simp call above may have proven the goal (unlikely, but we have to take this
      into account) -/
-  allGoals do
-    trace[BvTac] "Goal after `bvifyTacSimp`: {← getMainGoal}"
-    /- Figure out the bitwidth if the user didn't provide it -/
-    let n ← do
-      match n with
-      | some n => pure n
-      | none => getn
-    /- Then apply bvify -/
-    bvifyTac config.toConfig n Utils.Location.wildcard
-    trace[BvTac] "Goal after `bvifyTac`: {← getMainGoal}"
-    /- Call `simp_all ` to normalize the goal a bit -/
-    let simpLemmas ← bvifySimpExt.getTheorems
-    let simprocs ← bvifySimprocExt.getSimprocs
-    Utils.simpAll {dsimp := false, failIfUnchanged := false, maxDischargeDepth := 0} true
-                  {simprocs := #[simprocs], simpThms := #[simpLemmas]}
-    allGoals do
-    trace[BvTac] "Goal after `simp_all`: {← getMainGoal}"
+  trace[BvTac] "Goal after `bvifyTacSimp`: {← getMainGoal}"
+  /- Figure out the bitwidth if the user didn't provide it -/
+  let n ← do
+    match n with
+    | some n => pure n
+    | none => getn
+  /- Then apply bvify -/
+  bvifyTac config.toConfig n Utils.Location.wildcard
+  trace[BvTac] "Goal after `bvifyTac`: {← getMainGoal}"
+  /- Call `simp_all ` to normalize the goal a bit -/
+  let simpLemmas ← bvifySimpExt.getTheorems
+  let simprocs ← bvifySimprocExt.getSimprocs
+  Simp.simpAll {dsimp := false, failIfUnchanged := false, maxDischargeDepth := 0} true
+                {simprocs := #[simprocs], simpThms := #[simpLemmas]}
+  -- The simpAll may have solved the goal, so we need to be careful
+  allGoals do trace[BvTac] "Goal after `simp_all`: {← getMainGoal}"
 
 elab "bv_tac_preprocess" config:Parser.Tactic.optConfig n:(colGt term)? : tactic => do
   bvTacPreprocess (← elabConfig config) (← optElabTerm n)

backends/lean/Aeneas/Bvify/Bvify.lean

@@ -280,20 +280,24 @@ def bvifyAddSimpThms (n : Expr) : TacticM (Array FVarId) := do
 
 def bvifySimpConfig : Simp.Config := {maxDischargeDepth := 2, failIfUnchanged := false}
 
-def bvifyTacSimp (loc : Utils.Location) : TacticM Unit := do
+def bvifyTacSimp (loc : Utils.Location) : TacticM (Option (Array FVarId)) := do
   let args : ScalarTac.CondSimpArgs := {
       simpThms := #[← bvifySimpExt.getTheorems, ← SimpBoolProp.simpBoolPropSimpExt.getTheorems]
       simprocs := #[← bvifySimprocExt.getSimprocs, ← SimpBoolProp.simpBoolPropSimprocExt.getSimprocs]
     }
-  ScalarTac.condSimpTacSimp bvifySimpConfig args loc #[] false
+  ScalarTac.condSimpTacSimp bvifySimpConfig args loc #[] #[] none
 
 def bvifyTac (config : Config) (n : Expr) (loc : Utils.Location) : TacticM Unit := do
+  let hypsArgs : ScalarTac.CondSimpArgs := {
+      simpThms := #[← bvifyHypsSimpExt.getTheorems]
+      simprocs := #[← bvifyHypsSimprocExt.getSimprocs]
+    }
   let args : ScalarTac.CondSimpArgs := {
       simpThms := #[← bvifySimpExt.getTheorems, ← SimpBoolProp.simpBoolPropSimpExt.getTheorems]
       simprocs := #[← bvifySimprocExt.getSimprocs, ← SimpBoolProp.simpBoolPropSimprocExt.getSimprocs]
     }
   let config := { nonLin := config.nonLin, saturationPasses := config.saturationPasses }
-  ScalarTac.condSimpTac "bvify" config bvifySimpConfig args (bvifyAddSimpThms n) true loc
+  ScalarTac.condSimpTac "bvify" config bvifySimpConfig hypsArgs args (bvifyAddSimpThms n) true loc
 
 syntax (name := bvify) "bvify " colGt Parser.Tactic.optConfig term (location)? : tactic
 

backends/lean/Aeneas/Bvify/Init.lean

@@ -29,4 +29,18 @@ initialize bvifySimprocExt : Simp.SimprocExtension ←
     The `bvify_simps_proc` attribute registers simp procedures to be used by `bvify`
     during its preprocessing phase." none --(some bvifySimprocsRef)
 
+/-- The `bvify_hyps_simps` simp attribute. -/
+initialize bvifyHypsSimpExt : SimpExtension ←
+  registerSimpAttr `bvify_hyps_simps "\
+    The `bvify_hyps_simps` attribute registers simp lemmas to be used by `bvify`."
+
+-- TODO: initialization fails with this, while the same works for `scalar_tac`??
+--initialize bvifySimprocsRef : IO.Ref Simprocs ← IO.mkRef {}
+
+/-- The `bvify_hyps_simps_proc` simp attribute for the simp rocs. -/
+initialize bvifyHypsSimprocExt : Simp.SimprocExtension ←
+  Simp.registerSimprocAttr `bvify_hyps_simps_proc "\
+    The `bvify_hyps_simps_proc` attribute registers simp procedures to be used by `bvify`
+    during its preprocessing phase." none --(some bvifySimprocsRef)
+
 end Aeneas.Bvify

backends/lean/Aeneas/Let/Let.lean

@@ -15,7 +15,7 @@ def opaque_refold (h x : Name) (e : Expr) : TacticM Unit :=
   withMainContext do
   /- Retrieve the list of propositions in the context -/
   let ctx ← getLCtx
-  let props ← (← ctx.getDecls).filterM fun x => do pure (← inferType x.type).isProp
+  let props ← (← ctx.getDecls).filterM fun x => isProp x.type
   /- Generalize -/
   let goal ← getMainGoal
   let (_, _, ngoal) ← goal.generalizeHyp #[{expr := e, xName? := x, hName? := h}] (props.map LocalDecl.fvarId).toArray
@@ -66,7 +66,7 @@ def transparent_refold (x : Name) (e : Expr) : TacticM Unit :=
   withMainContext do
   /- Retrieve the list of propositions in the context -/
   let ctx ← getLCtx
-  let props ← (← ctx.getDecls).filterM fun x => do pure (← inferType x.type).isProp
+  let props ← (← ctx.getDecls).filterM fun x => isProp x.type
   /- List the assumptions which contain the declaration that we want to refold -/
   let mut toRevert := #[]
   for decl in props.reverse do

backends/lean/Aeneas/List/List.lean

@@ -588,7 +588,7 @@ theorem setSlice!_drop {α} (l : List α) (i : ℕ) :
   by_cases h1: j < l.length <;> simp_lists
   simp_scalar
 
-@[simp_lists_simps]
+@[simp_lists_hyps_simps]
 def Inhabited_getElem_eq_getElem! {α} [Inhabited α] (l : List α) (i : ℕ) (hi : i < l.length) :
   l[i] = l[i]! := by
   simp only [List.getElem!_eq_getElem?_getD, List.getElem?_eq_getElem, Option.getD_some, hi]

backends/lean/Aeneas/Progress/Init.lean

@@ -65,33 +65,33 @@ structure ProgressSpecDesc where
   postcond? : Option Expr
 
 section Methods
-  variable [MonadLiftT MetaM m] [MonadControlT MetaM m] [Monad m] [MonadOptions m]
+  variable {m} [MonadLiftT MetaM m] [MonadControlT MetaM m] [Monad m] [MonadOptions m]
   variable [MonadTrace m] [MonadLiftT IO m] [MonadRef m] [AddMessageContext m]
   variable [MonadError m]
   variable {a : Type}
 
 
-/-- Given ty := ∀ xs.., ∃ zs.., program = res ∧ post?, destruct and run continuation -/
-def programTelescope[Inhabited (m α)] [Nonempty (m α)] (ty: Expr)
-  (k: (xs:Array (MVarId × BinderInfo)) → (zs:Array FVarId) → (program:Expr) → (res:Expr) → (post:Option Expr) → m α)
-: m α := do
-  let ty := ty.consumeMData
-  unless ←isProp ty do
-    throwError "Expected a proposition, got {←inferType ty}"
-  -- ty == ∀ xs, ty₂
-  let (xs, xs_bi, ty₂) ← forallMetaTelescope ty
-  trace[Progress] "Universally quantified arguments and assumptions: {xs}"
-  -- ty₂ == ∃ zs, ty₃ ≃ Exists {α} (fun zs => ty₃)
-  existsTelescope ty₂.consumeMData fun zs ty₃ => do
-    trace[Progress] "Existentials: {zs}"
-    trace[Progress] "Proposition after stripping the quantifiers: {ty₃}"
-    -- ty₃ == ty₄ ∧ post?
-    let (ty₄, post?) ← Utils.optSplitConj ty₃.consumeMData
-    trace[Progress] "After splitting the conjunction:\n- eq: {ty₄}\n- post: {post?}"
-    -- ty₄ == (program = res)
-    let (program, res) ← Utils.destEq ty₄.consumeMData
-    trace[Progress] "After splitting the equality:\n- lhs: {program}\n- rhs: {res}"
-    k (xs.map (·.mvarId!) |>.zip xs_bi) (zs.map (·.fvarId!)) program res post?
+  /-- Given ty := ∀ xs.., ∃ zs.., program = res ∧ post?, destruct and run continuation -/
+  def monadTelescope {α} [Inhabited (m α)] [Nonempty (m α)] (ty: Expr)
+    (k: (xs:Array (MVarId × BinderInfo)) → (zs:Array FVarId) → (program:Expr) → (res:Expr) → (post:Option Expr) → m α)
+  : m α := do
+    let ty := ty.consumeMData
+    unless ←isProp ty do
+      throwError "Expected a proposition, got {←inferType ty}"
+    -- ty == ∀ xs, ty₂
+    let (xs, xs_bi, ty₂) ← forallMetaTelescope ty
+    trace[Progress] "Universally quantified arguments and assumptions: {xs}"
+    -- ty₂ == ∃ zs, ty₃ ≃ Exists {α} (fun zs => ty₃)
+    existsTelescope ty₂.consumeMData fun zs ty₃ => do
+      trace[Progress] "Existentials: {zs}"
+      trace[Progress] "Proposition after stripping the quantifiers: {ty₃}"
+      -- ty₃ == ty₄ ∧ post?
+      let (ty₄, post?) ← Utils.optSplitConj ty₃.consumeMData
+      trace[Progress] "After splitting the conjunction:\n- eq: {ty₄}\n- post: {post?}"
+      -- ty₄ == (program = res)
+      let (program, res) ← Utils.destEq ty₄.consumeMData
+      trace[Progress] "After splitting the equality:\n- lhs: {program}\n- rhs: {res}"
+      k (xs.map (·.mvarId!) |>.zip xs_bi) (zs.map (·.fvarId!)) program res post?
 
   /- Analyze a goal or a progress theorem to decompose its arguments.
 
@@ -113,7 +113,7 @@ def programTelescope[Inhabited (m α)] [Nonempty (m α)] (ty: Expr)
   def withProgressSpec [Inhabited (m a)] [Nonempty (m a)]
     (isGoal : Bool) (th : Expr) (k : ProgressSpecDesc → m a) :
     m a := do
-    programTelescope th fun xs evars mExpr ret post => do
+    monadTelescope th fun xs evars mExpr ret post => do
     -- Recursively destruct the monadic application to dive into the binds,
     -- if necessary (this is for when we use `withProgressSpec` inside of the `progress` tactic),
     -- and destruct the application to get the function name