Compare commits

...

4 Commits

Author SHA1 Message Date
Leonardo de Moura
4da34f6f98 feat: insertPattern 2026-01-03 11:13:02 -08:00
Leonardo de Moura
8d942ee415 refactor: move annotation support 2026-01-03 10:04:09 -08:00
Leonardo de Moura
be10299bfc refactor: move DiscrTree.insertCore and rename to DiscrTree.insertKeyValue 2026-01-03 09:44:09 -08:00
Leonardo de Moura
cbd8035200 chore: missing 2026-01-03 09:33:41 -08:00
15 changed files with 282 additions and 75 deletions

View File

@@ -193,7 +193,7 @@ def mkSpecTheoremFromStx (ref : Syntax) (proof : Expr) (prio : Nat := eval_prio
mkSpecTheorem type (.stx ( mkFreshId) ref proof) prio
def addSpecTheoremEntry (d : SpecTheorems) (e : SpecTheorem) : SpecTheorems :=
{ d with specs := d.specs.insertCore e.keys e }
{ d with specs := d.specs.insertKeyValue e.keys e }
def addSpecTheorem (ext : SpecExtension) (declName : Name) (prio : Nat) (attrKind : AttributeKind) : MetaM Unit := do
let thm mkSpecTheoremFromConst declName prio

View File

@@ -29,7 +29,7 @@ partial def headBetaUnderLambda (f : Expr) : Expr := Id.run do
builtin_initialize monotoneExt :
SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name)
registerSimpleScopedEnvExtension {
addEntry := fun dt (n, ks) => dt.insertCore ks n
addEntry := fun dt (n, ks) => dt.insertKeyValue ks n
initial := {}
}

View File

