/-
Copyright (c) 2018 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
module

prelude
public import Init.Data.Array.QSort
public import Lean.Data.PersistentHashMap
public import Lean.Data.PersistentHashSet
public import Lean.Hygiene
public import Lean.Data.Name
public import Lean.Data.Format
public import Init.Data.Option.Coe
public import Std.Data.TreeSet.Basic

public section

def Nat.imax (n m : Nat) : Nat :=
  if m = 0 then 0 else Nat.max n m

namespace Lean

/--
 Cached hash code, cached results, and other data for `Level`.
   hash      : 32-bits
   hasMVar   : 1-bit
   hasParam  : 1-bit
   depth     : 24-bits -/
@[expose] def Level.Data := UInt64

instance : Inhabited Level.Data :=
  inferInstanceAs (Inhabited UInt64)

def Level.Data.hash (c : Level.Data) : UInt64 :=
  c.toUInt32.toUInt64

instance : BEq Level.Data :=
  ⟨fun (a b : UInt64) => a == b⟩

def Level.Data.depth (c : Level.Data) : UInt32 :=
  (c.shiftRight 40).toUInt32

def Level.Data.hasMVar (c : Level.Data) : Bool :=
  ((c.shiftRight 32).land 1) == 1

def Level.Data.hasParam (c : Level.Data) : Bool :=
  ((c.shiftRight 33).land 1) == 1

@[extern "lean_level_mk_data"]
opaque Level.mkData (h : UInt64) (depth : Nat := 0) (hasMVar hasParam : Bool := false) : Level.Data

instance : Repr Level.Data where
  reprPrec v prec := Id.run do
    let mut r := "Level.mkData " ++ toString v.hash
    if v.depth != 0 then
      r := r ++ " (depth := " ++ toString v.depth ++ ")"
    if v.hasMVar then
      r := r ++ " (hasMVar := " ++ toString v.hasMVar ++ ")"
    if v.hasParam then
      r := r ++ " (hasParam := " ++ toString v.hasParam ++ ")"
    Repr.addAppParen r prec

open Level

/-- Universe level metavariable Id   -/
structure LevelMVarId where
  name : Name
  deriving Inhabited, BEq, Hashable, Repr

/-- Short for `LevelMVarId` -/
abbrev LMVarId := LevelMVarId

instance : Repr LMVarId where
  reprPrec n p := reprPrec n.name p

@[expose] def LMVarIdSet := Std.TreeSet LMVarId (Name.quickCmp ·.name ·.name)
  deriving Inhabited, EmptyCollection

instance : ForIn m LMVarIdSet LMVarId := inferInstanceAs (ForIn _ (Std.TreeSet _ _) ..)

@[expose] def LMVarIdMap (α : Type) := Std.TreeMap LMVarId α (Name.quickCmp ·.name ·.name)

instance : EmptyCollection (LMVarIdMap α) := inferInstanceAs (EmptyCollection (Std.TreeMap _ _ _))

instance : ForIn m (LMVarIdMap α) (LMVarId × α) := inferInstanceAs (ForIn _ (Std.TreeMap _ _ _) ..)

instance : Inhabited (LMVarIdMap α) where
  default := {}

inductive Level where
  | zero   : Level
  | succ   : Level → Level
  | max    : Level → Level → Level
  | imax   : Level → Level → Level
  | param  : Name → Level
  | mvar   : LMVarId → Level
with
  @[computed_field] data : Level → Data
    | .zero => mkData 2221 0 false false
    | .mvar mvarId => mkData (mixHash 2237 <| hash mvarId) 0 true false
    | .param name => mkData (mixHash 2239 <| hash name) 0 false true
    | .succ u => mkData (mixHash 2243 <| u.data.hash) (u.data.depth.toNat + 1) u.data.hasMVar u.data.hasParam
    | .max u v => mkData (mixHash 2251 <| mixHash (u.data.hash) (v.data.hash)) (Nat.max u.data.depth.toNat v.data.depth.toNat + 1)
       (u.data.hasMVar || v.data.hasMVar) (u.data.hasParam || v.data.hasParam)
    | .imax u v => mkData (mixHash 2267 <| mixHash (u.data.hash) (v.data.hash)) (Nat.max u.data.depth.toNat v.data.depth.toNat + 1)
       (u.data.hasMVar || v.data.hasMVar) (u.data.hasParam || v.data.hasParam)

deriving Inhabited, Repr

namespace Level

protected def hash (u : Level) : UInt64 :=
  u.data.hash

