Skip to content

Commit c01494c

Browse files
EdAyersEdAyers
committed
feat: more expr traversal functions (#263)
- split Mathlib/Lean/Expr into multiple files. - implement a version of replaceRecM that correctly instantiates free vars. - some additonal Expr functions These are prereqs of #234 and #229 Co-authored-by: Ed Ayers <EdAyers@users.noreply.github.com>
1 parent 45731d6 commit c01494c

File tree

5 files changed

+229
-60
lines changed

5 files changed

+229
-60
lines changed

‎Mathlib.lean‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ import Mathlib.Init.Logic
4343
import Mathlib.Init.Set
4444
import Mathlib.Init.SetNotation
4545
import Mathlib.Lean.Expr
46+
import Mathlib.Lean.Expr.Basic
4647
import Mathlib.Lean.Expr.ReplaceRec
4748
import Mathlib.Lean.Expr.Traverse
4849
import Mathlib.Lean.LocalContext

‎Mathlib/Lean/Expr.lean‎

Lines changed: 3 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -4,63 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
44
Authors: Mario Carneiro, Simon Hudon, Scott Morrison, Keeley Hoek, Robert Y. Lewis, Floris van Doorn
55
-/
66
import Lean
7-
/-!
8-
# Additional operations on Expr and related types
9-
10-
This file defines basic operations on the types expr, name, declaration, level, environment.
11-
12-
This file is mostly for non-tactics.
13-
-/
14-
15-
open Lean Meta
16-
17-
namespace Lean
18-
19-
namespace BinderInfo
20-
21-
/-! ### Declarations about `BinderInfo` -/
22-
23-
/-- The brackets corresponding to a given `BinderInfo`. -/
24-
def brackets : BinderInfo → String × String
25-
| BinderInfo.implicit => ("{", "}")
26-
| BinderInfo.strictImplicit => ("{{", "}}")
27-
| BinderInfo.instImplicit => ("[", "]")
28-
| _ => ("(", ")")
29-
30-
end BinderInfo
31-
32-
namespace Name
33-
34-
/-! ### Declarations about `name` -/
35-
36-
/-- Find the largest prefix `n` of a `Name` such that `f n != none`, then replace this prefix
37-
with the value of `f n`. -/
38-
def mapPrefix (f : Name → Option Name) (n : Name) : Name := Id.run do
39-
if let some n' := f n then return n'
40-
match n with
41-
| anonymous => anonymous
42-
| str n' s _ => mkStr (mapPrefix f n') s
43-
| num n' i _ => mkNum (mapPrefix f n') i
44-
45-
end Name
46-
47-
namespace Expr
48-
49-
/-! ### Declarations about `Expr` -/
50-
51-
/-- If the expression is a constant, return that name. Otherwise return `Name.anonymous`. -/
52-
def constName (e : Expr) : Name :=
53-
e.constName?.getD Name.anonymous
54-
55-
/-- Return the function (name) and arguments of an application. -/
56-
def getAppFnArgs (e : Expr) : Name × Array Expr :=
57-
withApp e λ e a => (e.constName, a)
58-
59-
/-- Turn an expression that is a natural number literal into a natural number. -/
60-
def natLit! : Expr → Nat
61-
| lit (Literal.natVal v) _ => v
62-
| _ => panic! "nat literal expected"
63-
64-
end Expr
65-
66-
end Lean
7+
import Mathlib.Lean.Expr.Basic
8+
import Mathlib.Lean.Expr.Traverse
9+
import Mathlib.Lean.Expr.ReplaceRec

