Introduce 'Premise's to capture hypotheses which are not rewrites

Lean/Egg.lean

@@ -9,18 +9,20 @@ import Egg.Core.Gen.TcSpecs
 import Egg.Core.MVars.Ambient
 import Egg.Core.MVars.Basic
 import Egg.Core.MVars.Subst
+import Egg.Core.Premise.Basic
+import Egg.Core.Premise.Facts
+import Egg.Core.Premise.Rewrites
 import Egg.Core.Config
 import Egg.Core.Congr
 import Egg.Core.Directions
 import Egg.Core.Guides
 import Egg.Core.Request
-import Egg.Core.Rewrites
 import Egg.Core.Source
 import Egg.Tactic.Config.Modifier
 import Egg.Tactic.Config.Option
 import Egg.Tactic.Base
 import Egg.Tactic.Basic
 import Egg.Tactic.Explanation
 import Egg.Tactic.Guides
-import Egg.Tactic.Rewrites
+import Egg.Tactic.Premises
 import Egg.Tactic.Trace

Lean/Egg/Core/Encode/Rewrites.lean

@@ -1,5 +1,5 @@
 import Egg.Core.Encode.Basic
-import Egg.Core.Rewrites
+import Egg.Core.Premise.Rewrites
 import Lean
 open Lean
 

Lean/Egg/Core/Explanation/Proof.lean

@@ -1,6 +1,6 @@
  import Egg.Core.Explanation.Basic
 import Egg.Core.Explanation.Congr
-import Egg.Core.Rewrites
+import Egg.Core.Premise.Rewrites
 open Lean Meta
 
 namespace Egg.Explanation

Lean/Egg/Core/Gen/TcProjs.lean

@@ -1,4 +1,4 @@
-import Egg.Core.Rewrites
+import Egg.Core.Premise.Basic
 import Egg.Core.Guides
 import Lean
 open Lean Meta
@@ -24,7 +24,7 @@ private def TcProj.reductionRewrite?
   if proj == reducedNorm then return none
   let eq ← mkEq proj reducedNorm
   let proof ← mkEqRefl proj
-  let some rw ← Rewrite.from? proof eq src none amb
+  let .rw rw ← Premise.from proof eq src none amb
     | throwError "egg: internal error in 'TcProj.reductionRewrite?'"
   return rw
 

Lean/Egg/Core/Gen/TcSpecs.lean

@@ -1,4 +1,4 @@
-import Egg.Core.Rewrites
+import Egg.Core.Premise.Rewrites
 import Std.Tactic.Exact
 import Lean
 open Lean Meta

Lean/Egg/Core/Premise/Basic.lean

@@ -0,0 +1,37 @@
+import Egg.Core.Premise.Facts
+import Egg.Core.Premise.Rewrites
+import Lean
+open Lean Meta
+
+namespace Egg
+
+inductive Premise where
+  | rw   (rw : Rewrite)
+  | fact (f : Fact)
+
+namespace Premise
+
+-- Note: It isn't sufficient to take the `args` as a rewrite's holes, as implicit arguments will
+--       already be instantiated as mvars during type inference. For example, the type of
+--       `theorem t : ∀ {x}, x + 0 = 0 + x := Nat.add_comm _ _` will be directly inferred as
+--       `?x + 0 = 0 + ?x`. On the other hand, we might be collecting too many mvars right now as a
+--       rewrite could possibly contain mvars which weren't quantified (e.g. if it comes from the
+--       local context). Also, we need to "catch loose args", that is, those which are preconditions
+--       for the rewrite, but don't appear in the body (as in conditional rewrites).
+--
+-- Note: We must instantiate mvars of the rewrite's type. For an example that breaks otherwise, see
+--       leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Different.20elab.20results
+def «from»
+    (proof : Expr) (type : Expr) (src : Source) (normalize : Option Config.Normalization)
+    (amb : MVars.Ambient) : MetaM Premise := do
+  let mut (args, _, type) ← forallMetaTelescope (← instantiateMVars type)
+  type ← if let some cfg := normalize then Egg.normalize type cfg else pure type
+  let proof := mkAppN proof args
+  let some cgr ← Congr.from? type | return .fact { src, type, proof }
+  let lhsMVars := (← MVars.collect cgr.lhs).remove amb
+  let rhsMVars := (← MVars.collect cgr.rhs).remove amb
+  let conds := looseArgs args lhsMVars rhsMVars
+  return .rw { cgr with proof, src, conds, lhsMVars, rhsMVars }
+where
+  looseArgs (args : Array Expr) (lhsMVars rhsMVars : MVars) : Array Expr :=
+    args.filter fun a => !lhsMVars.expr.contains a.mvarId! && !rhsMVars.expr.contains a.mvarId!