instance : Hashable Level := ⟨Level.hash⟩

def depth (u : Level) : Nat :=
  u.data.depth.toNat

def hasMVar (u : Level) : Bool :=
  u.data.hasMVar

def hasParam (u : Level) : Bool :=
  u.data.hasParam

@[export lean_level_hash] def hashEx (u : Level) : UInt32 := hash u |>.toUInt32
@[export lean_level_has_mvar] def hasMVarEx : Level → Bool := hasMVar
@[export lean_level_has_param] def hasParamEx : Level → Bool := hasParam
@[export lean_level_depth] def depthEx (u : Level) : UInt32 := u.data.depth

end Level

@[expose] def levelZero :=
  Level.zero

def mkLevelMVar (mvarId : LMVarId) :=
  Level.mvar mvarId

def mkLevelParam (name : Name) :=
  Level.param name

@[expose] def mkLevelSucc (u : Level) :=
  Level.succ u

def mkLevelMax (u v : Level) :=
  Level.max u v

def mkLevelIMax (u v : Level) :=
  Level.imax u v

def levelOne := mkLevelSucc levelZero

@[export lean_level_mk_zero] def mkLevelZeroEx : Unit → Level := fun _ => levelZero
@[export lean_level_mk_succ] def mkLevelSuccEx : Level → Level := mkLevelSucc
@[export lean_level_mk_mvar] def mkLevelMVarEx : LMVarId → Level := mkLevelMVar
@[export lean_level_mk_param] def mkLevelParamEx : Name → Level := mkLevelParam
@[export lean_level_mk_max] def mkLevelMaxEx : Level → Level → Level := mkLevelMax
@[export lean_level_mk_imax] def mkLevelIMaxEx : Level → Level → Level := mkLevelIMax

namespace Level

def isZero : Level → Bool
  | zero   => true
  | _      => false

def isSucc : Level → Bool
  | succ .. => true
  | _       => false

def isMax : Level → Bool
  | max .. => true
  | _      => false

def isIMax : Level → Bool
  | imax .. => true
  | _       => false

def isMaxIMax : Level → Bool
  | max ..  => true
  | imax .. => true
  | _       => false

def isParam : Level → Bool
  | param .. => true
  | _        => false

def isMVar : Level → Bool
  | mvar .. => true
  | _       => false

def mvarId! : Level → LMVarId
  | mvar mvarId => mvarId
  | _           => panic! "metavariable expected"

/-- If result is true, then forall assignments `A` which assigns all parameters and metavariables occurring
    in `l`, `l[A] != zero` -/
def isNeverZero : Level → Bool
  | zero         => false
  | param ..     => false
  | mvar ..      => false
  | succ ..      => true
  | max l₁ l₂    => isNeverZero l₁ || isNeverZero l₂
  | imax _  l₂   => isNeverZero l₂

/--
Returns true if and only if `l` evaluates to zero for all instantiations of parameters and
meta-variables.
-/
def isAlwaysZero : Level → Bool
  | zero         => true
  | param ..     => false
  | mvar ..      => false
  | succ ..      => false
  | max l₁ l₂    => isAlwaysZero l₁ && isAlwaysZero l₂
  | imax _  l₂   => isAlwaysZero l₂

@[expose] def ofNat : Nat → Level
  | 0   => levelZero
  | n+1 => mkLevelSucc (ofNat n)

instance instOfNat (n : Nat) : OfNat Level n where
  ofNat := ofNat n

def addOffsetAux : Nat → Level → Level
  | 0,     u => u
  | (n+1), u => addOffsetAux n (mkLevelSucc u)

def addOffset (u : Level) (n : Nat) : Level :=
  u.addOffsetAux n

def isExplicit : Level → Bool
  | zero     => true
  | succ u   => !u.hasMVar && !u.hasParam && isExplicit u
  | _        => false

def getOffsetAux : Level → Nat → Nat
  | succ u  , r => getOffsetAux u (r+1)
  | _,        r => r

def getOffset (lvl : Level) : Nat :=
  getOffsetAux lvl 0

def getLevelOffset : Level → Level
  | succ u   => getLevelOffset u
  | u        => u

def toNat (lvl : Level) : Option Nat :=
  match lvl.getLevelOffset with
  | zero   => lvl.getOffset
  | _      => none

@[extern "lean_level_eq"]
protected opaque beq (a : @& Level) (b : @& Level) : Bool

instance : BEq Level := ⟨Level.beq⟩