‎Mathlib/Lean/Expr/Basic.lean‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
/-
2+
Copyright (c) 2019 Robert Y. Lewis. All rights reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: Mario Carneiro, Simon Hudon, Scott Morrison, Keeley Hoek, Robert Y. Lewis, Floris van Doorn, E.W.Ayers
5+
-/
6+
import Lean
7+
8+
/-!
9+
# Additional operations on Expr and related types
10+
11+
This file defines basic operations on the types expr, name, declaration, level, environment.
12+
13+
This file is mostly for non-tactics.
14+
-/
15+
16+
namespace Lean
17+
18+
namespace BinderInfo
19+
20+
/-! ### Declarations about `BinderInfo` -/
21+
22+
/-- The brackets corresponding to a given `BinderInfo`. -/
23+
def brackets : BinderInfo → String × String
24+
| BinderInfo.implicit => ("{", "}")
25+
| BinderInfo.strictImplicit => ("{{", "}}")
26+
| BinderInfo.instImplicit => ("[", "]")
27+
| _ => ("(", ")")
28+
29+
end BinderInfo
30+
31+
namespace Name
32+
33+
/-! ### Declarations about `name` -/
34+
35+
/-- Find the largest prefix `n` of a `Name` such that `f n != none`, then replace this prefix
36+
with the value of `f n`. -/
37+
def mapPrefix (f : Name → Option Name) (n : Name) : Name := Id.run do
38+
if let some n' := f n then return n'
39+
match n with
40+
| anonymous => anonymous
41+
| str n' s _ => mkStr (mapPrefix f n') s
42+
| num n' i _ => mkNum (mapPrefix f n') i
43+
44+
end Name
45+
46+
47+
namespace ConstantInfo
48+
49+
def isDef : ConstantInfo → Bool
50+
| defnInfo _ => true
51+
| _ => false
52+
53+
def isThm : ConstantInfo → Bool
54+
| thmInfo _ => true
55+
| _ => false
56+
57+
def updateName : ConstantInfo → Name → ConstantInfo
58+
| defnInfo info, n => defnInfo {info with name := n}
59+
| axiomInfo info, n => axiomInfo {info with name := n}
60+
| thmInfo info, n => thmInfo {info with name := n}
61+
| opaqueInfo info, n => opaqueInfo {info with name := n}
62+
| quotInfo info, n => quotInfo {info with name := n}
63+
| inductInfo info, n => inductInfo {info with name := n}
64+
| ctorInfo info, n => ctorInfo {info with name := n}
65+
| recInfo info, n => recInfo {info with name := n}
66+
67+
def updateType : ConstantInfo → Expr → ConstantInfo
68+
| defnInfo info, y => defnInfo {info with type := y}
69+
| axiomInfo info, y => axiomInfo {info with type := y}
70+
| thmInfo info, y => thmInfo {info with type := y}
71+
| opaqueInfo info, y => opaqueInfo {info with type := y}
72+
| quotInfo info, y => quotInfo {info with type := y}
73+
| inductInfo info, y => inductInfo {info with type := y}
74+
| ctorInfo info, y => ctorInfo {info with type := y}
75+
| recInfo info, y => recInfo {info with type := y}
76+
77+
def updateValue : ConstantInfo → Expr → ConstantInfo
78+
| defnInfo info, v => defnInfo {info with value := v}
79+
| thmInfo info, v => thmInfo {info with value := v}
80+
| opaqueInfo info, v => opaqueInfo {info with value := v}
81+
| d, v => d
82+
83+
def toDeclaration! : ConstantInfo → Declaration
84+
| defnInfo info => Declaration.defnDecl info
85+
| thmInfo info => Declaration.thmDecl info
86+
| axiomInfo info => Declaration.axiomDecl info
87+
| opaqueInfo info => Declaration.opaqueDecl info
88+
| quotInfo info => panic! "toDeclaration for quotInfo not implemented"
89+
| inductInfo info => panic! "toDeclaration for inductInfo not implemented"
90+
| ctorInfo info => panic! "toDeclaration for ctorInfo not implemented"
91+
| recInfo info => panic! "toDeclaration for recInfo not implemented"
92+
93+
end ConstantInfo
94+
95+
namespace Expr
96+
97+
/-! ### Declarations about `Expr` -/
98+
99+
/-- If the expression is a constant, return that name. Otherwise return `Name.anonymous`. -/
100+
def constName (e : Expr) : Name :=
101+
e.constName?.getD Name.anonymous
102+
103+
def bvarIdx? : Expr → Option Nat
104+
| bvar idx _ => some idx
105+
| _ => none
106+
107+
/-- Return the function (name) and arguments of an application. -/
108+
def getAppFnArgs (e : Expr) : Name × Array Expr :=
109+
withApp e λ e a => (e.constName, a)
110+
111+
/-- Turn an expression that is a natural number literal into a natural number. -/
112+
def natLit! : Expr → Nat
113+
| lit (Literal.natVal v) _ => v
114+
| _ => panic! "nat literal expected"
115+
116+
/-- Returns a `NameSet` of all constants in an expression starting with a certain prefix. -/
117+
def listNamesWithPrefix (pre : Name) (e : Expr) : NameSet :=
118+
e.foldConsts ∅ fun n l => if n.getPrefix == pre then l.insert n else l
119+
120+
def modifyAppArgM [Functor M] [Pure M] (modifier : Expr → M Expr) : Expr → M Expr
121+
| app f a _ => mkApp f <$> modifier a
122+
| e => pure e
123+
124+
def modifyAppArg (modifier : Expr → Expr) : Expr → Expr :=
125+
modifyAppArgM (M := Id) modifier
126+
127+
def modifyRevArg (modifier : Expr → Expr): Nat → Expr → Expr
128+
| 0 => modifyAppArg modifier
129+
| (i+1) => modifyAppArg (modifyRevArg modifier i)
130+
131+
/-- Given `f a₀ a₁ ... aₙ₋₁`, runs `modifier` on the `i`th argument or returns the original expression if out of bounds. -/
132+
def modifyArg (modifier : Expr → Expr) (e : Expr) (i : Nat) (n := e.getAppNumArgs) : Expr :=
133+
modifyRevArg modifier (n - i - 1) e
134+
135+
def getRevArg? : Expr → Nat → Option Expr
136+
| app f a _, 0 => a
137+
| app f _ _, i+1 => getRevArg! f i
138+
| _, _ => none
139+
140+
/-- Given `f a₀ a₁ ... aₙ₋₁`, returns the `i`th argument or none if out of bounds. -/
141+
def getArg? (e : Expr) (i : Nat) (n := e.getAppNumArgs): Option Expr :=
142+
getRevArg? e (n - i - 1)
143+
144+
/-- Given `f a₀ a₁ ... aₙ₋₁`, runs `modifier` on the `i`th argument.
145+
An argument `n` may be provided which says how many arguments we are expecting `e` to have. -/
146+
def modifyArgM [Monad M] (modifier : Expr → M Expr) (e : Expr) (i : Nat) (n := e.getAppNumArgs) : M Expr := do
147+
let some a := getArg? e i | return e
148+
let a ← modifier a
149+
return modifyArg (fun _ => a) e i n
150+
151+
end Expr
152+
153+
end Lean

