Compute goal diffs for all rules

Aesop/BuiltinRules/ApplyHyps.lean

@@ -13,9 +13,14 @@ open Lean.Meta
 
 def applyHyp (hyp : FVarId) (goal : MVarId) (md : TransparencyMode) :
     MetaM RuleApplication := do
-  let (_, #[step]) ← applyS goal (.fvar hyp) none md |>.run
+  let (goals, #[step]) ← applyS goal (.fvar hyp) none md |>.run
     | throwError "aesop: internal error in applyHyps: multiple steps"
-  return .ofLazyScriptStep step none
+  return {
+    goals := goals.map ({ mvarId := ·, diff := ∅ })
+    postState := step.postState
+    scriptSteps? := #[step]
+    successProbability? := none
+  }
 
 @[aesop unsafe 75% tactic (rule_sets := [builtin])]
 def applyHyps : RuleTac := λ input =>

Aesop/BuiltinRules/Assumption.lean

@@ -44,7 +44,12 @@ def assumption : RuleTac := λ input => do
       let #[step] := steps
         | throwError "aesop: internal error in assumption: multiple steps"
       let proofHasMVar := (← fvarId.getType).hasMVar
-      let app := RuleApplication.ofLazyScriptStep step none
+      let app := {
+        goals := #[]
+        postState := step.postState
+        scriptSteps? := #[step]
+        successProbability? := none
+      }
       return some (app, proofHasMVar)
 
 end Aesop.BuiltinRules

Aesop/BuiltinRules/DestructProducts.lean

@@ -120,6 +120,9 @@ where
 partial def destructProducts : RuleTac := RuleTac.ofSingleRuleTac λ input => do
   let md := input.options.destructProductsTransparency
   let (goal, steps) ← destructProductsCore input.goal md
+  -- TODO we can construct a better diff here, and it would be important to do
+  -- so because `destructProducts` often renames hypotheses.
+  let goal := { mvarId := goal, diff := ← diffGoals input.goal goal ∅ }
   return (#[goal], steps, none)
 
 end Aesop.BuiltinRules

Aesop/BuiltinRules/Ext.lean

@@ -22,6 +22,7 @@ def extCore (goal : MVarId) : ScriptM (Option (Array MVarId)) :=
 def ext : RuleTac := RuleTac.ofSingleRuleTac λ input => do
   let (some goals, steps) ← extCore input.goal |>.run
     | throwError "found no applicable ext lemma"
+  let goals ← goals.mapM (mvarIdToSubgoal (parentMVarId := input.goal) · ∅)
   return (goals, steps, none)
 
 end Aesop.BuiltinRules

Aesop/BuiltinRules/Intros.lean

@@ -20,6 +20,9 @@ def intros : RuleTac := RuleTac.ofSingleRuleTac λ input => do
       | some md => introsUnfoldingS input.goal md |>.run
     if newFVarIds.size == 0 then
       throwError "nothing to introduce"
-    return (#[goal], steps, none)
+    let addedFVars := newFVarIds.foldl (init := ∅) λ set fvarId =>
+      set.insert fvarId
+    let diff := { addedFVars, removedFVars := ∅, fvarSubst := ∅ }
+    return (#[{ mvarId := goal, diff }], steps, none)
 
 end Aesop.BuiltinRules

Aesop/BuiltinRules/Split.lean

@@ -15,6 +15,7 @@ namespace Aesop.BuiltinRules
 def splitTarget : RuleTac := RuleTac.ofSingleRuleTac λ input => do
   let (some goals, steps) ← splitTargetS? input.goal |>.run | throwError
     "nothing to split in target"
+  let goals ← goals.mapM (mvarIdToSubgoal input.goal · ∅)
   return (goals, steps, none)
 
 partial def splitHypothesesCore (goal : MVarId) :
@@ -34,6 +35,7 @@ partial def splitHypothesesCore (goal : MVarId) :
 def splitHypotheses : RuleTac := RuleTac.ofSingleRuleTac λ input => do
   let (some goals, steps) ← splitHypothesesCore input.goal |>.run
     | throwError "no splittable hypothesis found"
+  let goals ← goals.mapM (mvarIdToSubgoal input.goal · ∅)
   return (goals, steps, none)
 
 end Aesop.BuiltinRules

Aesop/BuiltinRules/Subst.lean

@@ -90,6 +90,9 @@ def subst : RuleTac := RuleTac.ofSingleRuleTac λ input =>
       | _ => throwError "unexpected index match location"
     let (some goal, steps) ← substFVars input.goal hyps |>.run
       | throwError "no suitable hypothesis found"
+    -- TODO we can construct a better diff here, and doing so would be important
+    -- since `subst` often renames fvars.
+    let goal ← mvarIdToSubgoal input.goal goal ∅
     return (#[goal], steps, none)
 
 end Aesop.BuiltinRules

Aesop/RuleTac/Apply.lean

@@ -23,9 +23,15 @@ def applyExpr' (goal : MVarId) (e : Expr) (eStx : Term)
         pat.specializeRule patInst e
       else
         pure e
-    let (_, #[step]) ← applyS goal e eStx md |>.run
+    let (goals, #[step]) ← applyS goal e eStx md |>.run
       | throwError "aesop: internal error in applyExpr': multiple steps"
-    return .ofLazyScriptStep step none
+    let goals := goals.map λ mvarId => { mvarId, diff := ∅ }
+    return {
+      goals
+      postState := step.postState
+      scriptSteps? := #[step]
+      successProbability? := none
+    }
 
 def applyExpr (goal : MVarId) (e : Expr) (eStx : Term)
     (pat? : Option RulePattern) (patInsts : HashSet RulePatternInstantiation)

Aesop/RuleTac/Basic.lean

@@ -7,6 +7,8 @@ Authors: Jannis Limperg
 import Aesop.Index.Basic
 import Aesop.Options
 import Aesop.Percent
+import Aesop.RuleTac.GoalDiff
+import Aesop.RuleTac.FVarIdSubst
 import Aesop.Script.CtorNames
 import Aesop.Script.Step
 import Batteries.Lean.Meta.SavedState
@@ -20,6 +22,8 @@ namespace Aesop
 
 /-! # Rule Tactic Types -/
 
+-- TODO put docs on the structure fields instead of the structures
+
 /--
 Input for a rule tactic. Contains:
 
@@ -45,6 +49,23 @@ structure RuleTacInput where
   options : Options'
   deriving Inhabited
 
+/-- A subgoal produced by a rule. -/
+structure Subgoal where
+  /-- The goal mvar. -/
+  mvarId : MVarId
+  /--
+  A diff between the goal the rule was run on and this goal. Many `MetaM`
+  tactics report information that allows you to easily construct a `GoalDiff`.
+  If you don't have access to such information, use `diffGoals`, but note that
+  it may not give optimal results.
+  -/
+  diff : GoalDiff
+  deriving Inhabited
+
+def mvarIdToSubgoal (parentMVarId mvarId : MVarId) (fvarSubst : FVarIdSubst) :
+    MetaM Subgoal :=
+  return { mvarId, diff := ← diffGoals parentMVarId mvarId fvarSubst }
+
 /--
 A single rule application, representing the application of a tactic to the input
 goal. Must accurately report the following information:
@@ -59,27 +80,23 @@ goal. Must accurately report the following information:
   `none`, we use the success probability of the applied rule.
 -/
 structure RuleApplication where
-  goals : Array MVarId
+  goals : Array Subgoal
   postState : Meta.SavedState
   scriptSteps? : Option (Array Script.LazyStep)
   successProbability? : Option Percent
 
 namespace RuleApplication
 
-def check (r : RuleApplication) : MetaM (Option MessageData) :=
+def check (r : RuleApplication) (input : RuleTacInput) :
+    MetaM (Option MessageData) :=
   r.postState.runMetaM' do
     for goal in r.goals do
-      if ← goal.isAssignedOrDelayedAssigned then
-        return some m!"subgoal metavariable {goal.name} is already assigned."
+      if ← goal.mvarId.isAssignedOrDelayedAssigned then
+        return some m!"subgoal metavariable ?{goal.mvarId.name} is already assigned."
+      if let some err ← goal.diff.check input.goal goal.mvarId then
+        return some err
     return none
 
-def ofLazyScriptStep (step : Script.LazyStep)
-    (successProbability? : Option Percent) : RuleApplication where
-  goals := step.postGoals
-  postState := step.postState
-  scriptSteps? := #[step]
-  successProbability? := successProbability?
-
 end RuleApplication
 
 /--
@@ -102,7 +119,7 @@ A `RuleTac` which generates only a single `RuleApplication`.
 -/
 def SingleRuleTac :=
   RuleTacInput →
-  MetaM (Array MVarId × Option (Array Script.LazyStep) × Option Percent)
+  MetaM (Array Subgoal × Option (Array Script.LazyStep) × Option Percent)
 
 @[inline]
 def SingleRuleTac.toRuleTac (t : SingleRuleTac) : RuleTac := λ input => do
@@ -125,6 +142,7 @@ def RuleTac.ofTacticSyntax (t : RuleTacInput → MetaM Syntax.Tactic) : RuleTac
       tacticBuilders := #[return .unstructured stx]
       preState, postState, postGoals
     }
+    let postGoals ← postGoals.mapM (mvarIdToSubgoal input.goal · ∅)
     return (postGoals, some #[step], none)
 
 /--

Aesop/RuleTac/Cases.lean

@@ -36,7 +36,7 @@ namespace RuleTac
 partial def cases (target : CasesTarget) (md : TransparencyMode)
     (isRecursiveType : Bool) (ctorNames : Array CtorNames) : RuleTac :=
   SingleRuleTac.toRuleTac λ input => do
-    match ← go #[] #[] input.goal |>.run with
+    match ← go #[] #[] ∅ input.goal |>.run with
     | (none, _) => throwError "No matching hypothesis found."
     | (some goals, steps) => return (goals, steps, none)
   where
@@ -58,30 +58,25 @@ partial def cases (target : CasesTarget) (md : TransparencyMode)
           else
             return none
 
-    go (newGoals : Array MVarId) (excluded : Array FVarId) (goal : MVarId) :
-        ScriptM (Option (Array MVarId)) := do
+    go (newGoals : Array Subgoal) (excluded : Array FVarId) (diff : GoalDiff)
+        (goal : MVarId) : ScriptM (Option (Array Subgoal)) := do
       let some hyp ← findFirstApplicableHyp excluded goal
         | return none
       let some goals ← tryCasesS goal hyp ctorNames
         | return none
       let mut newGoals := newGoals
       for h : i in [:goals.size] do
         let g := goals[i]
-        let excluded :=
-          if ! isRecursiveType then
-            excluded
-          else
-            let excluded := excluded.map λ fvarId =>
-              match g.subst.find? fvarId with
-              | some (.fvar fvarId' ..) => fvarId'
-              | _ => fvarId
-            let fields := g.fields.filterMap λ
-              | (.fvar fvarId' ..) => some fvarId'
-              | _ => none
-            excluded ++ fields
-        match ← go newGoals excluded g.mvarId with
+        let newDiff ← diffGoals goal g.mvarId
+          (.ofFVarSubstIgnoringNonFVarIds g.subst)
+        let mut excluded := excluded
+        if isRecursiveType then
+          excluded :=
+            excluded.map diff.fvarSubst.get ++ newDiff.addedFVars.toArray
+        let diff := diff.comp newDiff
+        match ← go newGoals excluded diff g.mvarId with
         | some newGoals' => newGoals := newGoals'
-        | none => newGoals := newGoals.push g.mvarId
+        | none => newGoals := newGoals.push { mvarId := g.mvarId, diff }
       return some newGoals
 
 end Aesop.RuleTac

Aesop/RuleTac/FVarIdSubst.lean

@@ -0,0 +1,80 @@
+/-
+Copyright (c) 2024 Jannis Limperg. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jannis Limperg
+-/
+
+import Batteries.Data.HashMap
+
+open Lean Lean.Meta
+
+namespace Aesop
+
+structure FVarIdSubst where
+  map : Batteries.HashMap FVarId FVarId
+  deriving Inhabited
+
+namespace FVarIdSubst
+
+def isEmpty (s : FVarIdSubst) : Bool :=
+  s.map.isEmpty
+
+def contains (s : FVarIdSubst) (fvarId : FVarId) : Bool :=
+  s.map.contains fvarId
+
+def find? (s : FVarIdSubst) (fvarId : FVarId) : Option FVarId :=
+  s.map.find? fvarId
+
+def get (s : FVarIdSubst) (fvarId : FVarId) : FVarId :=
+  s.find? fvarId |>.getD fvarId
+
+def apply (s : FVarIdSubst) (e : Expr) : Expr :=
+  if s.isEmpty || ! e.hasFVar then
+    e
+  else
+    e.replace λ
+      | .fvar fvarId => s.find? fvarId |>.map mkFVar
+      | _ => none
+
+def applyToLocalDecl (s : FVarIdSubst) : LocalDecl → LocalDecl
+  | .cdecl i id n t bi k   => .cdecl i id n (s.apply t) bi k
+  | .ldecl i id n t v nd k => .ldecl i id n (s.apply t) (s.apply v) nd k
+
+def domain (s : FVarIdSubst) : HashSet FVarId :=
+  s.map.fold (init := ∅) λ r k _ => r.insert k
+
+def codomain (s : FVarIdSubst) : HashSet FVarId :=
+  s.map.fold (init := ∅) λ r _ v => r.insert v
+
+protected def empty : FVarIdSubst :=
+  ⟨∅⟩
+
+instance : EmptyCollection FVarIdSubst :=
+  ⟨.empty⟩
+
+def insert (s : FVarIdSubst) (old new : FVarId) : FVarIdSubst :=
+  let map : Batteries.HashMap _ _ := s.map.mapVal λ _ v =>
+    if v == old then new else v
+  ⟨map.insert old new⟩
+
+def erase (s : FVarIdSubst) (fvarId : FVarId) : FVarIdSubst :=
+  ⟨s.map.erase fvarId⟩
+
+def append (s t : FVarIdSubst) : FVarIdSubst :=
+  let map := s.map.mapVal λ _ v => t.get v
+  ⟨t.map.fold (init := map) λ s k v => s.insert k v⟩
+
+def ofFVarSubstIgnoringNonFVarIds (s : FVarSubst) : FVarIdSubst := .mk $
+  s.map.foldl (init := ∅) λ map k v =>
+    if let .fvar fvarId := v then map.insert k fvarId else map
+
+def ofFVarSubst? (s : FVarSubst) : Option FVarIdSubst := Id.run do
+  let mut result := ∅
+  for (k, v) in s.map do
+    if let .fvar fvarId := v then
+      result := result.insert k fvarId
+    else
+      return none
+  return some result
+
+end Aesop.FVarIdSubst

Aesop/RuleTac/Forward.lean

@@ -89,8 +89,8 @@ partial def makeForwardHyps (e : Expr) (pat? : Option RulePattern)
           return (proofsAcc, usedHypsAcc, proofTypesAcc)
 
 def assertForwardHyp (goal : MVarId) (hyp : Hypothesis) (depth : Nat)
-    (md : TransparencyMode) : ScriptM MVarId := do
-  withScriptStep goal (#[·]) (λ _ => true) tacticBuilder do
+    (md : TransparencyMode) : ScriptM (FVarId × MVarId) := do
+  withScriptStep goal (λ (_, g) => #[g]) (λ _ => true) tacticBuilder do
   withTransparency md do
     let hyp := {
       hyp with
@@ -103,14 +103,16 @@ def assertForwardHyp (goal : MVarId) (hyp : Hypothesis) (depth : Nat)
         binderInfo := .default
         kind := .implDetail
     }
-    (·.snd) <$> goal.assertHypotheses' #[hyp, implDetailHyp]
+    let (#[fvarId, _], goal) ← goal.assertHypotheses' #[hyp, implDetailHyp]
+      | throwError "aesop: internal error in assertForwardHyp: unexpected number of asserted fvars"
+    return (fvarId, goal)
 where
   tacticBuilder _ := Script.TacticBuilder.assertHypothesis goal hyp md
 
 def applyForwardRule (goal : MVarId) (e : Expr) (pat? : Option RulePattern)
     (patInsts : HashSet RulePatternInstantiation)
     (immediate : UnorderedArraySet Nat) (clear : Bool)
-    (md : TransparencyMode) (maxDepth? : Option Nat) : ScriptM MVarId :=
+    (md : TransparencyMode) (maxDepth? : Option Nat) : ScriptM Subgoal :=
   withTransparency md $ goal.withContext do
     let forwardHypData ← getForwardHypData
     let mut newHypProofs := #[]
@@ -136,12 +138,23 @@ def applyForwardRule (goal : MVarId) (e : Expr) (pat? : Option RulePattern)
       let type ← inferType proof
       newHyps := newHyps.push ({ value := proof, type, userName }, depth)
     let mut goal := goal
+    let mut addedFVars := ∅
     for (newHyp, depth) in newHyps do
-      goal ← assertForwardHyp goal newHyp depth md
+      let (fvarId, goal') ← assertForwardHyp goal newHyp depth md
+      goal := goal'
+      addedFVars := addedFVars.insert fvarId
+    let mut diff := {
+      addedFVars
+      removedFVars := ∅
+      fvarSubst := ∅
+    }
     if clear then
-      let (goal', _) ← tryClearManyS goal usedHyps
+      let (goal', removedFVars) ← tryClearManyS goal usedHyps
+      let removedFVars := removedFVars.foldl (init := ∅) λ set fvarId =>
+        set.insert fvarId
+      diff := { diff with removedFVars := removedFVars }
       goal := goal'
-    return goal
+    return { mvarId := goal, diff }
   where
     err {α} : MetaM α := throwError
       "found no instances of {e} (other than possibly those which had been previously added by forward rules)"