/-- `occurs u l` return `true` iff `u` occurs in `l`. -/
def occurs : Level → Level → Bool
  | u, v@(succ v₁  )     => u == v || occurs u v₁
  | u, v@(max v₁ v₂  )   => u == v || occurs u v₁ || occurs u v₂
  | u, v@(imax v₁ v₂  )  => u == v || occurs u v₁ || occurs u v₂
  | u, v                 => u == v

def ctorToNat : Level → Nat
  | zero ..  => 0
  | param .. => 1
  | mvar ..  => 2
  | succ ..  => 3
  | max ..   => 4
  | imax ..  => 5

def normLtAux : Level → Nat → Level → Nat → Bool
  | succ l₁, k₁, l₂, k₂ => normLtAux l₁ (k₁+1) l₂ k₂
  | l₁, k₁, succ l₂, k₂ => normLtAux l₁ k₁ l₂ (k₂+1)
  | l₁@(max l₁₁ l₁₂), k₁, l₂@(max l₂₁ l₂₂), k₂ =>
    if l₁ == l₂ then k₁ < k₂
    else if l₁₁ != l₂₁ then normLtAux l₁₁ 0 l₂₁ 0
    else normLtAux l₁₂ 0 l₂₂ 0
  | l₁@(imax l₁₁ l₁₂), k₁, l₂@(imax l₂₁ l₂₂), k₂ =>
    if l₁ == l₂ then k₁ < k₂
    else if l₁₁ != l₂₁ then normLtAux l₁₁ 0 l₂₁ 0
    else normLtAux l₁₂ 0 l₂₂ 0
  | param n₁, k₁, param n₂, k₂ => if n₁ == n₂ then k₁ < k₂ else Name.lt n₁ n₂ -- use `Name.lt` because it is lexicographical
  /-
    We also use `Name.lt` in the following case to make sure universe parameters in a declaration
    are not affected by shifted indices. We used to use `Name.quickLt` which is not stable over shifted indices (the hashcodes change),
    and changes to the elaborator could affect the universe parameters and break code that relies on an explicit order.
    Example: test `tests/lean/343.lean`.
   -/
  | mvar n₁, k₁, mvar n₂, k₂ => if n₁ == n₂ then k₁ < k₂ else Name.lt n₁.name n₂.name
  | l₁, k₁, l₂, k₂ => if l₁ == l₂ then k₁ < k₂ else ctorToNat l₁ < ctorToNat l₂

/--
  A total order on level expressions that has the following properties
  - `succ l` is an immediate successor of `l`.
  - `zero` is the minimal element.
 This total order is used in the normalization procedure. -/
def normLt (l₁ l₂ : Level) : Bool :=
  normLtAux l₁ 0 l₂ 0

def isAlreadyNormalizedCheap : Level → Bool
  | zero    => true
  | param _ => true
  | mvar _  => true
  | succ u  => isAlreadyNormalizedCheap u
  | _       => false

/- Auxiliary function used at `normalize` -/
private def mkIMaxAux : Level → Level → Level
  | _,    zero   => zero
  | zero, u      => u
  | succ zero, u => u
  | u₁,   u₂     => if u₁ == u₂ then u₁ else mkLevelIMax u₁ u₂

/- Auxiliary function used at `normalize` -/
@[specialize] private partial def getMaxArgsAux (normalize : Level → Level) : Level → Bool → Array Level → Array Level
  | max l₁ l₂, alreadyNormalized, lvls => getMaxArgsAux normalize l₂ alreadyNormalized (getMaxArgsAux normalize l₁ alreadyNormalized lvls)
  | l,           false,             lvls => getMaxArgsAux normalize (normalize l) true lvls
  | l,           true,              lvls => lvls.push l

private def accMax (result : Level) (prev : Level) (offset : Nat) : Level :=
  if result.isZero then prev.addOffset offset
  else mkLevelMax result (prev.addOffset offset)

/- Auxiliary function used at `normalize`.
   Remarks:
   - `lvls` are sorted using `normLt`
   - `extraK` is the outer offset of the `max` term. We will push it inside.
   - `i` is the current array index
   - `prev + prevK` is the "previous" level that has not been added to `result` yet.
   - `result` is the accumulator
 -/
private partial def mkMaxAux (lvls : Array Level) (extraK : Nat) (i : Nat) (prev : Level) (prevK : Nat) (result : Level) : Level :=
  if h : i < lvls.size then
    let lvl   := lvls[i]
    let curr  := lvl.getLevelOffset
    let currK := lvl.getOffset
    if curr == prev then
      mkMaxAux lvls extraK (i+1) curr currK result
    else
      mkMaxAux lvls extraK (i+1) curr currK (accMax result prev (extraK + prevK))
  else
    accMax result prev (extraK + prevK)

