-
Notifications
You must be signed in to change notification settings - Fork 307
/
DeriveToExpr.lean
229 lines (207 loc) · 10.5 KB
/
DeriveToExpr.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
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
/-
Copyright (c) 2023 Kyle Miller. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Kyle Miller
-/
import Mathlib.Tactic.ToLevel
/-!
# A `ToExpr` derive handler
This module defines a `ToExpr` derive handler for inductive types. It supports mutually inductive
types as well.
The `ToExpr` derive handlers support universe level polymorphism. This is implemented using the
`Lean.ToLevel` class. To use `ToExpr` in places where there is universe polymorphism, make sure
to have a `[ToLevel.{u}]` instance available.
**Warning:** Import `Mathlib.Tactic.ToExpr` instead of this one. This ensures that you are using
the universe polymorphic `ToExpr` instances that override the ones from Lean 4 core.
Implementation note: this derive handler was originally modeled after the `Repr` derive handler.
-/
namespace Mathlib.Deriving.ToExpr
open Lean Elab Lean.Parser.Term
open Meta Command Deriving
/-- Specialization of `Lean.Elab.Deriving.mkHeader` for `ToExpr`. -/
def mkToExprHeader (indVal : InductiveVal) : TermElabM Header := do
-- The auxiliary functions we produce are `indtype -> Expr`.
let header ← mkHeader ``ToExpr 1 indVal
return header
/-- Give a term that is equivalent to `(term| mkAppN $f #[$args,*])`.
As an optimization, `mkAppN` is pre-expanded out to use `Expr.app` directly. -/
def mkAppNTerm (f : Term) (args : Array Term) : MetaM Term :=
args.foldlM (fun a b => `(Expr.app $a $b)) f
/-- Create the body of the `toExpr` function
for the `ToExpr` instance, which is a `match` expression
that calls `toExpr` and `toTypeExpr` to assemble an expression for a given term.
For recursive inductive types, `auxFunName` refers to the `ToExpr` instance
for the current type.
For mutually recursive types, we rely on the local instances set up by `mkLocalInstanceLetDecls`. -/
def mkToExprBody (header : Header) (indVal : InductiveVal) (auxFunName : Name) :
TermElabM Term := do
let discrs ← mkDiscrs header indVal
let alts ← mkAlts
`(match $[$discrs],* with $alts:matchAlt*)
where
/-- Create the `match` cases, one per constructor. -/
mkAlts : TermElabM (Array (TSyntax ``matchAlt)) := do
let mut alts := #[]
for ctorName in indVal.ctors do
let ctorInfo ← getConstInfoCtor ctorName
let alt ← forallTelescopeReducing ctorInfo.type fun xs _ => do
let mut patterns := #[]
-- add `_` pattern for indices
for _ in [:indVal.numIndices] do
patterns := patterns.push (← `(_))
let mut ctorArgs := #[]
let mut rhsArgs : Array Term := #[]
let mkArg (x : Expr) (a : Term) : TermElabM Term := do
if (← inferType x).isAppOf indVal.name then
`($(mkIdent auxFunName) $a)
else if ← Meta.isType x then
`(toTypeExpr $a)
else
`(toExpr $a)
-- add `_` pattern for inductive parameters, which are inaccessible
for i in [:ctorInfo.numParams] do
let a := mkIdent header.argNames[i]!
ctorArgs := ctorArgs.push (← `(_))
rhsArgs := rhsArgs.push <| ← mkArg xs[i]! a
for i in [:ctorInfo.numFields] do
let a := mkIdent (← mkFreshUserName `a)
ctorArgs := ctorArgs.push a
rhsArgs := rhsArgs.push <| ← mkArg xs[ctorInfo.numParams + i]! a
patterns := patterns.push (← `(@$(mkIdent ctorName):ident $ctorArgs:term*))
let levels ← indVal.levelParams.toArray.mapM (fun u => `(toLevel.{$(mkIdent u)}))
let rhs : Term ←
mkAppNTerm (← `(Expr.const $(quote ctorInfo.name) [$levels,*])) rhsArgs
`(matchAltExpr| | $[$patterns:term],* => $rhs)
alts := alts.push alt
return alts
/-- Create the body of the `toTypeExpr` function for the `ToExpr` instance.
Calls `toExpr` and `toTypeExpr` to the arguments to the type constructor. -/
def mkToTypeExpr (argNames : Array Name) (indVal : InductiveVal) : TermElabM Term := do
let levels ← indVal.levelParams.toArray.mapM (fun u => `(toLevel.{$(mkIdent u)}))
forallTelescopeReducing indVal.type fun xs _ => do
let mut args : Array Term := #[]
for i in [:xs.size] do
let x := xs[i]!
let a := mkIdent argNames[i]!
if ← Meta.isType x then
args := args.push <| ← `(toTypeExpr $a)
else
args := args.push <| ← `(toExpr $a)
mkAppNTerm (← `((Expr.const $(quote indVal.name) [$levels,*]))) args
/--
For mutually recursive inductive types, the strategy is to have local `ToExpr` instances in scope
for each of the inductives when defining each instance.
This way, each instance can freely use `toExpr` and `toTypeExpr` for each of the other types.
Note that each instance gets its own definition of each of the others' `toTypeExpr` fields.
(This is working around the fact that the `Deriving.Context` API assumes
that each instance in mutual recursion only has a single auxiliary definition.
There are other ways to work around it, but `toTypeExpr` implementations
are very simple, so duplicating them seemed to be OK.) -/
def mkLocalInstanceLetDecls (ctx : Deriving.Context) (argNames : Array Name) :
TermElabM (Array (TSyntax ``Parser.Term.letDecl)) := do
let mut letDecls := #[]
for i in [:ctx.typeInfos.size] do
let indVal := ctx.typeInfos[i]!
let auxFunName := ctx.auxFunNames[i]!
let currArgNames ← mkInductArgNames indVal
let numParams := indVal.numParams
let currIndices := currArgNames[numParams:]
let binders ← mkImplicitBinders currIndices
let argNamesNew := argNames[:numParams] ++ currIndices
let indType ← mkInductiveApp indVal argNamesNew
let instName ← mkFreshUserName `localinst
let toTypeExpr ← mkToTypeExpr argNames indVal
let letDecl ← `(Parser.Term.letDecl| $(mkIdent instName):ident $binders:implicitBinder* :
ToExpr $indType :=
{ toExpr := $(mkIdent auxFunName), toTypeExpr := $toTypeExpr })
letDecls := letDecls.push letDecl
return letDecls
/-- Fix the output of `mkInductiveApp` to explicitly reference universe levels. -/
def fixIndType (indVal : InductiveVal) (t : Term) : TermElabM Term :=
match t with
| `(@$f $args*) =>
let levels := indVal.levelParams.toArray.map mkIdent
`(@$f.{$levels,*} $args*)
| _ => throwError "(internal error) expecting output of `mkInductiveApp`"
/-- Make `ToLevel` instance binders for all the level variables. -/
def mkToLevelBinders (indVal : InductiveVal) : TermElabM (TSyntaxArray ``instBinderF) := do
indVal.levelParams.toArray.mapM (fun u => `(instBinderF| [ToLevel.{$(mkIdent u)}]))
open TSyntax.Compat in
/-- Make a `toExpr` function for the given inductive type.
The implementations of each `toExpr` function for a (mutual) inductive type
are given as top-level private definitions.
These end up being assembled into `ToExpr` instances in `mkInstanceCmds`.
For mutual inductive types,
then each of the other types' `ToExpr` instances are provided as local instances,
to wire together the recursion (this necessitates these auxiliary definitions being `partial`). -/
def mkAuxFunction (ctx : Deriving.Context) (i : Nat) : TermElabM Command := do
let auxFunName := ctx.auxFunNames[i]!
let indVal := ctx.typeInfos[i]!
let header ← mkToExprHeader indVal
let mut body ← mkToExprBody header indVal auxFunName
if ctx.usePartial then
let letDecls ← mkLocalInstanceLetDecls ctx header.argNames
body ← mkLet letDecls body
-- We need to alter the last binder (the one for the "target") to have explicit universe levels
-- so that the `ToLevel` instance arguments can use them.
let addLevels binder :=
match binder with
| `(bracketedBinderF| ($a : $ty)) => do `(bracketedBinderF| ($a : $(← fixIndType indVal ty)))
| _ => throwError "(internal error) expecting inst binder"
let binders := header.binders.pop
++ (← mkToLevelBinders indVal)
++ #[← addLevels header.binders.back]
let levels := indVal.levelParams.toArray.map mkIdent
if ctx.usePartial then
`(private partial def $(mkIdent auxFunName):ident.{$levels,*} $binders:bracketedBinder* :
Expr := $body:term)
else
`(private def $(mkIdent auxFunName):ident.{$levels,*} $binders:bracketedBinder* :
Expr := $body:term)
/-- Create all the auxiliary functions using `mkAuxFunction` for the (mutual) inductive type(s).
Wraps the resulting definition commands in `mutual ... end`. -/
def mkMutualBlock (ctx : Deriving.Context) : TermElabM Syntax := do
let mut auxDefs := #[]
for i in [:ctx.typeInfos.size] do
auxDefs := auxDefs.push (← mkAuxFunction ctx i)
`(mutual $auxDefs:command* end)
open TSyntax.Compat in
/-- Assuming all of the auxiliary definitions exist, create all the `instance` commands
for the `ToExpr` instances for the (mutual) inductive type(s). -/
def mkInstanceCmds (ctx : Deriving.Context) (typeNames : Array Name) :
TermElabM (Array Command) := do
let mut instances := #[]
for i in [:ctx.typeInfos.size] do
let indVal := ctx.typeInfos[i]!
if typeNames.contains indVal.name then
let auxFunName := ctx.auxFunNames[i]!
let argNames ← mkInductArgNames indVal
let binders ← mkImplicitBinders argNames
let binders := binders ++ (← mkInstImplicitBinders ``ToExpr indVal argNames)
let binders := binders ++ (← mkToLevelBinders indVal)
let indType ← fixIndType indVal (← mkInductiveApp indVal argNames)
let toTypeExpr ← mkToTypeExpr argNames indVal
let levels := indVal.levelParams.toArray.map mkIdent
let instCmd ← `(instance $binders:implicitBinder* : ToExpr $indType where
toExpr := $(mkIdent auxFunName).{$levels,*}
toTypeExpr := $toTypeExpr)
instances := instances.push instCmd
return instances
/-- Returns all the commands generated by `mkMutualBlock` and `mkInstanceCmds`. -/
def mkToExprInstanceCmds (declNames : Array Name) : TermElabM (Array Syntax) := do
let ctx ← mkContext "toExpr" declNames[0]!
let cmds := #[← mkMutualBlock ctx] ++ (← mkInstanceCmds ctx declNames)
trace[Elab.Deriving.toExpr] "\n{cmds}"
return cmds
/-- The main entry point to the `ToExpr` derive handler. -/
def mkToExprInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
if (← declNames.allM isInductive) && declNames.size > 0 then
let cmds ← liftTermElabM <| mkToExprInstanceCmds declNames
cmds.forM elabCommand
return true
else
return false
initialize
registerDerivingHandler `Lean.ToExpr mkToExprInstanceHandler
registerTraceClass `Elab.Deriving.toExpr
end Mathlib.Deriving.ToExpr