Lean/Egg/Core/Premise/Facts.lean

@@ -0,0 +1,12 @@
+import Egg.Core.Source
+import Lean
+open Lean
+
+namespace Egg
+
+structure Fact where
+  src   : Source
+  type  : Expr
+  proof : Expr
+
+abbrev Facts := Array Fact

Lean/Egg/Core/Rewrites.lean → Lean/Egg/Core/Premise/Rewrites.lean

@@ -7,43 +7,20 @@ import Egg.Core.Source
 import Egg.Lean
 open Lean Meta
 
-namespace Egg.Rewrite
+namespace Egg
 
-structure _root_.Egg.Rewrite extends Congr where
-  private mk ::
+structure Rewrite extends Congr where
   proof    : Expr
   src      : Source
+  conds    : Array Expr
   lhsMVars : MVars
   rhsMVars : MVars
   deriving Inhabited
 
--- Note: It isn't sufficient to take the `args` as a rewrite's holes, as implicit arguments will
---       already be instantiated as mvars during type inference. For example, the type of
---       `theorem t : ∀ {x}, x + 0 = 0 + x := Nat.add_comm _ _` will be directly inferred as
---       `?x + 0 = 0 + ?x`. On the other hand, we might be collecting too many mvars right now as a
---       rewrite could possibly contain mvars which weren't quantified (e.g. if it comes from the
---       local context). Also, we need to "catch loose args", that is, those which are preconditions
---       for the rewrite, but don't appear in the body (as in conditional rewrites).
---
--- Note: We must instantiate mvars of the rewrite's type. For an example that breaks otherwise, see
---       leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Different.20elab.20results
-def from?
-    (proof : Expr) (type : Expr) (src : Source) (normalize : Option Config.Normalization)
-    (amb : MVars.Ambient) : MetaM (Option Rewrite) := do
-  let mut (args, _, type) ← forallMetaTelescope (← instantiateMVars type)
-  type ← if let some cfg := normalize then Egg.normalize type cfg else pure type
-  let proof := mkAppN proof args
-  let some cgr ← Congr.from? type | return none
-  let lhsMVars := (← MVars.collect cgr.lhs).remove amb
-  let rhsMVars := (← MVars.collect cgr.rhs).remove amb
-  catchLooseArgs args lhsMVars rhsMVars
-  return some { cgr with proof, src, lhsMVars, rhsMVars }
-where
-  catchLooseArgs (args : Array Expr) (lhsMVars rhsMVars : MVars) : MetaM Unit := do
-    for arg in args do
-      if lhsMVars.expr.contains arg.mvarId! then continue
-      if rhsMVars.expr.contains arg.mvarId! then continue
-      throwError m!"Rewrite {src.description} contains loose argument."
+namespace Rewrite
+
+def isConditional (rw : Rewrite) : Bool :=
+  !rw.conds.isEmpty
 
 def validDirs (rw : Rewrite) : Directions :=
   let exprDirs := Directions.satisfyingSuperset rw.lhsMVars.expr rw.rhsMVars.expr

Lean/Egg/Core/Request.lean

@@ -2,7 +2,6 @@ import Egg.Core.Encode.Rewrites
 import Egg.Core.Encode.Guides
 import Egg.Core.Config
 import Egg.Core.Explanation.Basic
-import Egg.Core.Rewrites
 open Lean
 
 namespace Egg.Request

Lean/Egg/Tactic/Basic.lean

@@ -7,7 +7,7 @@ import Egg.Tactic.Config.Modifier
 import Egg.Tactic.Explanation
 import Egg.Tactic.Base
 import Egg.Tactic.Guides