/-
  Auxiliary function for `normalize`. It assumes `lvls` has been sorted using `normLt`.
  It finds the first position that is not an explicit universe. -/
private partial def skipExplicit (lvls : Array Level) (i : Nat) : Nat :=
  if h : i < lvls.size then
    let lvl := lvls[i]
    if lvl.getLevelOffset.isZero then skipExplicit lvls (i+1) else i
  else
    i

/--
Auxiliary function for `normalize`.
`maxExplicit` is the maximum explicit universe level at `lvls`.
Return true if it finds a level with offset ≥ maxExplicit.
`i` starts at the first non explicit level.
It assumes `lvls` has been sorted using `normLt`.
-/
private partial def isExplicitSubsumedAux (lvls : Array Level) (maxExplicit : Nat) (i : Nat) : Bool :=
  if h : i < lvls.size then
    let lvl := lvls[i]
    if lvl.getOffset ≥ maxExplicit then true
    else isExplicitSubsumedAux lvls maxExplicit (i+1)
  else
    false

/- Auxiliary function for `normalize`. See `isExplicitSubsumedAux` -/
private def isExplicitSubsumed (lvls : Array Level) (firstNonExplicit : Nat) : Bool :=
  if firstNonExplicit == 0 then false
  else
    let max := lvls[firstNonExplicit - 1]!.getOffset
    isExplicitSubsumedAux lvls max firstNonExplicit

partial def normalize (l : Level) : Level :=
  if isAlreadyNormalizedCheap l then l
  else
    let k := l.getOffset
    let u := l.getLevelOffset
    match u with
    | max l₁ l₂ =>
      let lvls  := getMaxArgsAux normalize l₁ false #[]
      let lvls  := getMaxArgsAux normalize l₂ false lvls
      let lvls  := lvls.qsort normLt
      let firstNonExplicit := skipExplicit lvls 0
      let i := if isExplicitSubsumed lvls firstNonExplicit then firstNonExplicit else firstNonExplicit - 1
      let lvl₁  := lvls[i]!
      let prev  := lvl₁.getLevelOffset
      let prevK := lvl₁.getOffset
      mkMaxAux lvls k (i+1) prev prevK levelZero
    | imax l₁ l₂ =>
      if l₂.isNeverZero then addOffset (normalize (mkLevelMax l₁ l₂)) k
      else
        let l₁ := normalize l₁
        let l₂ := normalize l₂
        addOffset (mkIMaxAux l₁ l₂) k
    | _ => unreachable!

/--
Return true if `u` and `v` denote the same level.
Check is currently incomplete.
-/
def isEquiv (u v : Level) : Bool :=
  u == v || u.normalize == v.normalize

/-- Reduce (if possible) universe level by 1 -/
def dec : Level → Option Level
  | zero       => none
  | param _    => none
  | mvar _     => none
  | succ l     => l
  | max l₁ l₂  => return mkLevelMax (← dec l₁) (← dec l₂)
  /- Remark: `mkLevelMax` in the following line is not a typo.
     If `dec l₂` succeeds, then `imax l₁ l₂` is equivalent to `max l₁ l₂`. -/
  | imax l₁ l₂ => return mkLevelMax (←  dec l₁) (← dec l₂)


/- Level to Format/Syntax -/
namespace PP
inductive Result where
  | leaf      : Name → Result
  | num       : Nat → Result
  | offset    : Result → Nat → Result
  | maxNode   : List Result → Result
  | imaxNode  : List Result → Result

def Result.succ : Result → Result
  | Result.offset f k => Result.offset f (k+1)
  | Result.num k      => Result.num (k+1)
  | f                 => Result.offset f 1

def Result.max : Result → Result → Result
  | f, Result.maxNode Fs => Result.maxNode (f::Fs)
  | f₁, f₂               => Result.maxNode [f₁, f₂]

def Result.imax : Result → Result → Result
  | f, Result.imaxNode Fs => Result.imaxNode (f::Fs)
  | f₁, f₂                => Result.imaxNode [f₁, f₂]

