175 lines
7.7 KiB
Text
175 lines
7.7 KiB
Text
/-
|
||
Copyright (c) 2022 E.W.Ayers. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: E.W.Ayers
|
||
-/
|
||
import Lean.Meta.Basic
|
||
import Lean.SubExpr
|
||
|
||
/-!
|
||
|
||
# Expression Lenses
|
||
|
||
Functions for manipulating subexpressions using `SubExpr.Pos`.
|
||
|
||
-/
|
||
|
||
namespace Lean.Meta
|
||
|
||
section ExprLens
|
||
|
||
open Lean.SubExpr
|
||
|
||
variable {M} [Monad M] [MonadLiftT MetaM M] [MonadControlT MetaM M] [MonadError M]
|
||
|
||
/-- Given a constructor index for Expr, runs `g` on the value of that subexpression and replaces it.
|
||
If the subexpression is under a binder it will instantiate and abstract the binder body correctly.
|
||
Mdata is ignored. An index of 3 is interpreted as the type of the expression. An index of 3 will throw since we can't replace types.
|
||
|
||
See also `Lean.Meta.transform`, `Lean.Meta.traverseChildren`. -/
|
||
private def lensCoord (g : Expr → M Expr) : Nat → Expr → M Expr
|
||
| 0, e@(Expr.app f a _) => return e.updateApp! (← g f) a
|
||
| 1, e@(Expr.app f a _) => return e.updateApp! f (← g a)
|
||
| 0, e@(Expr.lam _ y b _) => return e.updateLambdaE! (← g y) b
|
||
| 1, (Expr.lam n y b c) => withLocalDecl n c.binderInfo y fun x => do mkLambdaFVars #[x] <|← g <| b.instantiateRev #[x]
|
||
| 0, e@(Expr.forallE _ y b _) => return e.updateForallE! (← g y) b
|
||
| 1, (Expr.forallE n y b c) => withLocalDecl n c.binderInfo y fun x => do mkForallFVars #[x] <|← g <| b.instantiateRev #[x]
|
||
| 0, e@(Expr.letE _ y a b _) => return e.updateLet! (← g y) a b
|
||
| 1, e@(Expr.letE _ y a b _) => return e.updateLet! y (← g a) b
|
||
| 2, (Expr.letE n y a b _) => withLetDecl n y a fun x => do mkLetFVars #[x] <|← g <| b.instantiateRev #[x]
|
||
| 0, e@(Expr.proj _ _ b _) => e.updateProj! <$> g b
|
||
| n, e@(Expr.mdata _ a _) => e.updateMData! <$> lensCoord g n a
|
||
| 3, _ => throwError "Lensing on types is not supported"
|
||
| c, e => throwError "Invalid coordinate {c} for {e}"
|
||
|
||
private def lensAux (g : Expr → M Expr) : List Nat → Expr → M Expr
|
||
| [], e => g e
|
||
| head::tail, e => lensCoord (lensAux g tail) head e
|
||
|
||
/-- Run the given `replace` function to replace the expression at the subexpression position. If the subexpression is below a binder
|
||
the bound variables will be appropriately instantiated with free variables and reabstracted after the replacement.
|
||
If the subexpression is invalid or points to a type then this will throw. -/
|
||
def replaceSubexpr (replace : (subexpr : Expr) → M Expr) (p : Pos) (root : Expr) : M Expr :=
|
||
lensAux replace p.toArray.toList root
|
||
|
||
/-- Runs `k` on the given coordinate, including handling binders properly.
|
||
The subexpression value passed to `k` is not instantiated with respect to the
|
||
array of free variables. -/
|
||
private def viewCoordAux (k : Array Expr → Expr → M α) (fvars: Array Expr) : Nat → Expr → M α
|
||
| 3, _ => throwError "Internal: Types should be handled by viewAux"
|
||
| 0, (Expr.app f _ _) => k fvars f
|
||
| 1, (Expr.app _ a _) => k fvars a
|
||
| 0, (Expr.lam _ y _ _) => k fvars y
|
||
| 1, (Expr.lam n y b c) => withLocalDecl n c.binderInfo (y.instantiateRev fvars) fun x => k (fvars.push x) b
|
||
| 0, (Expr.forallE _ y _ _) => k fvars y
|
||
| 1, (Expr.forallE n y b c) => withLocalDecl n c.binderInfo (y.instantiateRev fvars) fun x => k (fvars.push x) b
|
||
| 0, (Expr.letE _ y _ _ _) => k fvars y
|
||
| 1, (Expr.letE _ _ a _ _) => k fvars a
|
||
| 2, (Expr.letE n y a b _) => withLetDecl n (y.instantiateRev fvars) (a.instantiateRev fvars) fun x => k (fvars.push x) b
|
||
| 0, (Expr.proj _ _ b _) => k fvars b
|
||
| n, (Expr.mdata _ a _) => viewCoordAux k fvars n a
|
||
| c, e => throwError "Invalid coordinate {c} for {e}"
|
||
|
||
private def viewAux (k : Array Expr → Expr → M α) (fvars : Array Expr) : List Nat → Expr → M α
|
||
| [], e => k fvars <| e.instantiateRev fvars
|
||
| 3::tail, e => do
|
||
let y ← inferType <| e.instantiateRev fvars
|
||
viewAux (fun otherFvars => k (fvars ++ otherFvars)) #[] tail y
|
||
| head::tail, e => viewCoordAux (fun fvars => viewAux k fvars tail) fvars head e
|
||
|
||
/-- `view visit p e` runs `visit fvars s` where `s : Expr` is the subexpression of `e` at `p`.
|
||
and `fvars` are the free variables for the binders that `s` is under.
|
||
`s` is already instantiated with respect to these.
|
||
The role of the `visit` function is analogous to the `k` function in `Lean.Meta.forallTelescope`. -/
|
||
def viewSubexpr
|
||
(visit : (fvars : Array Expr) → (subexpr : Expr) → M α)
|
||
(p : Pos) (root : Expr) : M α :=
|
||
viewAux visit #[] p.toArray.toList root
|
||
|
||
private def foldAncestorsAux
|
||
(k : Array Expr → Expr → Nat → α → M α)
|
||
(acc : α) (address : List Nat) (fvars : Array Expr) (current : Expr) : M α :=
|
||
match address with
|
||
| [] => return acc
|
||
| 3 :: tail => do
|
||
let current := current.instantiateRev fvars
|
||
let y ← inferType current
|
||
let acc ← k fvars current 3 acc
|
||
foldAncestorsAux (fun otherFvars => k (fvars ++ otherFvars)) acc tail #[] y
|
||
| head :: tail => do
|
||
let acc ← k fvars (current.instantiateRev fvars) head acc
|
||
viewCoordAux (foldAncestorsAux k acc tail) fvars head current
|
||
|
||
/-- `foldAncestors k init p e` folds over the strict ancestor subexpressions of the given expression `e` above position `p`, starting at the root expression and working down.
|
||
The fold function `k` is given the newly instantiated free variables, the ancestor subexpression, and the coordinate
|
||
that will be explored next.-/
|
||
def foldAncestors
|
||
(k : (fvars: Array Expr) → (subexpr : Expr) → (nextCoord : Nat) → α → M α)
|
||
(init : α) (p : Pos) (e : Expr) : M α :=
|
||
foldAncestorsAux k init p.toArray.toList #[] e
|
||
|
||
end ExprLens
|
||
|
||
end Lean.Meta
|
||
|
||
namespace Lean.Core
|
||
|
||
open Lean.SubExpr
|
||
|
||
section ViewRaw
|
||
|
||
variable {M} [Monad M] [MonadError M]
|
||
|
||
/-- Get the raw subexpression without performing any instantiation. -/
|
||
private def viewCoordRaw: Expr → Nat → M Expr
|
||
| e , 3 => throwError "Can't viewRaw the type of {e}"
|
||
| (Expr.app f _ _) , 0 => pure f
|
||
| (Expr.app _ a _) , 1 => pure a
|
||
| (Expr.lam _ y _ _) , 0 => pure y
|
||
| (Expr.lam _ _ b _) , 1 => pure b
|
||
| (Expr.forallE _ y _ _), 0 => pure y
|
||
| (Expr.forallE _ _ b _), 1 => pure b
|
||
| (Expr.letE _ y _ _ _) , 0 => pure y
|
||
| (Expr.letE _ _ a _ _) , 1 => pure a
|
||
| (Expr.letE _ _ _ b _) , 2 => pure b
|
||
| (Expr.proj _ _ b _) , 0 => pure b
|
||
| (Expr.mdata _ a _) , n => viewCoordRaw a n
|
||
| e , c => throwError "Bad coordinate {c} for {e}"
|
||
|
||
|
||
/-- Given a valid SubExpr, will return the raw current expression without performing any instantiation.
|
||
If the SubExpr has a type subexpression coordinate then will error.
|
||
|
||
This is a cheaper version of `Lean.Meta.viewSubexpr` and can be used to quickly view the
|
||
subexpression at a position. Note that because the resulting expression will contain
|
||
loose bound variables it can't be used in any `MetaM` methods. -/
|
||
def viewSubexpr (p : Pos) (root : Expr) : M Expr :=
|
||
p.foldlM viewCoordRaw root
|
||
|
||
private def viewBindersCoord : Nat → Expr → Option (Name × Expr)
|
||
| 1, (Expr.lam n y _ _) => some (n, y)
|
||
| 1, (Expr.forallE n y _ _) => some (n, y)
|
||
| 2, (Expr.letE n y _ _ _) => some (n, y)
|
||
| _, _ => none
|
||
|
||
/-- `viewBinders p e` returns a list of all of the binders (name, type) above the given position `p` in the root expression `e` -/
|
||
def viewBinders (p : Pos) (root : Expr) : M (Array (Name × Expr)) := do
|
||
let (acc, _) ← p.foldlM (fun (acc, e) c => do
|
||
let e₂ ← viewCoordRaw e c
|
||
let acc :=
|
||
match viewBindersCoord c e with
|
||
| none => acc
|
||
| some b => acc.push b
|
||
return (acc, e₂)
|
||
) (#[], root)
|
||
return acc
|
||
|
||
/-- Returns the number of binders above a given subexpr position. -/
|
||
def numBinders (p : Pos) (e : Expr) : M Nat :=
|
||
Array.size <$> viewBinders p e
|
||
|
||
end ViewRaw
|
||
|
||
end Lean.Core
|
||
|
||
|