@@ -11,6 +11,12 @@ public import Lean.CoreM
public section
namespace Lean.Meta.DiscrTree
def mkNoindexAnnotation (e : Expr) : Expr :=
mkAnnotation `noindex e
def hasNoindexAnnotation (e : Expr) : Bool :=
annotation? `noindex e |>.isSome
instance : Inhabited (Trie α) where
default := .node #[] #[]
@@ -115,4 +121,60 @@ where
r := r.push ( go)
return r
private partial def createNodes (keys : Array Key) (v : α) (i : Nat) : Trie α :=
if h : i < keys.size then
let k := keys[i]
let c := createNodes keys v (i+1)
.node #[] #[(k, c)]
else
.node #[v] #[]
/--
If `vs` contains an element `v'` such that `v == v'`, then replace `v'` with `v`.
Otherwise, push `v`.
See issue #2155
Recall that `BEq α` may not be Lawful.
-/
private def insertVal [BEq α] (vs : Array α) (v : α) : Array α :=
loop 0
where
loop (i : Nat) : Array α :=
if h : i < vs.size then
if v == vs[i] then
vs.set i v
else
loop (i+1)
else
vs.push v
termination_by vs.size - i
private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat Trie α Trie α
| i, .node vs cs =>
if h : i < keys.size then
let k := keys[i]
let c := Id.run $ cs.binInsertM
(fun a b => a.1 < b.1)
(fun _, s => let c := insertAux keys v (i+1) s; (k, c)) -- merge with existing
(fun _ => let c := createNodes keys v (i+1); (k, c))
(k, default)
.node vs c
else
.node (insertVal vs v) cs
def insertKeyValue [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
if keys.isEmpty then panic! "invalid key sequence"
else
let k := keys[0]!
match d.root.find? k with
| none =>
let c := createNodes keys v 1
{ root := d.root.insert k c }
| some c =>
let c := insertAux keys v 1 c
{ root := d.root.insert k c }
@[deprecated insertKeyValue (since := "2026-01-02")]
def insertCore [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
insertKeyValue d keys v
end Lean.Meta.DiscrTree

View File

@@ -198,12 +198,6 @@ private def isOffset (fName : Name) (e : Expr) : MetaM Bool := do
private def shouldAddAsStar (fName : Name) (e : Expr) : MetaM Bool := do
isOffset fName e
def mkNoindexAnnotation (e : Expr) : Expr :=
mkAnnotation `noindex e
def hasNoindexAnnotation (e : Expr) : Bool :=
annotation? `noindex e |>.isSome
/--
Reduction procedure for the discrimination tree indexing.
-/
@@ -338,61 +332,9 @@ def mkPath (e : Expr) (noIndexAtArgs := false) : MetaM (Array Key) := do
let keys : Array Key := .mkEmpty initCapacity
mkPathAux (root := true) (todo.push e) keys noIndexAtArgs
private partial def createNodes (keys : Array Key) (v : α) (i : Nat) : Trie α :=
if h : i < keys.size then
let k := keys[i]
let c := createNodes keys v (i+1)
.node #[] #[(k, c)]
else
.node #[v] #[]
/--
If `vs` contains an element `v'` such that `v == v'`, then replace `v'` with `v`.
Otherwise, push `v`.
See issue #2155
Recall that `BEq α` may not be Lawful.
-/
private def insertVal [BEq α] (vs : Array α) (v : α) : Array α :=
loop 0
where
loop (i : Nat) : Array α :=
if h : i < vs.size then
if v == vs[i] then
vs.set i v
else
loop (i+1)
else
vs.push v
termination_by vs.size - i
private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat Trie α Trie α
| i, .node vs cs =>
if h : i < keys.size then
let k := keys[i]
let c := Id.run $ cs.binInsertM
(fun a b => a.1 < b.1)
(fun _, s => let c := insertAux keys v (i+1) s; (k, c)) -- merge with existing
(fun _ => let c := createNodes keys v (i+1); (k, c))
(k, default)
.node vs c
else
.node (insertVal vs v) cs
def insertCore [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
if keys.isEmpty then panic! "invalid key sequence"
else
let k := keys[0]!
match d.root.find? k with
| none =>
let c := createNodes keys v 1
{ root := d.root.insert k c }
| some c =>
let c := insertAux keys v 1 c
{ root := d.root.insert k c }
def insert [BEq α] (d : DiscrTree α) (e : Expr) (v : α) (noIndexAtArgs := false) : MetaM (DiscrTree α) := do
let keys mkPath e noIndexAtArgs
return d.insertCore keys v
return d.insertKeyValue keys v
/--
Inserts a value into a discrimination tree,
@@ -400,10 +342,10 @@ but only if its key is not of the form `#[*]` or `#[=, *, *, *]`.
-/
def insertIfSpecific [BEq α] (d : DiscrTree α) (e : Expr) (v : α) (noIndexAtArgs := false) : MetaM (DiscrTree α) := do
let keys mkPath e noIndexAtArgs
return if keys == #[Key.star] || keys == #[Key.const `Eq 3, Key.star, Key.star, Key.star] then
d
if keys == #[Key.star] || keys == #[Key.const `Eq 3, Key.star, Key.star, Key.star] then
return d
else
d.insertCore keys v
return d.insertKeyValue keys v
private def getKeyArgs (e : Expr) (isMatch root : Bool) : MetaM (Key × Array Expr) := do
let e reduceDT e root

View File

@@ -5,7 +5,7 @@ Authors: Leonardo de Moura
-/
module
prelude
public import Lean.ToExpr
public import Lean.Expr
public section
namespace Lean.Meta
/-! See file `DiscrTree.lean` for the actual implementation and documentation. -/

View File

@@ -75,8 +75,8 @@ structure Instances where
def addInstanceEntry (d : Instances) (e : InstanceEntry) : Instances :=
match e.globalName? with
| some n => { d with discrTree := d.discrTree.insertCore e.keys e, instanceNames := d.instanceNames.insert n e, erased := d.erased.erase n }
| none => { d with discrTree := d.discrTree.insertCore e.keys e }
| some n => { d with discrTree := d.discrTree.insertKeyValue e.keys e, instanceNames := d.instanceNames.insert n e, erased := d.erased.erase n }
| none => { d with discrTree := d.discrTree.insertKeyValue e.keys e }
def Instances.eraseCore (d : Instances) (declName : Name) : Instances :=
{ d with erased := d.erased.insert declName, instanceNames := d.instanceNames.erase declName }

View File

@@ -26,6 +26,7 @@ public import Lean.Meta.Sym.EqTrans
public import Lean.Meta.Sym.Congr
public import Lean.Meta.Sym.SimpResult
public import Lean.Meta.Sym.Simp
public import Lean.Meta.Sym.DiscrTree
/-!
# Symbolic simulation support.

View File

@@ -0,0 +1,139 @@
/-
Copyright (c) 2026 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
-/
module
prelude
public import Lean.Meta.Sym.Pattern
import Lean.Meta.DiscrTree.Basic
namespace Lean.Meta.Sym
open DiscrTree
/-!
# Discrimination Tree for the Symbolic Simulation Framework
This module provides pattern insertion into discrimination trees for the symbolic simulation
framework. Unlike the standard `DiscrTree` insertion which works with arbitrary expressions,
this module converts `Pattern` values (see `Sym/Pattern.lean`) into discrimination tree keys.
`Pattern` values have been preprocessed and use de Bruijn variables for wildcards instead of
metavariables.
## Key Design Decisions
The conversion from `Pattern` to `Array Key` respects the structural matching semantics of `Sym`:
1. **Proof and instance arguments are wildcards.** When building keys for a function application,
arguments marked as proofs or instances in `ProofInstInfo` are replaced with `Key.star`.
This ensures that during retrieval, terms differing only in proof/instance arguments will
match the same entries.
2. **Bound variables are wildcards.** Pattern variables (represented as de Bruijn indices in
`Pattern.pattern`) become `Key.star`, allowing them to match any subterm during retrieval.
3. **`noindex` annotations are respected.** Subterms annotated with `noindex` become `Key.star`.
## Usage
Use `insertPattern` to add a pattern to a discrimination tree:
```
let tree := insertPattern tree pattern value
```
Retrieval should use the standard `DiscrTree.getMatch` or similar, which will find all patterns
whose key sequence is compatible with the query term.
-/
/--
Returns the number of child keys for a given discrimination tree key.
**Note**: Unlike the standard `DiscrTree` module, `Key.arrow` has arity 2.
-/
def getKeyArity : Key Nat
| .const _ a => a
| .fvar _ a => a
| .arrow => 2
| _ => 0
/-- Returns `true` if argument at position `i` should be ignored (is a proof or instance). -/
def ignoreArg (infos : Array ProofInstArgInfo) (i : Nat) : Bool :=
if h : i < infos.size then
let info := infos[i]
info.isInstance || info.isProof
else
false
/-- Pushes all arguments of an application onto the todo stack in reverse order. -/
def pushAllArgs (e : Expr) (todo : Array Expr) : Array Expr :=
match e with
| .app f a => pushAllArgs f (todo.push a)
| _ => todo
/--
Pushes arguments of an application onto the todo stack, replacing proof/instance arguments
with dummy bound variables (which become `Key.star` wildcards).
-/
def pushArgsUsingInfo (infos : Array ProofInstArgInfo) (i : Nat) (e : Expr) (todo : Array Expr) : Array Expr :=
match e with
| .app f a =>
if ignoreArg infos i then
-- **Note**: We use a dummy bvar, it will be transformed into a `Key.star` later.
let dummyBVar := .bvar 1000000
pushArgsUsingInfo infos (i-1) f (todo.push dummyBVar)
else
pushArgsUsingInfo infos (i-1) f (todo.push a)
| _ => todo
/--
Computes the discrimination tree key for an expression and pushes its subterms onto the todo stack.
Returns `Key.star` for bound variables and `noindex`-annotated terms.
-/
def pushArgs (fnInfos : AssocList Name ProofInstInfo) (todo : Array Expr) (e : Expr) : Key × Array Expr :=
if hasNoindexAnnotation e then
(.star, todo)
else
let fn := e.getAppFn
match fn with
| .lit v => (.lit v, todo)
| .bvar _ => (.star, todo)
| .forallE _ d b _ => (.arrow, todo.push b |>.push d)
| .const declName _ =>
let numArgs := e.getAppNumArgs
let todo := if let some info := fnInfos.find? declName then
pushArgsUsingInfo info.argsInfo (numArgs - 1) e todo
else
pushAllArgs e todo
(.const declName numArgs, todo)
| .fvar fvarId =>
let numArgs := e.getAppNumArgs
let todo := pushAllArgs e todo
(.fvar fvarId numArgs, todo)
| _ => (.other, todo)
/-- Work-list based traversal that builds the key sequence for a pattern. -/
partial def mkPathAux (fnInfos : AssocList Name ProofInstInfo) (todo : Array Expr) (keys : Array Key) : Array Key :=
if todo.isEmpty then
keys
else
let e := todo.back!
let todo := todo.pop
let (k, todo) := pushArgs fnInfos todo e
mkPathAux fnInfos todo (keys.push k)
def initCapacity := 8
/-- Converts a `Pattern` into a discrimination tree key sequence. -/
public def Pattern.mkDiscrTreeKeys (p : Pattern) : Array Key :=
let todo : Array Expr := .mkEmpty initCapacity
let keys : Array Key := .mkEmpty initCapacity
mkPathAux p.fnInfos (todo.push p.pattern) keys
/-- Inserts a pattern into a discrimination tree, associating it with value `v`. -/
public def insertPattern [BEq α] (d : DiscrTree α) (p : Pattern) (v : α) : DiscrTree α :=
let keys := p.mkDiscrTreeKeys
d.insertKeyValue keys v
/-!
**TODO** Retrieval.
-/
end Lean.Meta.Sym

View File

@@ -40,7 +40,7 @@ builtin_initialize extExtension :
SimpleScopedEnvExtension ExtTheorem ExtTheorems
registerSimpleScopedEnvExtension {
addEntry := fun { tree, erased } thm =>
{ tree := tree.insertCore thm.keys thm, erased := erased.erase thm.declName }
{ tree := tree.insertKeyValue thm.keys thm, erased := erased.erase thm.declName }
initial := {}
}

View File

@@ -26,7 +26,7 @@ open Lean Meta
initialize reflExt :
SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name)
registerSimpleScopedEnvExtension {
addEntry := fun dt (n, ks) => dt.insertCore ks n
addEntry := fun dt (n, ks) => dt.insertKeyValue ks n
initial := {}
}

View File

@@ -572,9 +572,9 @@ def SimpTheorems.addSimpTheorem (d : SimpTheorems) (e : SimpTheorem) : SimpTheor
-- Erase the converse, if it exists
let d := eraseFwdIfBwd d e
if e.post then
{ d with post := d.post.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
{ d with post := d.post.insertKeyValue e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
else
{ d with pre := d.pre.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
{ d with pre := d.pre.insertKeyValue e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
where
updateLemmaNames (s : PHashSet Origin) : PHashSet Origin :=
s.insert e.origin

View File

@@ -153,9 +153,9 @@ def addSimprocAttrCore (ext : SimprocExtension) (declName : Name) (kind : Attrib
def Simprocs.addCore (s : Simprocs) (keys : Array SimpTheoremKey) (declName : Name) (post : Bool) (proc : Sum Simproc DSimproc) : Simprocs :=
let s := { s with simprocNames := s.simprocNames.insert declName, erased := s.erased.erase declName }
if post then
{ s with post := s.post.insertCore keys { declName, keys, post, proc } }
{ s with post := s.post.insertKeyValue keys { declName, keys, post, proc } }
else
{ s with pre := s.pre.insertCore keys { declName, keys, post, proc } }
{ s with pre := s.pre.insertKeyValue keys { declName, keys, post, proc } }
/--
Implements attributes `builtin_simproc` and `builtin_sevalproc`.

View File

@@ -26,7 +26,7 @@ namespace Lean.Meta.Symm
builtin_initialize symmExt :
SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name)
registerSimpleScopedEnvExtension {
addEntry := fun dt (n, ks) => dt.insertCore ks n
addEntry := fun dt (n, ks) => dt.insertKeyValue ks n
initial := {}
}

View File

@@ -32,7 +32,7 @@ private def config : ConfigWithKey :=
{ iota := false, proj := .no : Config }.toConfigWithKey
def UnificationHints.add (hints : UnificationHints) (e : UnificationHintEntry) : UnificationHints :=
{ hints with discrTree := hints.discrTree.insertCore e.keys e.val }
{ hints with discrTree := hints.discrTree.insertKeyValue e.keys e.val }
builtin_initialize unificationHintExtension : SimpleScopedEnvExtension UnificationHintEntry UnificationHints
registerSimpleScopedEnvExtension {

View File

@@ -1,4 +1,5 @@
import Lean.Meta.Sym
import Lean.Meta.DiscrTree.Basic
open Lean Meta Sym Grind
set_option grind.debug true
opaque p [Ring α] : α α Prop
@@ -58,3 +59,65 @@ info: ∀ (x : Nat), q (f (f x)) (f x + f (f 1))
-/
#guard_msgs in
#eval SymM.run' test₂
theorem forall_and_eq (P Q : α Prop) : ( x, P x Q x) = (( x, P x) ( x, Q x)):= by
grind
def logPatternKey (p : Pattern) : MetaM Unit := do
let k := p.mkDiscrTreeKeys
logInfo m!"{k.toList.map (·.format)}"
def logPatternKeyFor (declName : Name) : MetaM Unit := do
let (p, _) mkEqPatternFromDecl declName
logPatternKey p
/--
info: [HAdd.hAdd, Nat, Nat, Nat, *, OfNat.ofNat, Nat, 0, *, *]
---
info: [HMul.hMul, *, *, *, *, OfNat.ofNat, *, 0, *, *]
---
info: [∀, *, And, *, *]
---
info: [Array.eraseIdx, *, HAppend.hAppend, Array, *, Array, *, Array, *, *, *, *, *, *]
---
info: [List.map, *, *, *, List.map, *, *, *, *]
---
info: [Std.HashMap.insertMany,
*,
*,
*,
*,
List,
Prod,
*,
*,
*,
*,
HAppend.hAppend,
List,
Prod,
*,
*,
List,
Prod,
*,
*,
List,
Prod,
*,
*,
*,
*,
*]
---
info: [GetElem.getElem, Std.HashMap, *, *, *, *, *, *, ◾, *, Std.HashMap.insert, *, *, *, *, *, *, *, *, *]
-/
#guard_msgs in
#eval SymM.run' do
logPatternKeyFor ``Nat.zero_add
logPatternKeyFor ``Grind.Semiring.zero_mul
logPatternKeyFor ``forall_and_eq
logPatternKeyFor ``Array.eraseIdx_append
logPatternKeyFor ``List.map_map
logPatternKeyFor ``Std.HashMap.insertMany_append
logPatternKeyFor ``Std.HashMap.getElem_insert