feat: generalise forEachExpr

Lean.Meta.forEachExpr should be general over monads rather than restricted to the MetaM monad.
This is similar to the generalisation of Lean.Meta.transform
This commit is contained in:
E.W.Ayers 2022-08-23 14:28:14 +01:00 committed by Leonardo de Moura
parent e81ba951c6
commit d18667c484

View file

@ -8,60 +8,58 @@ import Lean.Util.MonadCache
import Lean.Meta.Basic
namespace Lean.Meta
namespace ForEachExpr
abbrev M := MonadCacheT Expr Unit MetaM
variable {m} [Monad m] [MonadLiftT MetaM m] [MonadControlT MetaM m]
mutual
private partial def visitBinder (fn : Expr → MetaM Bool) : Array Expr → Nat → Expr → M Unit
| fvars, j, Expr.lam n d b c => do
let d := d.instantiateRevRange j fvars.size fvars
visit fn d
withLocalDecl n c d fun x =>
visitBinder fn (fvars.push x) j b
| fvars, j, Expr.forallE n d b c => do
let d := d.instantiateRevRange j fvars.size fvars
visit fn d
withLocalDecl n c d fun x =>
visitBinder fn (fvars.push x) j b
| fvars, j, Expr.letE n t v b _ => do
let t := t.instantiateRevRange j fvars.size fvars
visit fn t
let v := v.instantiateRevRange j fvars.size fvars
visit fn v
withLetDecl n t v fun x =>
visitBinder fn (fvars.push x) j b
| fvars, j, e => visit fn $ e.instantiateRevRange j fvars.size fvars
partial def visit (fn : Expr → MetaM Bool) (e : Expr) : M Unit :=
checkCache e fun _ => do
if (← liftM (fn e)) then
match e with
| .forallE .. => visitBinder fn #[] 0 e
| .lam .. => visitBinder fn #[] 0 e
| .letE .. => visitBinder fn #[] 0 e
| .app f a => visit fn f; visit fn a
| .mdata _ b => visit fn b
| .proj _ _ b => visit fn b
| _ => return ()
end
end ForEachExpr
/-- Similar to `Expr.forEach'`, but creates free variables whenever going inside of a binder. -/
def forEachExpr' (e : Expr) (f : Expr → MetaM Bool) : MetaM Unit :=
ForEachExpr.visit f e |>.run
/-- Similar to `Expr.forEach'`, but creates free variables whenever going inside of a binder.
If the inner function returns `false`, deeper subexpressions will not be visited.
-/
partial def forEachExpr'
(input : Expr)
(fn : Expr → m Bool)
: m Unit := do
let _ : STWorld IO.RealWorld m := ⟨⟩
let _ : MonadLiftT (ST IO.RealWorld) m := { monadLift := fun x => liftM (m := MetaM) (liftM (m := ST IO.RealWorld) x) }
let rec visit (e : Expr) : MonadCacheT Expr Unit m Unit :=
let rec visitBinder (fvars : Array Expr) (j : Nat) : Expr → MonadCacheT Expr Unit m Unit
| Expr.lam n d b c => do
let d := d.instantiateRevRange j fvars.size fvars
visit d
withLocalDecl n c d fun x =>
visitBinder (fvars.push x) j b
| Expr.forallE n d b c => do
let d := d.instantiateRevRange j fvars.size fvars
visit d
withLocalDecl n c d fun x =>
visitBinder (fvars.push x) j b
| Expr.letE n t v b _ => do
let t := t.instantiateRevRange j fvars.size fvars
visit t
let v := v.instantiateRevRange j fvars.size fvars
visit v
withLetDecl n t v fun x =>
visitBinder (fvars.push x) j b
| e => visit $ e.instantiateRevRange j fvars.size fvars
checkCache e fun _ => do
if (← liftM (fn e)) then
match e with
| .forallE .. => visitBinder #[] 0 e
| .lam .. => visitBinder #[] 0 e
| .letE .. => visitBinder #[] 0 e
| .app f a => visit f; visit a
| .mdata _ b => visit b
| .proj _ _ b => visit b
| _ => return ()
visit input |>.run
/-- Similar to `Expr.forEach`, but creates free variables whenever going inside of a binder. -/
def forEachExpr (e : Expr) (f : Expr → MetaM Unit) : MetaM Unit :=
def forEachExpr (e : Expr) (f : Expr → m Unit) : m Unit :=
forEachExpr' e fun e => do
f e
return true
/-- Return true iff `x` is a metavariable with an anonymous user facing name. -/
private def shouldInferBinderName (x : Expr) : MetaM Bool := do
private def shouldInferBinderName (x : Expr) : m Bool := do
match x with
| .mvar mvarId => return (← mvarId.getDecl).userName.isAnonymous
| _ => return false