/-
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 Level

partial def replace (f? : Level → Option Level) (u : Level) : Level :=
  match f? u with
  | some v => v
  | none   => match u with
    | max v₁ v₂  _ => mkLevelMax' (replace f? v₁) (replace f? v₂)
    | imax v₁ v₂ _ => mkLevelIMax' (replace f? v₁) (replace f? v₂)
    | succ v _     => mkLevelSucc (replace f? v)
    | _            => u

end Level

namespace Expr

namespace ReplaceLevelImpl

abbrev cacheSize : USize := 8192

structure State where
  keys    : Array Expr -- Remark: our "unsafe" implementation relies on the fact that `()` is not a valid Expr
  results : Array Expr

abbrev ReplaceM := StateM State

@[inline] unsafe def cache (i : USize) (key : Expr) (result : Expr) : ReplaceM Expr := do
  modify fun s => { keys := s.keys.uset i key lcProof, results := s.results.uset i result lcProof };
  pure result

@[specialize] unsafe def replaceUnsafeM (f? : Level → Option Level) (size : USize) (e : Expr) : ReplaceM Expr := do
  let rec visit (e : Expr) := do
    let c ← get
    let h := ptrAddrUnsafe e
    let i := h % size
    if ptrAddrUnsafe (c.keys.uget i lcProof) == h then
      pure <| c.results.uget i lcProof
    else match e with
        | Expr.forallE _ d b _   => cache i e <| e.updateForallE! (← visit d) (← visit b)
        | Expr.lam _ d b _       => cache i e <| e.updateLambdaE! (← visit d) (← visit b)
        | Expr.mdata _ b _       => cache i e <| e.updateMData! (← visit b)
        | Expr.letE _ t v b _    => cache i e <| e.updateLet! (← visit t) (← visit v) (← visit b)
        | Expr.app f a _         => cache i e <| e.updateApp! (← visit f) (← visit a)
        | Expr.proj _ _ b _      => cache i e <| e.updateProj! (← visit b)
        | Expr.sort u _          => cache i e <| e.updateSort! (u.replace f?)
        | Expr.const n us _      => cache i e <| e.updateConst! (us.map (Level.replace f?))
        | e                      => pure e
  visit e

unsafe def initCache : State :=
  { keys    := mkArray cacheSize.toNat (cast lcProof ()), -- `()` is not a valid `Expr`
    results := mkArray cacheSize.toNat arbitrary }

@[inline] unsafe def replaceUnsafe (f? : Level → Option Level) (e : Expr) : Expr :=
  (replaceUnsafeM f? cacheSize e).run' initCache

end ReplaceLevelImpl

@[implementedBy ReplaceLevelImpl.replaceUnsafe]
partial def replaceLevel (f? : Level → Option Level) : Expr → Expr
  | e@(Expr.forallE _ d b _)   => let d := replaceLevel f? d; let b := replaceLevel f? b; e.updateForallE! d b
  | e@(Expr.lam _ d b _)       => let d := replaceLevel f? d; let b := replaceLevel f? b; e.updateLambdaE! d b
  | e@(Expr.mdata _ b _)       => let b := replaceLevel f? b; e.updateMData! b
  | e@(Expr.letE _ t v b _)    => let t := replaceLevel f? t; let v := replaceLevel f? v; let b := replaceLevel f? b; e.updateLet! t v b
  | e@(Expr.app f a _)         => let f := replaceLevel f? f; let a := replaceLevel f? a; e.updateApp! f a
  | e@(Expr.proj _ _ b _)      => let b := replaceLevel f? b; e.updateProj! b
  | e@(Expr.sort u _)          => e.updateSort! (u.replace f?)
  | e@(Expr.const n us _)      => e.updateConst! (us.map (Level.replace f?))
  | e                          => e

end Expr
end Lean