‎Mathlib/Lean/Expr/ReplaceRec.lean‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,26 @@ def replaceRec (f? : (Expr → Expr) → Expr → Option Expr) : Expr → Expr :
3030
| some x => x
3131
| none => traverseChildren (M := Id) r e
3232

33+
/-- replaceRec under a monad. -/
34+
def replaceRecM [Monad M] (f? : (Expr → M Expr) → Expr → M (Option Expr)) : Expr → M Expr :=
35+
fix fun r e => do
36+
match ← f? r e with
37+
| some x => return x
38+
| none => traverseChildren r e
39+
40+
/-- Similar to `replaceRecM` except that bound variables are instantiated with free variables
41+
(like `Lean.Meta.transform`).
42+
This means that MetaM tactics can be used inside the replacement function.
43+
44+
If you don't need recursive calling, you should prefer using `Lean.Meta.transform` because it also caches visits.
45+
-/
46+
def replaceRecMeta [Monad M] [MonadLiftT MetaM M] [MonadControlT MetaM M]
47+
(f? : (Expr → M Expr) → Expr → M (Option Expr)) : Expr → M Expr :=
48+
fix fun r e => do
49+
match ← f? r e with
50+
| some x => return x
51+
| none => Lean.Meta.traverseChildren r e
52+
3353
/-- A version of `Expr.replace` where we can use recursive calls even if we replace a subexpression.
3454
When reaching a subexpression `e` we call `traversal e` to see if we want to do anything with this
3555
expression. If `traversal e = none` we proceed to the children of `e`. If