def toResult (l : Level) (mvars : Bool) : Result :=
  match l with
  | zero       => Result.num 0
  | succ l     => Result.succ (toResult l mvars)
  | max l₁ l₂  => Result.max (toResult l₁ mvars) (toResult l₂ mvars)
  | imax l₁ l₂ => Result.imax (toResult l₁ mvars) (toResult l₂ mvars)
  | param n    => Result.leaf n
  | mvar n     =>
    if mvars then
      Result.leaf <| n.name.replacePrefix `_uniq (Name.mkSimple "?u")
    else
      Result.leaf `_

private def parenIfFalse : Format → Bool → Format
  | f, true  => f
  | f, false => f.paren

mutual
  private partial def Result.formatLst : List Result → Format
    | []    => Format.nil
    | r::rs => Format.line ++ format r false ++ formatLst rs

  partial def Result.format : Result → Bool → Format
    | Result.leaf n,         _ => Std.format n
    | Result.num k,          _ => toString k
    | Result.offset f 0,     r => format f r
    | Result.offset f (k+1), r =>
      let f' := format f false;
      parenIfFalse (f' ++ "+" ++ Std.format (k+1)) r
    | Result.maxNode fs,    r => parenIfFalse (Format.group <| "max"  ++ formatLst fs) r
    | Result.imaxNode fs,   r => parenIfFalse (Format.group <| "imax" ++ formatLst fs) r
end

protected partial def Result.quote (r : Result) (prec : Nat) : Syntax.Level :=
  let addParen (s : Syntax.Level) :=
    if prec > 0 then Unhygienic.run `(level| ( $s )) else s
  match r with
  | Result.leaf n         => Unhygienic.run `(level| $(mkIdent n):ident)
  | Result.num  k         => Unhygienic.run `(level| $(quote k):num)
  | Result.offset r 0     => Result.quote r prec
  | Result.offset r (k+1) => addParen <| Unhygienic.run `(level| $(Result.quote r 65) + $(quote (k+1)):num)
  | Result.maxNode rs     => addParen <| Unhygienic.run `(level| max $(rs.toArray.map (Result.quote · max_prec))*)
  | Result.imaxNode rs    => addParen <| Unhygienic.run `(level| imax $(rs.toArray.map (Result.quote · max_prec))*)

end PP

protected def format (u : Level) (mvars : Bool) : Format :=
  (PP.toResult u mvars).format true

instance : ToFormat Level where
  format u := Level.format u (mvars := true)

instance : ToString Level where
  toString u := Format.pretty (format u)

protected def quote (u : Level) (prec : Nat := 0) (mvars : Bool := true) : Syntax.Level :=
  (PP.toResult u (mvars := mvars)).quote prec

instance : Quote Level `level where
  quote := Level.quote

end Level

@[inline] private def mkLevelMaxCore (u v : Level) (elseK : Unit → Level) : Level :=
  let subsumes (u v : Level) : Bool :=
    if v.isExplicit && u.getOffset ≥ v.getOffset then true
    else match u with
      | Level.max u₁ u₂ => v == u₁ || v == u₂
      | _ => false
  if u == v then u
  else if u.isZero then v
  else if v.isZero then u
  else if subsumes u v then u
  else if subsumes v u then v
  else if u.getLevelOffset == v.getLevelOffset then
    if u.getOffset ≥ v.getOffset then u else v
  else
    elseK ()

/- Similar to `mkLevelMax`, but applies cheap simplifications -/
def mkLevelMax' (u v : Level) : Level :=
  mkLevelMaxCore u v fun _ => mkLevelMax u v

def simpLevelMax' (u v : Level) (d : Level) : Level :=
  mkLevelMaxCore u v fun _ => d

@[inline] private def mkLevelIMaxCore (u v : Level) (elseK : Unit → Level) : Level :=
  if v.isNeverZero then mkLevelMax' u v
  else if v.isZero then v
  else if u.isZero then v
  else if u == v then u
  else elseK ()

/- Similar to `mkLevelIMax`, but applies cheap simplifications -/
def mkLevelIMax' (u v : Level) : Level :=
  mkLevelIMaxCore u v fun _ => mkLevelIMax u v

def simpLevelIMax' (u v : Level) (d : Level) :=
  mkLevelIMaxCore u v fun _ => d

namespace Level

/-!
The update functions try to avoid allocating new values using pointer equality.
Note that if the `update*!` functions are used under a match-expression,
the compiler will eliminate the double-match.
-/

@[inline] private unsafe def updateSucc!Impl (lvl : Level) (newLvl : Level) : Level :=
  match lvl with
  | succ l => if ptrEq l newLvl then lvl else mkLevelSucc newLvl
  | _      => panic! "succ level expected"

@[implemented_by updateSucc!Impl]
def updateSucc! (lvl : Level) (newLvl : Level) : Level :=
  match lvl with
  | succ _ => mkLevelSucc newLvl
  | _      => panic! "succ level expected"

@[inline] private unsafe def updateMax!Impl (lvl : Level) (newLhs : Level) (newRhs : Level) : Level :=
  match lvl with
  | max lhs rhs => if ptrEq lhs newLhs && ptrEq rhs newRhs then simpLevelMax' newLhs newRhs lvl else mkLevelMax' newLhs newRhs
  | _           => panic! "max level expected"

@[implemented_by updateMax!Impl]
def updateMax! (lvl : Level) (newLhs : Level) (newRhs : Level) : Level :=
  match lvl with
  | max _ _ => mkLevelMax' newLhs newRhs
  | _       => panic! "max level expected"

@[inline] private unsafe def updateIMax!Impl (lvl : Level) (newLhs : Level) (newRhs : Level) : Level :=
  match lvl with
  | imax lhs rhs => if ptrEq lhs newLhs && ptrEq rhs newRhs then simpLevelIMax' newLhs newRhs lvl else mkLevelIMax' newLhs newRhs
  | _            => panic! "imax level expected"

@[implemented_by updateIMax!Impl]
def updateIMax! (lvl : Level) (newLhs : Level) (newRhs : Level) : Level :=
  match lvl with
  | imax _ _ => mkLevelIMax' newLhs newRhs
  | _        => panic! "imax level expected"

def mkNaryMax : List Level → Level
  | []    => levelZero
  | [u]   => u
  | u::us => mkLevelMax' u (mkNaryMax us)

@[specialize] def substParams (u : Level) (s : Name → Option Level) : Level :=
  go u
where
  go (u : Level) : Level :=
    match u with
    | .zero       => u
    | .succ v     => if u.hasParam then u.updateSucc! (go v) else u
    | .max v₁ v₂  => if u.hasParam then u.updateMax! (go v₁) (go v₂) else u
    | .imax v₁ v₂ => if u.hasParam then u.updateIMax! (go v₁) (go v₂) else u
    | .param n    => match s n with
      | some u' => u'
      | none    => u
    | u => u

def getParamSubst : List Name → List Level → Name → Option Level
  | p::ps, u::us, p' => if p == p' then some u else getParamSubst ps us p'
  | _,     _,     _  => none

def instantiateParams (u : Level) (paramNames : List Name) (vs : List Level) : Level :=
  u.substParams (getParamSubst paramNames vs)

def geq (u v : Level) : Bool :=
  go u.normalize v.normalize
where
  go (u v : Level) : Bool :=
    u == v ||
    let k := fun () =>
      match v with
      | imax v₁ v₂ => go u v₁ && go u v₂
      | _          =>
        let v' := v.getLevelOffset
        (u.getLevelOffset == v' || v'.isZero)
        && u.getOffset ≥ v.getOffset
    match u, v with
    | _,          zero      => true
    | u,          max v₁ v₂ => go u v₁ && go u v₂
    | max u₁ u₂,  v         => go u₁ v || go u₂ v || k ()
    | imax _  u₂, v         => go u₂ v
    | succ u,     succ v    => go u v
    | _,          _         => k ()
  termination_by (u, v)

end Level

abbrev LevelMap (α : Type)  := Std.HashMap Level α
abbrev PersistentLevelMap (α : Type) := PHashMap Level α
abbrev LevelSet := Std.HashSet Level
abbrev PersistentLevelSet := PHashSet Level
abbrev PLevelSet := PersistentLevelSet

def Level.collectMVars (u : Level) (s : LMVarIdSet := {}) : LMVarIdSet :=
  match u with
  | succ v   => collectMVars v s
  | max u v  => collectMVars u (collectMVars v s)
  | imax u v => collectMVars u (collectMVars v s)
  | mvar n   => s.insert n
  | _        => s

def Level.find? (u : Level) (p : Level → Bool) : Option Level :=
  let rec visit (u : Level) : Option Level :=
    if p u then
      return u
    else match u with
      | succ v   => visit v
      | max u v  => visit u <|> visit v
      | imax u v => visit u <|> visit v
      | _          => failure
  visit u

def Level.any (u : Level) (p : Level → Bool) : Bool :=
  u.find? p |>.isSome

end Lean

abbrev Nat.toLevel (n : Nat) : Lean.Level :=
  Lean.Level.ofNat n