lean4-htt/src/Lean/Util/FindLevelMVar.lean
2022-06-07 16:37:45 -07:00

46 lines
1.6 KiB
Text

/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Daniel Selsam
-/
import Lean.Expr
namespace Lean
namespace FindLevelMVar
abbrev Visitor := Option MVarId → Option MVarId
mutual
partial def visit (p : MVarId → Bool) (e : Expr) : Visitor := fun s =>
if s.isSome || !e.hasLevelMVar then s else main p e s
partial def main (p : MVarId → Bool) : Expr → Visitor
| Expr.sort l _ => visitLevel p l
| Expr.const _ ls _ => ls.foldr (init := id) fun l acc => visitLevel p l ∘ acc
| Expr.forallE _ d b _ => visit p b ∘ visit p d
| Expr.lam _ d b _ => visit p b ∘ visit p d
| Expr.letE _ t v b _ => visit p b ∘ visit p v ∘ visit p t
| Expr.app f a _ => visit p a ∘ visit p f
| Expr.mdata _ b _ => visit p b
| Expr.proj _ _ e _ => visit p e
| _ => id
partial def visitLevel (p : MVarId → Bool) (l : Level) : Visitor := fun s =>
if s.isSome || !l.hasMVar then s else mainLevel p l s
partial def mainLevel (p : MVarId → Bool) : Level → Visitor
| Level.zero _ => id
| Level.succ l _ => visitLevel p l
| Level.max l₁ l₂ _ => visitLevel p l₁ ∘ visitLevel p l₂
| Level.imax l₁ l₂ _ => visitLevel p l₁ ∘ visitLevel p l₂
| Level.param _ _ => id
| Level.mvar mvarId _ => fun s => if p mvarId then some mvarId else s
end
end FindLevelMVar
@[inline] def Expr.findLevelMVar? (e : Expr) (p : MVarId → Bool) : Option MVarId :=
FindLevelMVar.main p e none
end Lean