‎Mathlib/Lean/Expr/Traverse.lean‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ Authors: E.W.Ayers
55
-/
66

77
import Lean
8+
import Mathlib.Lean.Expr.Basic
89

910
/-!
1011
# Traversal functions for expressions.
@@ -22,4 +23,55 @@ def traverseChildren [Applicative M] (f : Expr → M Expr) : Expr → M Expr
2223
| e@(proj _ _ b _) => e.updateProj! <$> f b
2324
| e => pure e
2425

26+
/-- Given `e = fn a₁ ... aₙ`, runs `f` on `fn` and each of the arguments `aᵢ` and
27+
makes a new function application with the results. -/
28+
def traverseApp {M} [Monad M]
29+
(f : Expr → M Expr) (e : Expr) : M Expr :=
30+
e.withApp fun fn args => (pure mkAppN) <*> (f fn) <*> (args.mapM f)
31+
32+
2533
end Lean.Expr
34+
35+
namespace Lean.Meta
36+
37+
variable {M} [Monad M] [MonadLiftT MetaM M] [MonadControlT MetaM M]
38+
39+
/-- Given an expression `fun (x₁ : α₁) ... (xₙ : αₙ) => b`, will run
40+
`f` on each of the variable types `αᵢ` and `b` with the correct MetaM context,
41+
replacing each expression with the output of `f` and creating a new lambda.
42+
(that is, correctly instantiating bound variables and repackaging them after) -/
43+
def traverseLambda
44+
(f : Expr → M Expr) (fvars : Array Expr := #[]) (e : Expr) : M Expr :=
45+
match e with
46+
| (Expr.lam n d b c) => do withLocalDecl n c.binderInfo (← f (d.instantiateRev fvars)) fun x => traverseLambda f (fvars.push x) b
47+
| e => do mkLambdaFVars (usedLetOnly := false) fvars (← f (e.instantiateRev fvars))
48+
49+
/-- Given an expression ` (x₁ : α₁) → ... → (xₙ : αₙ) → b`, will run
50+
`f` on each of the variable types `αᵢ` and `b` with the correct MetaM context,
51+
replacing the expression with the output of `f` and creating a new forall expression.
52+
(that is, correctly instantiating bound variables and repackaging them after) -/
53+
def traverseForall
54+
(f : Expr → M Expr) (fvars : Array Expr := #[]) (e : Expr) : M Expr :=
55+
match e with
56+
| (Expr.forallE n d b c) => do withLocalDecl n c.binderInfo (← f (d.instantiateRev fvars)) fun x => traverseForall f (fvars.push x) b
57+
| e => do mkForallFVars (usedLetOnly := false) fvars (← f (e.instantiateRev fvars))
58+
59+
/-- Similar to traverseLambda and traverseForall but with let binders. -/
60+
def traverseLet
61+
(f : Expr → M Expr) (fvars : Array Expr := #[]) (e : Expr) : M Expr := do
62+
match e with
63+
| Expr.letE n t v b _ =>
64+
withLetDecl n (← f (t.instantiateRev fvars)) (← f (v.instantiateRev fvars)) fun x =>
65+
traverseLet f (fvars.push x) b
66+
| e => mkLetFVars (usedLetOnly := false) fvars (← f (e.instantiateRev fvars))
67+
68+
/-- A version of traverseChildren that automatically makes free variables etc. -/
69+
def traverseChildren (f : Expr → M Expr) (e : Expr) : M Expr :=
70+
match e with
71+
| Expr.forallE .. => traverseForall f #[] e
72+
| Expr.lam .. => traverseLambda f #[] e
73+
| Expr.letE .. => traverseLet f #[] e
74+
| Expr.app .. => Lean.Expr.traverseApp f e
75+
| _ => Lean.Expr.traverseChildren f e
76+
77+
end Lean.Meta

0 commit comments

Comments
 (0)