-
Notifications
You must be signed in to change notification settings - Fork 337
/
Canonicalizer.lean
149 lines (132 loc) · 6.39 KB
/
Canonicalizer.lean
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
/-
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 Lean.Util.ShareCommon
import Lean.Data.HashMap
import Lean.Meta.Basic
import Lean.Meta.FunInfo
namespace Lean.Meta
namespace Canonicalizer
/-!
Applications have implicit arguments. Thus, two terms that may look identical when pretty-printed can be structurally different.
For example, `@id (Id Nat) x` and `@id Nat x` are structurally different but are both pretty-printed as `id x`.
Moreover, these two terms are definitionally equal since `Id Nat` reduces to `Nat`. This may create situations
that are counterintuitive to our users. Furthermore, several tactics (e.g., `omega`) need to collect unique atoms in a goal.
One simple approach is to maintain a list of atoms found so far, and whenever a new atom is discovered, perform a
linear scan to test whether it is definitionally equal to a previously found one. However, this method is too costly,
even if the definitional equality test were inexpensive.
This module aims to efficiently identify terms that are structurally different, definitionally equal, and structurally equal
when we disregard implicit arguments like `@id (Id Nat) x` and `@id Nat x`. The procedure is straightforward. For each atom,
we create a new abstracted atom by erasing all implicit information. We refer to this abstracted atom as a 'key.' For the two
terms mentioned, the key would be `@id _ x`, where `_` denotes a placeholder for a dummy term. To preserve any
pre-existing directed acyclic graph (DAG) structure and prevent exponential blowups while constructing the key, we employ
unsafe techniques, such as pointer equality. Additionally, we maintain a mapping from keys to lists of terms, where each
list contains terms sharing the same key but not definitionally equal. We posit that these lists will be small in practice.
-/
/--
Auxiliary structure for creating a pointer-equality mapping from `Expr` to `Key`.
We use this mapping to ensure we preserve the dag-structure of input expressions.
-/
structure ExprVisited where
e : Expr
deriving Inhabited
unsafe instance : BEq ExprVisited where
beq a b := ptrAddrUnsafe a == ptrAddrUnsafe b
unsafe instance : Hashable ExprVisited where
hash a := USize.toUInt64 (ptrAddrUnsafe a)
abbrev Key := ExprVisited
/--
State for the `CanonM` monad.
-/
structure State where
/-- "Set" of all keys created so far. This is a hash-consing helper structure available in Lean. -/
keys : ShareCommon.State.{0} Lean.ShareCommon.objectFactory := ShareCommon.State.mk Lean.ShareCommon.objectFactory
/-- Mapping from `Expr` to `Key`. See comment at `ExprVisited`. -/
-- We use `HashMapImp` to ensure we don't have to tag `State` as `unsafe`.
cache : HashMapImp ExprVisited Key := mkHashMapImp
/--
Given a key `k` and `keyToExprs.find? k = some es`, we have that all `es` share key `k`, and
are not definitionally equal modulo the transparency setting used. -/
keyToExprs : HashMapImp Key (List Expr) := mkHashMapImp
instance : Inhabited State where
default := {}
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) {}
private def shareCommon (a : α) : CanonM α :=
modifyGet fun { keys, cache, keyToExprs } =>
let (a, keys) := ShareCommon.State.shareCommon keys a
(a, { keys, cache, keyToExprs })
private partial def mkKey (e : Expr) : CanonM Key := do
if let some key := unsafe (← get).cache.find? { e } then
return key
else
let key ← match e with
| .sort .. | .fvar .. | .bvar .. | .lit .. =>
pure { e := (← shareCommon e) }
| .const n _ =>
pure { e := (← shareCommon (.const n [])) }
| .mvar .. =>
-- We instantiate assigned metavariables because the
-- pretty-printer also instantiates them.
let eNew ← instantiateMVars e
if eNew == e then pure { e := (← shareCommon e) }
else mkKey eNew
| .mdata _ a => mkKey a
| .app .. =>
let f := e.getAppFn
if f.isMVar then
let eNew ← instantiateMVars e
unless eNew == e do
return (← mkKey eNew)
let info ← getFunInfo f
let args ← e.getAppArgs.mapIdxM fun i arg => do
if h : i < info.paramInfo.size then
let info := info.paramInfo[i]
if info.isExplicit && !info.isProp then
pure (← mkKey arg).e
else
pure (mkSort 0) -- some dummy value for erasing implicit
else
pure (← mkKey arg).e
let f' := (← mkKey f).e
pure { e := (← shareCommon (mkAppN f' args)) }
| .lam n t b i =>
pure { e := (← shareCommon (.lam n (← mkKey t).e (← mkKey b).e i)) }
| .forallE n t b i =>
pure { e := (← shareCommon (.forallE n (← mkKey t).e (← mkKey b).e i)) }
| .letE n t v b d =>
pure { e := (← shareCommon (.letE n (← mkKey t).e (← mkKey v).e (← mkKey b).e d)) }
| .proj t i s =>
pure { e := (← shareCommon (.proj t i (← mkKey s).e)) }
unsafe modify fun { keys, cache, keyToExprs} => { keys, keyToExprs, cache := cache.insert { e } key |>.1 }
return key
/--
"Canonicalize" the given expression.
-/
def canon (e : Expr) : CanonM Expr := do
let k ← mkKey e
-- Find all expressions canonicalized before that have the same key.
if let some es' := unsafe (← get).keyToExprs.find? k then
withTransparency (← read) do
for e' in es' do
-- Found an expression `e'` that is definitionally equal to `e` and share the same key.
if (← isDefEq e e') then
return e'
-- `e` is not definitionally equal to any expression in `es'`. We claim this should be rare.
unsafe modify fun { keys, cache, keyToExprs } => { keys, cache, keyToExprs := keyToExprs.insert k (e :: es') |>.1 }
return e
else
-- `e` is the first expression we found with key `k`.
unsafe modify fun { keys, cache, keyToExprs } => { keys, cache, keyToExprs := keyToExprs.insert k [e] |>.1 }
return e
end Canonicalizer
export Canonicalizer (CanonM canon)
end Lean.Meta