From 419100d42b7b98ff76f25f7cc4d4ceddea1cc5a7 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Thu, 19 Oct 2023 08:49:44 -0700
Subject: [PATCH] feat: add `Simp.Config.ground` for simplifying nested ground
 terms

This is an experimental new feature. We need more bells and whistles,
and `cbv` tactic for improving its performance.
---
 src/Init/Meta.lean                  |  5 +++
 src/Lean/Meta/Tactic/Simp.lean      |  1 +
 src/Lean/Meta/Tactic/Simp/Main.lean | 51 +++++++++++++++++++++++++++--
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/src/Init/Meta.lean b/src/Init/Meta.lean
index 86f9335dea..aaa2b20263 100644
--- a/src/Init/Meta.lean
+++ b/src/Init/Meta.lean
@@ -1261,6 +1261,10 @@ structure Config where
   /-- If `failIfUnchanged := true`, then calls to `simp`, `dsimp`, or `simp_all`
   will fail if they do not make progress. -/
   failIfUnchanged   : Bool := true
+  /-- If `ground := true`, then ground terms are reduced. A term is ground when
+  it does not contain free or meta variables. Reduction is interrupted at a function application `f ...`
+  if `f` is marked to not be unfolded. -/
+  ground            : Bool := false
   deriving Inhabited, BEq, Repr
 
 -- Configuration object for `simp_all`
@@ -1276,6 +1280,7 @@ def neutralConfig : Simp.Config := {
   decide            := false
   arith             := false
   autoUnfold        := false
+  ground            := false
 }
 
 end Simp
diff --git a/src/Lean/Meta/Tactic/Simp.lean b/src/Lean/Meta/Tactic/Simp.lean
index 2e61e9bcfa..dca62c697c 100644
--- a/src/Lean/Meta/Tactic/Simp.lean
+++ b/src/Lean/Meta/Tactic/Simp.lean
@@ -17,6 +17,7 @@ builtin_initialize registerTraceClass `Meta.Tactic.simp.congr (inherited := true
 builtin_initialize registerTraceClass `Meta.Tactic.simp.discharge (inherited := true)
 builtin_initialize registerTraceClass `Meta.Tactic.simp.rewrite (inherited := true)
 builtin_initialize registerTraceClass `Meta.Tactic.simp.unify (inherited := true)
+builtin_initialize registerTraceClass `Meta.Tactic.simp.ground (inherited := true)
 builtin_initialize registerTraceClass `Meta.Tactic.simp.numSteps
 builtin_initialize registerTraceClass `Meta.Tactic.simp.heads
 builtin_initialize registerTraceClass `Debug.Meta.Tactic.simp
diff --git a/src/Lean/Meta/Tactic/Simp/Main.lean b/src/Lean/Meta/Tactic/Simp/Main.lean
index 4fca5307a0..a1b2c1f63b 100644
--- a/src/Lean/Meta/Tactic/Simp/Main.lean
+++ b/src/Lean/Meta/Tactic/Simp/Main.lean
@@ -7,6 +7,7 @@ import Lean.Meta.Transform
 import Lean.Meta.Tactic.Replace
 import Lean.Meta.Tactic.UnifyEq
 import Lean.Meta.Tactic.Simp.Rewrite
+import Lean.Meta.Match.Value
 
 namespace Lean.Meta
 namespace Simp
@@ -145,15 +146,59 @@ where
       let f := e.getAppFn
       f.isConst && isMatcherCore env f.constName!
 
+/--
+Auxiliary function for implementing `ctx.config.ground`: evaluate ground terms eagerly.
+We currently use `whnf` to implement this feature, but we want to stop ground evaluation at symbols marked with the `-` modifier.
+For example, `simp (config := { ground := true }) [-f]` should not unfold `f` even if the goal contains a ground term such as `f 2`.
+-/
+private def canUnfoldAtSimpGround (erased : SimpTheoremsArray) (_ : Meta.Config) (info : ConstantInfo) : CoreM Bool := do
+  return !erased.isErased (.decl info.name)
+
+/--
+Try to unfold `e`.
+-/
 private def unfold? (e : Expr) : SimpM (Option Expr) := do
   let f := e.getAppFn
   if !f.isConst then
     return none
   let fName := f.constName!
-  if (← isProjectionFn fName) then
-    return none -- should be reduced by `reduceProjFn?`
   let ctx ← read
-  if ctx.config.autoUnfold then
+  -- TODO: remove `rec` after we switch to new code generator
+  let rec unfoldGround? : SimpM (Option Expr) := do
+      unless ctx.config.ground do return none
+      -- We are assuming that assigned metavariables are going to be instantiated by the main simp loop.
+      if e.hasExprMVar || e.hasFVar then return none
+      if ctx.simpTheorems.isErased (.decl fName) then return none
+      -- TODO: check whether we need more filters
+      if (← isType e) then return none -- we don't unfold types
+      if (← isProof e) then return none -- we don't unfold proofs
+      if (← isInstance fName) then return none -- we don't unfold instances
+      -- TODO: we must have a notion of `simp` value, or more general solution for Lean
+      if Meta.isMatchValue e || isOfNatNatLit e then return none
+      if e.isConst then
+        -- We don't unfold constants that take arguments
+        -- TODO: add support for skipping partial applications too.
+        if let .forallE .. ← whnfD (← inferType e) then
+          return none
+      /-
+      We are currently using `whnf` with a custom `canUnfold?` predicate to reduce ground terms.
+      This can be inefficient, and produce proofs that are too expensive to type check in the Kernel. Some reasons:
+      - Functions defined by Well-founded recursion are expensive to reduce here and in the kernel.
+      - The kernel does not know we may have controlled reduction using `canUnfold?`.
+
+      It would be great to reduce the ground term using a to-be-implemented `cbv` tactic which produces a
+      proof that can be efficiently checked by the kernel.
+      -/
+      let eNew ← withDefault <|
+        withTheReader Meta.Context (fun c => { c with canUnfold? := canUnfoldAtSimpGround ctx.simpTheorems }) <| whnf e
+      if eNew == e then return none
+      trace[Meta.Tactic.simp.ground] "{e}\n---->\n{eNew}"
+      return some eNew
+  if let some eNew ← unfoldGround? then
+    return some eNew
+  else if (← isProjectionFn fName) then
+    return none -- should be reduced by `reduceProjFn?`
+  else if ctx.config.autoUnfold then
     if ctx.simpTheorems.isErased (.decl fName) then
       return none
     else if hasSmartUnfoldingDecl (← getEnv) fName then