67 lines
2.4 KiB
Text
67 lines
2.4 KiB
Text
/-
|
|
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
|
|
Released under Apache 2.0 license as described in the file LICENSE.
|
|
Authors: Leonardo de Moura
|
|
-/
|
|
import Lean.Expr
|
|
|
|
namespace Lean
|
|
|
|
namespace CollectLevelParams
|
|
|
|
structure State :=
|
|
(visitedLevel : LevelSet := {})
|
|
(visitedExpr : ExprSet := {})
|
|
(params : Array Name := #[])
|
|
|
|
instance State.inhabited : Inhabited State := ⟨{}⟩
|
|
|
|
abbrev Visitor := State → State
|
|
|
|
@[inline] def visitLevel (f : Level → Visitor) (u : Level) : Visitor :=
|
|
fun s =>
|
|
if !u.hasParam || s.visitedLevel.contains u then s
|
|
else f u { s with visitedLevel := s.visitedLevel.insert u }
|
|
|
|
partial def collect : Level → Visitor
|
|
| Level.succ v _ => visitLevel collect v
|
|
| Level.max u v _ => visitLevel collect v ∘ visitLevel collect u
|
|
| Level.imax u v _ => visitLevel collect v ∘ visitLevel collect u
|
|
| Level.param n _ => fun s => { s with params := s.params.push n }
|
|
| _ => id
|
|
|
|
@[inline] def visitExpr (f : Expr → Visitor) (e : Expr) : Visitor :=
|
|
fun s =>
|
|
if !e.hasLevelParam then s
|
|
else if s.visitedExpr.contains e then s
|
|
else f e { s with visitedExpr := s.visitedExpr.insert e }
|
|
|
|
partial def main : Expr → Visitor
|
|
| Expr.proj _ _ s _ => visitExpr main s
|
|
| Expr.forallE _ d b _ => visitExpr main b ∘ visitExpr main d
|
|
| Expr.lam _ d b _ => visitExpr main b ∘ visitExpr main d
|
|
| Expr.letE _ t v b _ => visitExpr main b ∘ visitExpr main v ∘ visitExpr main t
|
|
| Expr.app f a _ => visitExpr main a ∘ visitExpr main f
|
|
| Expr.mdata _ b _ => visitExpr main b
|
|
| Expr.const _ us _ => fun s => us.foldl (fun s u => visitLevel collect u s) s
|
|
| Expr.sort u _ => visitLevel collect u
|
|
| _ => id
|
|
|
|
private partial def getUnusedLevelParamAux (s : CollectLevelParams.State) (pre : Name) : Nat → Level
|
|
| i =>
|
|
let v := mkLevelParam (pre.appendIndexAfter i);
|
|
if s.visitedLevel.contains v then getUnusedLevelParamAux (i+1) else v
|
|
|
|
def State.getUnusedLevelParam (s : CollectLevelParams.State) (pre : Name := `v) : Level :=
|
|
let v := mkLevelParam pre;
|
|
if s.visitedLevel.contains v then getUnusedLevelParamAux s pre 1 else v
|
|
|
|
end CollectLevelParams
|
|
|
|
def collectLevelParams (s : CollectLevelParams.State) (e : Expr) : CollectLevelParams.State :=
|
|
CollectLevelParams.main e s
|
|
|
|
def CollectLevelParams.State.collect (s : CollectLevelParams.State) (e : Expr) : CollectLevelParams.State :=
|
|
collectLevelParams s e
|
|
|
|
end Lean
|