feat: improve grind preprocessor

src/Init/Grind.lean

@@ -6,3 +6,4 @@ Authors: Leonardo de Moura
 prelude
 import Init.Grind.Norm
 import Init.Grind.Tactics
+import Init.Grind.Lemmas

src/Init/Grind/Lemmas.lean

@@ -0,0 +1,14 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Init.Core
+
+namespace Lean.Grind
+
+theorem intro_with_eq (p p' q : Prop) (he : p = p') (h : p' → q) : p → q :=
+  fun hp => h (he.mp hp)
+
+end Lean.Grind

src/Lean/Elab/Tactic/Omega/OmegaM.lean

@@ -71,7 +71,7 @@ abbrev OmegaM := StateRefT Cache OmegaM'
 
 /-- Run a computation in the `OmegaM` monad, starting with no recorded atoms. -/
 def OmegaM.run (m : OmegaM α) (cfg : OmegaConfig) : MetaM α :=
-  m.run' HashMap.empty |>.run' {} { cfg } |>.run
+  m.run' HashMap.empty |>.run' {} { cfg } |>.run'
 
 /-- Retrieve the user-specified configuration options. -/
 def cfg : OmegaM OmegaConfig := do pure (← read).cfg

src/Lean/Meta/Canonicalizer.lean

@@ -63,8 +63,11 @@ abbrev CanonM := ReaderT TransparencyMode $ StateRefT State MetaM
 The definitionally equality tests are performed using the given transparency mode.
 We claim `TransparencyMode.instances` is a good setting for most applications.
 -/
-def CanonM.run (x : CanonM α) (transparency := TransparencyMode.instances) : MetaM α :=
-  StateRefT'.run' (x transparency) {}
+def CanonM.run' (x : CanonM α) (transparency := TransparencyMode.instances) (s : State := {}) : MetaM α :=
+  StateRefT'.run' (x transparency) s
+
+def CanonM.run (x : CanonM α) (transparency := TransparencyMode.instances) (s : State := {}) : MetaM (α × State) :=
+  StateRefT'.run (x transparency) s
 
 private partial def mkKey (e : Expr) : CanonM UInt64 := do
   if let some hash := unsafe (← get).cache.find? { e } then

src/Lean/Meta/Tactic/Grind.lean

@@ -6,6 +6,6 @@ Authors: Leonardo de Moura
 prelude
 import Lean.Meta.Tactic.Grind.Attr
 import Lean.Meta.Tactic.Grind.RevertAll
-import Lean.Meta.Tactic.Grind.EnsureNoMVar
 import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.Tactic.Grind.Preprocessor
+import Lean.Meta.Tactic.Grind.Util

src/Lean/Meta/Tactic/Grind/Preprocessor.lean

@@ -4,14 +4,15 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import Init.Grind.Lemmas
 import Lean.Meta.Canonicalizer
 import Lean.Meta.Tactic.Util
 import Lean.Meta.Tactic.Intro
 import Lean.Meta.Tactic.Simp.Main
 import Lean.Meta.Tactic.Grind.Attr
 import Lean.Meta.Tactic.Grind.RevertAll
-import Lean.Meta.Tactic.Grind.EnsureNoMVar
 import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.Util
 
 namespace Lean.Meta.Grind
 namespace Preprocessor
@@ -25,99 +26,113 @@ structure Context where
   deriving Inhabited
 
 structure State where
-  canon : Canonicalizer.State := {}
-  todo  : List Goal := []
-  goals : PArray Goal := {}
   simpStats : Simp.Stats := {}
   deriving Inhabited
 
-structure Result where
-  canon : Canonicalizer.State := {}
-  goals : PArray Goal := {}
-  deriving Inhabited
-
-abbrev M := ReaderT Context $ StateRefT State MetaM
-
-def mkInitialState (mvarId : MVarId) : State :=
-  { todo := [ mkGoal mvarId ] }
+abbrev PreM := ReaderT Context $ StateRefT State GrindM
 
-def M.run (x : M α) (mvarId : MVarId) : MetaM α := do
+def PreM.run (x : PreM α) : GrindM α := do
   let thms ← grindNormExt.getTheorems
   let simprocs := #[(← grindNormSimprocExt.getSimprocs)]
   let simp : Simp.Context := {
+    config := { arith := true }
     simpTheorems := #[thms]
     congrTheorems := (← getSimpCongrTheorems)
   }
-  x { simp, simprocs } |>.run' (mkInitialState mvarId)
+  x { simp, simprocs } |>.run' {}
 
-def simpHyp? (mvarId : MVarId) (fvarId : FVarId) : M (Option (FVarId × MVarId)) := do
+def simp (e : Expr) : PreM Simp.Result := do
+  let simpStats := (← get).simpStats
+  let (r, simpStats) ← Meta.simp e (← read).simp (← read).simprocs (stats := simpStats)
+  modify fun s => { s with simpStats }
+  return r
+
+def simpHyp? (mvarId : MVarId) (fvarId : FVarId) : PreM (Option (FVarId × MVarId)) := do
   let simpStats := (← get).simpStats
   let (result, simpStats) ← simpLocalDecl mvarId fvarId (← read).simp (← read).simprocs (stats := simpStats)
   modify fun s => { s with simpStats }
   return result
 
-def getNextGoal? : M (Option Goal) := do
-  match (← get).todo with
-  | [] => return none
-  | goal :: todo =>
-    modify fun s => { s with todo }
-    return some goal
-
 inductive IntroResult where
-  | done | closed
+  | done
   | newHyp (fvarId : FVarId) (goal : Goal)
   | newLocal (fvarId : FVarId) (goal : Goal)
 
-def introNext (goal : Goal) : M IntroResult := do
+def introNext (goal : Goal) : PreM IntroResult := do
   let target ← goal.mvarId.getType
   if target.isArrow then
-    let (fvarId, mvarId) ← goal.mvarId.intro1P
-    -- TODO: canonicalize subterms
-    mvarId.withContext do
-    if (← isProp (← fvarId.getType)) then
-      let some (fvarId, mvarId) ← simpHyp? mvarId fvarId | return .closed
-      return .newHyp fvarId { goal with mvarId }
-    else
-      return .newLocal fvarId { goal with mvarId }
+    goal.mvarId.withContext do
+      let p := target.bindingDomain!
+      if !(← isProp p) then
+        let (fvarId, mvarId) ← goal.mvarId.intro1P
+        return .newLocal fvarId { goal with mvarId }
+      else
+        let tag ← goal.mvarId.getTag
+        let q := target.bindingBody!
+        let r ← simp p
+        let p' := r.expr
+        let p' ← canon p'
+        let p' ← shareCommon p'
+        let fvarId ← mkFreshFVarId
+        let lctx := (← getLCtx).mkLocalDecl fvarId target.bindingName! p' target.bindingInfo!
+        let mvarNew ← mkFreshExprMVarAt lctx (← getLocalInstances) q .syntheticOpaque tag
+        let mvarIdNew := mvarNew.mvarId!
+        mvarIdNew.withContext do
+          let h ← mkLambdaFVars #[mkFVar fvarId] mvarNew
+          match r.proof? with
+          | some he =>
+            let hNew := mkAppN (mkConst ``Lean.Grind.intro_with_eq) #[p, p', q, he, h]
+            goal.mvarId.assign hNew
+            return .newHyp fvarId { goal with mvarId := mvarIdNew }
+          | none =>
+            -- `p` and `p'` are definitionally equal
+            goal.mvarId.assign h
+            return .newHyp fvarId { goal with mvarId := mvarIdNew }
   else if target.isLet || target.isForall then
-    -- TODO: canonicalize subterms
-    -- TODO: If forall is of the form `∀ h : <proposition>, A h`, generalize `h`.
     let (fvarId, mvarId) ← goal.mvarId.intro1P
-    return .newLocal fvarId { goal with mvarId }
+    mvarId.withContext do
+      let localDecl ← fvarId.getDecl
+      if (← isProp localDecl.type) then
+        -- Add a non-dependent copy
+        let mvarId ← mvarId.assert localDecl.userName localDecl.type (mkFVar fvarId)
+        return .newLocal fvarId { goal with mvarId }
+      else
+        return .newLocal fvarId { goal with mvarId }
   else
     return .done
 
-def pushTodo (goal : Goal) : M Unit :=
-  modify fun s => { s with todo := goal :: s.todo }
+def pushResult (goal : Goal) : PreM Unit :=
+  modifyThe Grind.State fun s => { s with goals := s.goals.push goal }
 
-def pushResult (goal : Goal) : M Unit :=
-  modify fun s => { s with goals := s.goals.push goal }
+partial def preprocess (goal : Goal) : PreM Unit := do
+  trace[Meta.debug] "{goal.mvarId}"
+  match (← introNext goal) with
+  | .done =>
+    if let some mvarId ← goal.mvarId.byContra? then
+      preprocess { goal with mvarId }
+    else
+      pushResult goal
+  | .newHyp fvarId goal =>
+    -- TODO: apply eliminators
+    let clause ← goal.mvarId.withContext do mkInputClause fvarId
+    preprocess { goal with clauses := goal.clauses.push clause }
+  | .newLocal _ goal =>
+    -- TODO: apply eliminators
+    preprocess goal
 
-partial def main (mvarId : MVarId) : MetaM Result := do
+end Preprocessor
+
+open Preprocessor
+
+partial def main (mvarId : MVarId) (mainDeclName : Name) : MetaM Grind.State := do
+  mvarId.ensureProp
   mvarId.ensureNoMVar
   let mvarId ← mvarId.revertAll
   mvarId.ensureNoMVar
-  let s ← (loop *> get) |>.run mvarId
-  return { s with }
-where
-  loop : M Unit := do
-    let some goal ← getNextGoal? | return ()
-    trace[Meta.debug] "{goal.mvarId}"
-    match (← introNext goal) with
-    | .closed => loop
-    | .done =>
-      -- TODO: apply `byContradiction`
-      pushResult goal
-      return ()
-    | .newHyp fvarId goal =>
-      -- TODO: apply eliminators
-      let clause ← goal.mvarId.withContext do mkInputClause fvarId
-      pushTodo { goal with clauses := goal.clauses.push clause }
-      loop
-    | .newLocal _ goal =>
-      -- TODO: apply eliminators
-      pushTodo goal
-      loop
+  let mvarId ← mvarId.abstractNestedProofs mainDeclName
+  let mvarId ← mvarId.unfoldReducible
+  let mvarId ← mvarId.betaReduce
+  let s ← (preprocess { mvarId } *> getThe Grind.State) |>.run |>.run mainDeclName
+  return s
 
-end Preprocessor
 end Lean.Meta.Grind

src/Lean/Meta/Tactic/Grind/Types.lean

@@ -4,10 +4,46 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import Lean.Util.ShareCommon
 import Lean.Meta.Basic
+import Lean.Meta.AbstractNestedProofs
+import Lean.Meta.Canonicalizer
 import Lean.Meta.Tactic.Util
 
 namespace Lean.Meta.Grind
+/--
+Stores information for a node in the egraph.
+Each internalized expression `e` has an `ENode` associated with it.
+-/
+structure ENode where
+  /-- Next element in the equivalence class. -/
+  next : Expr
+  /-- Root (aka canonical representative) of the equivalence class -/
+  root : Expr
+  /-- Root of the congruence class. This is field is a don't care if `e` is not an application. -/
+  cgRoot : Expr
+  /--
+  When `e` was added to this equivalence class because of an equality `h : e = target`,
+  then we store `target` here, and `h` at `proof?`.
+  -/
+  target? : Option Expr := none
+  proof? : Option Expr := none
+  /-- Proof has been flipped. -/
+  flipped : Bool := false
+  /-- Number of elements in the equivalence class, this field is meaningless if node is not the root. -/
+  size : Nat := 1
+  /-- `interpreted := true` if node should be viewed as an abstract value. -/
+  interpreted : Bool := false
+  /-- `ctor := true` if the head symbol is a constructor application. -/
+  ctor : Bool := false
+  /-- `hasLambdas := true` if equivalence class contains lambda expressions. -/
+  hasLambdas : Bool := false
+  /--
+  If `heqProofs := true`, then some proofs in the equivalence class are based
+  on heterogeneous equality.
+  -/
+  heqProofs : Bool := false
+  -- TODO: see Lean 3 implementation
 
 structure Clause where
   expr  : Expr
@@ -20,6 +56,8 @@ def mkInputClause (fvarId : FVarId) : MetaM Clause :=
 structure Goal where
   mvarId  : MVarId
   clauses : PArray Clause := {}
+  enodes  : PHashMap UInt64 ENode := {}
+  -- TODO: occurrences for propagation, etc
   deriving Inhabited
 
 def mkGoal (mvarId : MVarId) : Goal :=
@@ -28,4 +66,45 @@ def mkGoal (mvarId : MVarId) : Goal :=
 def Goal.admit (goal : Goal) : MetaM Unit :=
   goal.mvarId.admit
 
+structure Context where
+  mainDeclName : Name
+
+structure State where
+  canon      : Canonicalizer.State := {}
+  /-- `ShareCommon` (aka `Hashconsing`) state. -/
+  scState    : ShareCommon.State.{0} ShareCommon.objectFactory := ShareCommon.State.mk _
+  /-- Next index for creating auxiliary theorems. -/
+  nextThmIdx : Nat := 1
+  goals      : PArray Goal := {}
+
+abbrev GrindM := ReaderT Context $ StateRefT State MetaM
+
+def GrindM.run (x : GrindM α) (mainDeclName : Name) : MetaM α :=
+  x { mainDeclName } |>.run' {}
+
+def abstractNestedProofs (e : Expr) : GrindM Expr := do
+  let nextIdx := (← get).nextThmIdx
+  let (e, s') ← AbstractNestedProofs.visit e |>.run { baseName := (← read).mainDeclName } |>.run |>.run { nextIdx }
+  modify fun s => { s with nextThmIdx := s'.nextIdx }
+  return e
+
+def shareCommon (e : Expr) : GrindM Expr := do
+  modifyGet fun { canon, scState, nextThmIdx, goals } =>
+    let (e, scState) := ShareCommon.State.shareCommon scState e
+    (e, { canon, scState, nextThmIdx, goals })
+
+def canon (e : Expr) : GrindM Expr := do
+  let canonS ← modifyGet fun s => (s.canon, { s with canon := {} })
+  let (e, canonS) ← Canonicalizer.CanonM.run (canonRec e) (s := canonS)
+  modify fun s => { s with canon := canonS }
+  return e
+where
+  canonRec (e : Expr) : CanonM Expr := do
+    let post (e : Expr) : CanonM TransformStep := do
+      if e.isApp then
+        return .done (← Meta.canon e)
+      else
+        return .done e
+    transform e post
+
 end Lean.Meta.Grind