-import Egg.Tactic.Rewrites
+import Egg.Tactic.Premises
 import Egg.Tactic.Trace
 import Std.Tactic.Exact
 import Lean
@@ -42,24 +42,26 @@ where
     else
       throwError "expected goal to be of type '=' or '↔', but found:\n{← ppExpr goalType}"
 
-private def traceRewrites
-    (basic : Rewrites) (stx : Array Syntax) (tc : Rewrites) (cfg : Config.Gen) : TacticM Unit := do
+private def tracePremises (ps : Premises) (tc : Rewrites) (cfg : Config.Gen) : TacticM Unit := do
   let cls := `egg.rewrites
   withTraceNode cls (fun _ => return "Rewrites") do
-    withTraceNode cls (fun _ => return m!"Basic ({basic.size})") do basic.trace stx cls
+    withTraceNode cls (fun _ => return m!"Basic ({ps.rws.size})") do ps.rws.trace ps.rwsStx cls
     withTraceNode cls (fun _ => return m!"Generated ({tc.size})") do tc.trace #[] cls
     withTraceNode cls (fun _ => return "Definitional") do
       if cfg.genBetaRw    then Lean.trace cls fun _ => "β-Reduction"
       if cfg.genEtaRw     then Lean.trace cls fun _ => "η-Reduction"
       if cfg.genNatLitRws then Lean.trace cls fun _ => "Natural Number Literals"
+    withTraceNode cls (fun _ => return m!"Hypotheses ({ps.facts.size})") do
+      ps.facts.trace ps.factsStx cls
 
 private partial def genRewrites
-    (goal : Goal) (rws : TSyntax `egg_rws) (guides : Guides) (cfg : Config) (amb : MVars.Ambient) :
+    (goal : Goal) (ps : TSyntax `egg_prems) (guides : Guides) (cfg : Config) (amb : MVars.Ambient) :
     TacticM Rewrites := do
-  let (rws, stx) ← Rewrites.parse cfg.toNormalization amb rws
-  let tcRws ← genTcRws rws
-  traceRewrites rws stx tcRws cfg.toGen
-  return rws ++ tcRws
+  let ps ← Premises.parse cfg.toNormalization amb ps
+  let tcRws ← genTcRws ps.rws
+  tracePremises ps tcRws cfg.toGen
+  throwOnConditionalRw ps.rws ps.rwsStx
+  return ps.rws ++ tcRws
 where
   genTcRws (rws : Rewrites) : TacticM Rewrites := do
     let mut projTodo := #[]
@@ -81,6 +83,10 @@ where
         tcRws    := tcRws ++ specRws
     return tcRws
 
+  throwOnConditionalRw (rws : Rewrites) (stxs : Array Syntax) : TacticM Unit := do
+    for rw in rws, stx in stxs do
+      if rw.isConditional then throwErrorAt stx "egg does not currently support conditional rewrites"
+
 private def processRawExpl
     (rawExpl : Explanation.Raw) (goal : Goal) (rws : Rewrites) (amb : MVars.Ambient) :
     TacticM Expr := do
@@ -107,7 +113,7 @@ private def traceRequest (req : Request) : TacticM Unit := do
 
 open Config.Modifier (egg_cfg_mod)
 
-elab "egg " mod:egg_cfg_mod rws:egg_rws base:(egg_base)? guides:(egg_guides)? : tactic => do
+elab "egg " mod:egg_cfg_mod rws:egg_prems base:(egg_base)? guides:(egg_guides)? : tactic => do
   let goal ← getMainGoal
   let mod  ← Config.Modifier.parse mod
   let cfg := (← Config.fromOptions).modify mod

Lean/Egg/Tactic/Rewrites.lean → Lean/Egg/Tactic/Premises.lean

@@ -1,34 +1,55 @@
-import Egg.Core.Rewrites
+import Egg.Core.Premise.Basic
 import Lean
 
 open Lean Meta Elab Tactic
 
 namespace Egg
 
-declare_syntax_cat egg_rws_arg
-syntax "*"  : egg_rws_arg
-syntax term : egg_rws_arg
+declare_syntax_cat egg_prems_arg
+syntax "*"  : egg_prems_arg
+syntax term : egg_prems_arg
 
-declare_syntax_cat egg_rws_args
-syntax "[" egg_rws_arg,* "]": egg_rws_args
+declare_syntax_cat egg_prems_args
+syntax "[" egg_prems_arg,* "]": egg_prems_args
 
-declare_syntax_cat egg_rws
-syntax (egg_rws_args)? : egg_rws
+declare_syntax_cat egg_prems
+syntax (egg_prems_args)? : egg_prems
 
-namespace Rewrites
+structure Premises where
+  rws      : Rewrites     := #[]
+  rwsStx   : Array Syntax := #[]
+  facts    : Facts        := #[]
+  factsStx : Array Syntax := #[]
+
+namespace Premises
+
+private def push (ps : Premises) (stx : Syntax) : Premise → Premises
+  | .rw rw  => { ps with rws := ps.rws.push rw, rwsStx := ps.rwsStx.push stx }
+  | .fact f => { ps with facts := ps.facts.push f, factsStx := ps.factsStx.push stx }
+
+private def singleton (stx : Syntax)  (p : Premise) : Premises :=
+  ({} : Premises).push stx p
+
+private def append (ps₁ ps₂ : Premises) : Premises where
+  rws      := ps₁.rws.append ps₂.rws
+  rwsStx   := ps₁.rwsStx.append ps₂.rwsStx
+  facts    := ps₁.facts.append ps₂.facts
+  factsStx := ps₁.factsStx.append ps₂.factsStx
+
+instance : Append Premises where
+  append := append
 
 -- Note: We must use `Tactic.elabTerm`, not `Term.elabTerm`. Otherwise elaborating `‹...›` doesn't
 --       work correctly. Cf. https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Elaborate.20.E2.80.B9.2E.2E.2E.E2.80.BA
-partial def explicit
-    (arg : Term) (argIdx : Nat) (norm : Config.Normalization) (amb : MVars.Ambient) :
-    TacticM Rewrites := do
+partial def explicit (arg : Term) (argIdx : Nat) (norm : Config.Normalization) (amb : MVars.Ambient) :
+    TacticM Premises := do
   match ← elabArg arg with
-  | .inl (e, ty?) => return #[← mkRw e ty? none]
+  | .inl (e, ty?) => return Premises.singleton arg (← mkPremise e ty? none)
   | .inr eqns =>
-    let mut result : Rewrites := #[]
+    let mut result : Premises := {}
     for eqn in eqns, eqnIdx in [:eqns.size] do
       let e ← Tactic.elabTerm eqn none
-      result := result.push (← mkRw e none eqnIdx)
+      result := result.push arg (← mkPremise e none eqnIdx)
     return result
   -- We don't just elaborate the `arg` directly as:
   -- (1) this can cause problems for global constants with typeclass arguments, as Lean sometimes
@@ -39,12 +60,11 @@ where
   -- Note: When we infer the type of `e` it might not have the syntactic form we expect. For
   --       example, if `e` is `congrArg (fun x => x + 1) (_ : a = b)` then its type will be inferred
   --       as `a + 1 = b + 1` instead of `(fun x => x + 1) a = (fun x => x + 1) b`.
-  mkRw (e : Expr) (ty? : Option Expr) (eqnIdx? : Option Nat) : TacticM Rewrite := do
+  mkPremise (e : Expr) (ty? : Option Expr) (eqnIdx? : Option Nat) : TacticM Premise := do
     let src := .explicit argIdx eqnIdx?
     let ty := ty?.getD (← inferType e)
-    let some rw ← Rewrite.from? e ty src norm amb
-      | throwErrorAt arg "egg requires arguments to be equalities, equivalences or (non-propositional) definitions"
-    return rw
+    Premise.from e ty src norm amb
+
   elabArg (arg : Term) : TacticM (Sum (Expr × Option Expr) (Array Ident)) := do
     if let some hyp ← optional (getFVarId arg) then
       -- `arg` is a local declaration.
@@ -72,29 +92,26 @@ where
 -- Note: We need to filter out auxiliary declaration and implementation details, as they are not
 --       visible in the proof context and, for example, contain the declaration being defined itself
 --       (to enable recursive calls). Cf. https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/local.20context.20without.20current.20decl
-def star (norm : Config.Normalization) (amb : MVars.Ambient) : MetaM Rewrites := do
-  let mut result : Rewrites := #[]
+def star (stx : Syntax) (norm : Config.Normalization) (amb : MVars.Ambient) : MetaM Premises := do
+  let mut result : Premises := {}
   for decl in ← getLCtx do
     if decl.isImplementationDetail || decl.isAuxDecl then continue
-    if let some rw ← Rewrite.from? decl.toExpr decl.type (.star decl.fvarId) norm amb
-    then result := result.push rw
+    let src := Source.star decl.fvarId
+    result := result.push stx (← Premise.from decl.toExpr decl.type src norm amb)
   return result
 
-def parse (norm : Config.Normalization) (amb : MVars.Ambient) :
-    (TSyntax `egg_rws) → TacticM (Rewrites × Array Syntax)
-  | `(egg_rws|)          => return (#[], #[])
-  | `(egg_rws|[$args,*]) => do
-    let mut result : Rewrites := #[]
+def parse (norm : Config.Normalization) (amb : MVars.Ambient) : (TSyntax `egg_prems) → TacticM Premises
+  | `(egg_prems|)          => return {}
+  | `(egg_prems|[$args,*]) => do
+    let mut result : Premises := {}
     let mut noStar := true
     for arg in args.getElems, idx in [:args.getElems.size] do
       match arg with
-      | `(egg_rws_arg|$arg:term) =>
-        result := result ++ (← explicit arg idx norm amb)
-      | `(egg_rws_arg|*%$tk) =>
+      | `(egg_prems_arg|$arg:term) => result := result ++ (← explicit arg idx norm amb)
+      | `(egg_prems_arg|*%$tk) =>
         unless noStar do throwErrorAt tk "duplicate '*' in arguments to egg"
         noStar := false
-        result := result ++ (← star norm amb)
-      | _ =>
-        throwUnsupportedSyntax
-    return (result, args)
+        result := result ++ (← star tk norm amb)
+      | _ => throwUnsupportedSyntax
+    return result
   | _ => throwUnsupportedSyntax

Lean/Egg/Tactic/Trace.lean

@@ -1,6 +1,7 @@
 import Egg.Core.Request
 import Egg.Core.Explanation.Proof
 import Egg.Core.MVars.Ambient
+import Egg.Tactic.Premises
 import Lean
 open Lean Meta Elab Tactic Std Format
 
@@ -44,12 +45,22 @@ def Rewrite.trace (rw : Rewrite) (stx? : Option Syntax) (cls : Name) : TacticM U
   if let some stx := stx? then header := m!"{header}: {stx}"
   withTraceNode cls (fun _ => return header) do
     traceM cls fun _ => return m!"{← rw.toCongr.format}"
+    if !rw.conds.isEmpty then
+      withTraceNode cls (fun _ => return "Conditions") (collapsed := false) do
+        for cond in rw.conds do
+          traceM cls fun _ => return m!"{← cond.mvarId!.getType}"
     traceM cls fun _ => return m!"LHS MVars\n{← rw.lhsMVars.format}"
     traceM cls fun _ => return m!"RHS MVars\n{← rw.rhsMVars.format}"
 
 def Rewrites.trace (rws : Rewrites) (stx : Array Syntax) (cls : Name) : TacticM Unit := do
   for rw in rws, idx in [:rws.size] do rw.trace stx[idx]? cls
 
+nonrec def Fact.trace (f : Fact) (stx : Syntax) (cls : Name) : TacticM Unit := do
+  trace cls fun _ => m!"{f.src.description}: {stx} : {f.type}"
+
+def Facts.trace (fs : Facts) (stx : Array Syntax) (cls : Name) : TacticM Unit := do
+  for f in fs, s in stx do f.trace s cls
+
 def Rewrite.Encoded.trace (rw : Rewrite.Encoded) (cls : Name) : TacticM Unit := do
   let header := m!"{rw.name}({rw.dirs.format})"
   withTraceNode cls (fun _ => return header) do