feat: support for ground terms at DiscrTree.lean

WIP for discussing with @joehendrix

src/Lean/Meta/DiscrTree.lean

@@ -48,14 +48,19 @@ namespace Lean.Meta.DiscrTree
   2- Distinguish partial applications `f a`, `f a b`, and `f a b c`.
 -/
 
+/--
+Auxiliary function used to implement `Key.lt`.
+Note that the `DiscrTree` implementation assumes that the `.star` key is the smallest one, and `.ground` key is the biggest.
+-/
 def Key.ctorIdx : Key s → Nat
-  | .star     => 0
-  | .other    => 1
-  | .lit ..   => 2
-  | .fvar ..  => 3
-  | .const .. => 4
-  | .arrow    => 5
-  | .proj ..  => 6
+  | .star      => 0
+  | .other     => 1
+  | .lit ..    => 2
+  | .fvar ..   => 3
+  | .const ..  => 4
+  | .arrow     => 5
+  | .proj ..   => 6
+  | .ground .. => 7
 
 def Key.lt : Key s → Key s → Bool
   | .lit v₁,        .lit v₂        => v₁ < v₂
@@ -69,6 +74,7 @@ instance (a b : Key s) : Decidable (a < b) := inferInstanceAs (Decidable (Key.lt
 
 def Key.format : Key s → Format
   | .star                   => "*"
+  | .ground                 => "*g"
   | .other                  => "◾"
   | .lit (Literal.natVal v) => Std.format v
   | .lit (Literal.strVal v) => repr v
@@ -247,6 +253,12 @@ def mkNoindexAnnotation (e : Expr) : Expr :=
 def hasNoindexAnnotation (e : Expr) : Bool :=
   annotation? `noindex e |>.isSome
 
+def mkGroundAnnotation (e : Expr) : Expr :=
+  mkAnnotation `ground e
+
+def hasGroundAnnotation (e : Expr) : Bool :=
+  annotation? `ground e |>.isSome
+
 /--
 Reduction procedure for the discrimination tree indexing.
 The parameter `simpleReduce` controls how aggressive the term is reduced.
@@ -328,6 +340,10 @@ def reduceDT (e : Expr) (root : Bool) (simpleReduce : Bool) : MetaM Expr :=
 private def pushArgs (root : Bool) (todo : Array Expr) (e : Expr) : MetaM (Key s × Array Expr) := do
   if hasNoindexAnnotation e then
     return (.star, todo)
+  else if hasGroundAnnotation e then
+    if root then
+      throwError "`ground` pattern modifier cannot be use in root terms{indentExpr e}"
+    return (.ground, todo)
   else
     let e ← reduceDT e root (simpleReduce := s)
     let fn := e.getAppFn
@@ -539,7 +555,7 @@ private abbrev getUnifyKeyArgs (e : Expr) (root : Bool) : MetaM (Key s × Array
 private def getStarResult (d : DiscrTree α s) : Array α :=
   let result : Array α := .mkEmpty initCapacity
   match d.root.find? .star with
-  | none                  => result
+  | none              => result
   | some (.node vs _) => result ++ vs
 
 private abbrev findKey (cs : Array (Key s × Trie α s)) (k : Key s) : Option (Key s × Trie α s) :=
@@ -555,21 +571,32 @@ private partial def getMatchLoop (todo : Array Expr) (c : Trie α s) (result : A
     else
       let e     := todo.back
       let todo  := todo.pop
-      let first := cs[0]! /- Recall that `Key.star` is the minimal key -/
       let (k, args) ← getMatchKeyArgs e (root := false)
       /- We must always visit `Key.star` edges since they are wildcards.
          Thus, `todo` is not used linearly when there is `Key.star` edge
          and there is an edge for `k` and `k != Key.star`. -/
       let visitStar (result : Array α) : MetaM (Array α) :=
+        let first := cs[0]! /- Recall that `Key.star` is the minimal key -/
         if first.1 == .star then
           getMatchLoop todo first.2 result
         else
           return result
+      let visitGround (result : Array α) : MetaM (Array α) := do
+        let last  := cs.back /- Recall that `Key.ground` is the maximal key -/
+        if last.1 == .ground then
+          let e ← instantiateMVars e /- TODO(Leo): check whether this is perf bottleneck. -/
+          if e.hasFVar || e.hasExprMVar then
+            return result
+          else
+            getMatchLoop todo last.2 result
+        else
+          return result
       let visitNonStar (k : Key s) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
         match findKey cs k with
         | none   => return result
         | some c => getMatchLoop (todo ++ args) c.2 result
       let result ← visitStar result
+      let result ← visitGround result
       match k with
       | .star  => return result
       /-
@@ -666,6 +693,18 @@ where
             process 0 todo first.2 result
           else
             return result
+        let visitGround (result : Array α) : MetaM (Array α) := do
+          let last := cs.back
+          if last.1 == .ground then
+            let e ← instantiateMVars e /- TODO(Leo): check whether this is perf bottleneck. -/
+            if e.hasFVar || e.hasExprMVar then
+              return result
+            else
+              process 0 todo last.2 result
+          else
+            return result
+        let visitStarAndGround (result : Array α) : MetaM (Array α) := do
+           visitGround (← visitStar result)
         let visitNonStar (k : Key s) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
           match findKey cs k with
           | none   => return result
@@ -673,7 +712,7 @@ where
         match k with
         | .star  => cs.foldlM (init := result) fun result ⟨k, c⟩ => process k.arity todo c result
         -- See comment a `getMatch` regarding non-dependent arrows vs dependent arrows
-        | .arrow => visitNonStar .other #[] (← visitNonStar k args (← visitStar result))
-        | _      => visitNonStar k args (← visitStar result)
+        | .arrow => visitNonStar .other #[] (← visitNonStar k args (← visitStarAndGround result))
+        | _      => visitNonStar k args (← visitStarAndGround result)
 
 end Lean.Meta.DiscrTree

src/Lean/Meta/DiscrTreeTypes.lean

@@ -11,16 +11,42 @@ namespace Lean.Meta
 
 namespace DiscrTree
 /--
-Discrimination tree key. See `DiscrTree`
+(Imperfect) Discrimination tree key. See `DiscrTree`.
+The parameter `simpleReduce` controls how aggressive the term is reduced.
 -/
 inductive Key (simpleReduce : Bool) where
-  | const : Name → Nat → Key simpleReduce
-  | fvar  : FVarId → Nat → Key simpleReduce
-  | lit   : Literal → Key simpleReduce
-  | star  : Key simpleReduce
-  | other : Key simpleReduce
-  | arrow : Key simpleReduce
-  | proj  : Name → Nat → Nat → Key simpleReduce
+  /-- A constant application with name `declName` and `arity` arguments. -/
+  | const (declName : Name) (arity : Nat)
+  /--
+  Free variables (and arity). Thus, an entry in the discrimination tree
+  may reference hypotheses from the local context.
+  -/
+  | fvar (fvarId : FVarId) (arity : Nat)
+  /-- Literal. -/
+  | lit (litVal : Literal)
+  /--
+  Star or wildcard. We use them to represent metavariables and terms
+  we want to ignore. We ignore implicit arguments and proofs.
+  -/
+  | star
+  /--
+  Other terms. We use to represent other kinds of terms
+  (e.g., nested lambda, forall, sort, etc).
+  -/
+  | other
+  /-- Arrow (aka non dependent arrows). -/
+  | arrow
+  /-- Projection (applications). -/
+  | proj (structName : Name) (projIdx : Nat) (arity : Nat)
+  /--
+  Ground terms.
+  We use to implement the `ground p` pattern annotation.
+  When the `DiscrTree` is trying to match a term `e` with key the `.ground`,
+  it succeeds if `e` does not contain free or meta variables.
+  Note that, in the pattern `ground p`, the term `p` is ignored.
+  We can also view `ground` as a variant of `star` that matches all ground terms.
+  -/
+  | ground : Key simpleReduce
   deriving Inhabited, BEq, Repr
 
 protected def Key.hash : Key s → UInt64
@@ -30,6 +56,7 @@ protected def Key.hash : Key s → UInt64
   | Key.star        => 7883
   | Key.other       => 2411
   | Key.arrow       => 17
+  | Key.ground      => 11
   | Key.proj s i a  =>  mixHash (hash a) $ mixHash (hash s) (hash i)
 
 instance : Hashable (Key s) := ⟨Key.hash⟩

tests/lean/run/ground.lean

@@ -0,0 +1,22 @@
+import Lean
+
+-- TODO: move elaborator to core
+open Lean Meta Elab Term
+elab "ground " e:term : term <= expectedType =>
+  return DiscrTree.mkGroundAnnotation (← elabTerm e expectedType)
+
+def consn (n : Nat) (a : α) (as : List α) : List α :=
+  match n with
+  | 0 => as
+  | n+1 => consn n a (a :: as)
+
+@[simp] theorem consn_0 : consn 0 a as = as := rfl
+@[simp] theorem consn_succ : consn (ground n+1) a as = consn n a (a::as) := rfl
+
+example : consn (n+2) a as = consn n a (a::a::as) := by
+  fail_if_success simp -- Should not fire consn_succ
+  simp [consn]
+  done
+
+example : consn 2 a as = a::a::as := by
+  simp -- should fire consn