feat: proof reconstruction for rewrite rules

Blaster.lean

@@ -4,5 +4,6 @@
 import Blaster.Command
 import Blaster.Logging
 import Blaster.Optimize
+import Blaster.Reconstruct
 import Blaster.Smt
 import Blaster.StateMachine

Blaster/Command/Tactic.lean

@@ -36,12 +36,28 @@ def blasterTacticImp : Tactic := fun stx =>
    let sOpts ← parseSolveOptions opts default
    let (goal, nbQuantifiers) ← revertHypotheses (← getMainGoal)
    let env := {(default : TranslateEnv) with optEnv.options.solverOptions := sOpts}
-   let ((result, optExpr), _) ←
-     withTheReader Core.Context (fun ctx => { ctx with maxHeartbeats := 0 }) $ do
-       IO.setNumHeartbeats 0
-       Translate.main (← goal.getType >>= instantiateMVars') (logUndetermined := false) |>.run env
+   let ((result, (optExpr, proof)), _) ←
+     withTheReader Core.Context
+      (fun ctx => { ctx with maxHeartbeats := 0, maxRecDepth := max ctx.maxRecDepth 4096 }) $ do
+        IO.setNumHeartbeats 0
+        Translate.main (← goal.getType >>= instantiateMVars') (logUndetermined := false) |>.run env
    match result with
-   | .Valid => goal.admit -- TODO: replace with proof reconstruction
+   | .Valid =>
+        match proof with
+        | some p =>
+            -- verify certificate type matches goal before assigning
+            let goalType ← goal.getType
+            let pType ← inferType p
+            if (← isDefEq pType goalType) then
+              goal.assign p
+            else
+              logWarning "blaster: proof reconstruction failed, closing with admit"
+              goal.admit
+        | none =>
+            try goal.refl
+            catch _ =>
+              logWarning "blaster: proof reconstruction failed, closing with admit"
+              goal.admit
    | .Falsified cex => throwTacticEx `blaster goal "Goal was falsified (see counterexample above)"
    | .Undetermined =>
         -- Replace the goal with the optimized expression

Blaster/Optimize.lean

@@ -1,3 +1,3 @@
-
 import Blaster.Optimize.Basic
 import Blaster.Optimize.Lemmas
+import Blaster.Optimize.Types

Blaster/Optimize/Env.lean

@@ -1,9 +1,10 @@
 import Lean
+import Blaster.Command.Options
 import Blaster.Optimize.Expr
 import Blaster.Optimize.MatchInfo
 import Blaster.Optimize.Opaque
+import Blaster.Optimize.Types
 import Blaster.Smt.Term
-import Blaster.Command.Options
 
 open Lean Meta Blaster.Smt Blaster.Options
 
@@ -83,7 +84,7 @@ inductive MatchEntry where
  deriving Repr
 
 abbrev HypothesisMap := Std.HashMap Lean.Expr Lean.Expr
-abbrev RewriteCacheMap := Std.HashMap Lean.Expr Lean.Expr
+abbrev RewriteCacheMap := Std.HashMap Lean.Expr OptimizeResult
 abbrev MatchEntryMap := Std.HashMap Lean.Expr MatchEntry -- with key corresponding to a match pattern
 abbrev MatchContextMap := Std.HashMap Lean.Expr MatchEntryMap  -- with key corresponding to a match discriminator
 abbrev EqualityMap := Std.HashMap Lean.Expr Lean.Expr -- with key corresponding to expression to be replaced.
@@ -289,6 +290,11 @@ structure OptimizeEnv where
   -/
   restart : Bool
 
+  /-- Set of expressions whose proof certificates were stripped during caching
+      due to containing free variables in a global context.
+  -/
+  strippedProofExprs : Std.HashSet Lean.Expr
+
   /-- local declaration context -/
   ctx : LocalDeclContext
 
@@ -307,6 +313,7 @@ instance : Inhabited OptimizeEnv where
      memCache := default,
      options := default,
      restart := false,
+     strippedProofExprs := Std.HashSet.emptyWithCapacity,
      ctx := default
    }
 
@@ -612,20 +619,20 @@ def isInFunApp : TranslateEnvT Bool :=
   return (← get).optEnv.options.inFunApp
 
 @[always_inline, inline]
-def findGlobalCache (a : Expr) : TranslateEnvT (Option Expr) := do
+def findGlobalCache (a : Expr) : TranslateEnvT (Option OptimizeResult) := do
  return (← get).optEnv.globalRewriteCache.get? a
 
 @[always_inline, inline]
-def findLocalCache (a : Expr) : TranslateEnvT (Option Expr) := do
+def findLocalCache (a : Expr) : TranslateEnvT (Option OptimizeResult) := do
  return (← get).optEnv.localRewriteCache.get? a
 
 /-- Update global rewrite cache with `a := b`. -/
-def updateGlobalRewriteCache (a : Expr) (b : Expr) : TranslateEnvT Unit := do
-  modify (fun env => { env with optEnv.globalRewriteCache := env.optEnv.globalRewriteCache.insert a b })
+def updateGlobalRewriteCache (a : Expr) (r : OptimizeResult) : TranslateEnvT Unit := do
+  modify (fun env => { env with optEnv.globalRewriteCache := env.optEnv.globalRewriteCache.insert a r })
 
 /-- Update local rewrite cache with `a := b`. -/
-def updateLocalRewriteCache (a : Expr) (b : Expr) : TranslateEnvT Unit := do
-  modify (fun env => { env with optEnv.localRewriteCache := env.optEnv.localRewriteCache.insert a b })
+def updateLocalRewriteCache (a : Expr) (r : OptimizeResult) : TranslateEnvT Unit := do
+  modify (fun env => { env with optEnv.localRewriteCache := env.optEnv.localRewriteCache.insert a r })
 
 /-- Update synthesize decidable instance cache with `a := b`. -/
 @[always_inline, inline]
@@ -655,29 +662,30 @@ def withSynthInstanceCache (a : Expr) (f: Unit → TranslateEnvT (Option Expr))
 @[always_inline, inline]
 def mkExpr (a : Expr) (cacheResult := true) : TranslateEnvT Expr := do
    match (← findGlobalCache a) with
-   | some a' => return a'
+   | some r => return r.optExpr
    | none => do
-       if cacheResult then updateGlobalRewriteCache a a
+       if cacheResult then updateGlobalRewriteCache a ⟨a, none⟩
        return a
 
 /-- Return `true` only when both hypothesisMap and matchInContext are empty and isRefHyp flag is not set -/
 @[always_inline, inline]
 def isGlobalContext : TranslateEnvT Bool := do
-  let ⟨_, ⟨_, _, _, _, _, _, _, _, hypothesisContext, matchInContext, _, _, _, _⟩⟩ ← get
+  let ⟨_, ⟨_, _, _, _, _, _, _, _, hypothesisContext, matchInContext, _, _, _, _, _⟩⟩ ← get
   return hypothesisContext.hypothesisMap.size == 0 && matchInContext.size == 0
 
 /-- Perform the following:
       - When isGlobal
-         - Add entry `a := b` to `globalRewriteCache`
+         - Add entry `a := r` to `globalRewriteCache`
       - Otherwise
-         - Add entry `a := b` to `localRewriteCache`
+         - Add entry `a := r` to `localRewriteCache`
 -/
 @[always_inline, inline]
-def updateOptimizeEnvCache (a : Expr) (b : Expr) (isGlobal : Bool) : TranslateEnvT Unit := do
-  -- trace[Optimize.cacheExpr] "cacheExpr {← ppExpr a} ===> {← ppExpr b}"
+def updateOptimizeEnvCache (a : Expr) (r : OptimizeResult) (isGlobal : Bool) :
+    TranslateEnvT Unit := do
+  -- trace[Optimize.cacheExpr] "cacheExpr {← ppExpr a} ===> {← ppExpr r.optExpr}"
   if isGlobal
-  then updateGlobalRewriteCache a b
-  else updateLocalRewriteCache a b
+  then updateGlobalRewriteCache a r
+  else updateLocalRewriteCache a r
 
 /-- Perform the following:
       - When isGlobal
@@ -690,7 +698,7 @@ def updateOptimizeEnvCache (a : Expr) (b : Expr) (isGlobal : Bool) : TranslateEn
          - Otherwise `none`
 -/
 @[always_inline, inline]
-def isInOptimizeCache? (a : Expr) (isGlobal : Bool) : TranslateEnvT (Option Expr) := do
+def isInOptimizeCache? (a : Expr) (isGlobal : Bool) : TranslateEnvT (Option OptimizeResult) := do
  if isGlobal
  then findGlobalCache a
  else findLocalCache a
@@ -1712,7 +1720,7 @@ where
     An error is triggered if no corresponding entry can be found in `recFunMap`.
 -/
 def hasRecFunInst? (instApp : Expr) : TranslateEnvT (Option Expr) := do
-  let ⟨_, ⟨_, _, _, _, _,recFunInstCache,_,recFunMap, _, _, _, _, _, _⟩⟩ ← get
+  let ⟨_, ⟨_, _, _, _, _,recFunInstCache,_,recFunMap, _, _, _, _, _, _, _⟩⟩ ← get
   match recFunInstCache.get? instApp with
   | some fbody =>
      -- retrieve function application from recFunMap