module Quox.Untyped.Erase

import Quox.Definition as Q
import Quox.Pretty
import Quox.Syntax.Term.Base as Q
import Quox.Syntax.Term.Subst
import Quox.Typing
import Quox.Untyped.Syntax as U
import Quox.Whnf

import Quox.EffExtra
import Data.List1
import Data.Singleton
import Data.SnocVect
import Language.Reflection

%default total
%language ElabReflection

%hide TT.Name
%hide AppView.(.head)


public export
data IsErased = Erased | Kept

public export
isErased : Qty -> IsErased
isErased Zero = Erased
isErased One  = Kept
isErased Any  = Kept


public export
ErasureContext : Nat -> Nat -> Type
ErasureContext = TyContext


public export
TypeError : Type
TypeError = Typing.Error.Error
%hide Typing.Error.Error

public export
data Error =
  CompileTimeOnly (ErasureContext d n) (Q.Term d n)
| WrapTypeError   TypeError
| Postulate       Loc Name
| WhileErasing    Name Q.Definition Error
| MainIsErased    Loc Name
%name Error err

private %inline
notInScope : Loc -> Name -> Error
notInScope = WrapTypeError .: NotInScope

export
Located Error where
  (CompileTimeOnly _ s).loc  = s.loc
  (WrapTypeError err).loc    = err.loc
  (Postulate loc _).loc      = loc
  (WhileErasing _ def e).loc = e.loc `or` def.loc
  (MainIsErased loc _).loc   = loc


parameters {opts : LayoutOpts} (showContext : Bool)
  export
  prettyErrorNoLoc : Error -> Eff Pretty (Doc opts)
  prettyErrorNoLoc (CompileTimeOnly ctx s) =
    inTContext showContext ctx $
    sep ["the term", !(prettyTerm ctx.dnames ctx.tnames s),
         "only exists at compile time"]
  prettyErrorNoLoc (WrapTypeError err) =
    prettyErrorNoLoc showContext err
  prettyErrorNoLoc (Postulate _ x) =
    pure $ sep [!(prettyFree x), "is a postulate with no definition"]
  prettyErrorNoLoc (WhileErasing x def err) = pure $
    vsep [hsep ["while erasing the definition", !(prettyFree x)],
          !(prettyErrorNoLoc err)]
  prettyErrorNoLoc (MainIsErased _ x) =
    pure $ hsep [!(prettyFree x), "is marked #[main] but is erased"]

  export
  prettyError : Error -> Eff Pretty (Doc opts)
  prettyError err = sep <$> sequence
    [prettyLoc err.loc, indentD =<< prettyErrorNoLoc err]


public export
Erase : List (Type -> Type)
Erase = [Except Error, DefsReader, NameGen]

export
liftWhnf : Eff Whnf a -> Eff Erase a
liftWhnf act = lift $ wrapErr WrapTypeError act

export covering
computeElimType : ErasureContext d n -> SQty -> Elim d n -> Eff Erase (Term d n)
computeElimType ctx sg e = do
  defs <- askAt DEFS
  let ctx = toWhnfContext ctx
  liftWhnf $ do
    Element e _ <- whnf defs ctx sg e
    computeElimType defs ctx sg e


private %macro
wrapExpect : TTImp ->
             Elab (TyContext d n -> Loc -> Term d n -> Eff Erase a)
wrapExpect f_ = do
  f <- check `(\x => ~(f_) x)
  pure $ \ctx, loc, s => liftWhnf $ f !(askAt DEFS) ctx SZero loc s


public export
record EraseElimResult d n where
  constructor EraRes
  type : Lazy (Q.Term d n)
  term : U.Term n


-- "Ψ | Γ | Σ ⊢ s ⤋ s' ⇐ A" for `s' <- eraseTerm (Ψ,Γ,Σ) A s`
--
-- in the below comments, Ψ, Γ, Σ are implicit and
-- only their extensions are written
export covering
eraseTerm : ErasureContext d n ->
            (ty, tm : Q.Term d n) -> Eff Erase (U.Term n)


-- "Ψ | Γ | Σ ⊢ e ⤋ e' ⇒ A" for `EraRes A e' <- eraseElim (Ψ,Γ,Σ) e`
export covering
eraseElim : ErasureContext d n -> (tm : Q.Elim d n) ->
            Eff Erase (EraseElimResult d n)

eraseTerm ctx _ s@(TYPE {}) =
  throw $ CompileTimeOnly ctx s

eraseTerm ctx _ s@(IOState {}) =
  throw $ CompileTimeOnly ctx s

eraseTerm ctx _ s@(Pi {}) =
  throw $ CompileTimeOnly ctx s

--      x : A | 0.x ⊢ s ⤋ s' ⇐ B
-- -------------------------------------
--  (λ x ⇒ s) ⤋ s'[⌷/x] ⇐ 0.(x : A) → B
--
--    x : A | π.x ⊢ s ⤋ s' ⇐ B   π ≠ 0
-- ----------------------------------------
--  (λ x ⇒ s) ⤋ (λ x ⇒ s') ⇐ π.(x : A) → B
eraseTerm ctx ty (Lam body loc) = do
  let x = body.name
  (qty, arg, res) <- wrapExpect `(expectPi) ctx loc ty
  body            <- eraseTerm (extendTy qty x arg ctx) res.term body.term
  pure $ case isErased qty of
    Kept   => U.Lam x body loc
    Erased => sub1 (Erased loc) body

eraseTerm ctx _ s@(Sig {}) =
  throw $ CompileTimeOnly ctx s

--   s ⤋ s' ⇐ A    t ⤋ t' ⇐ B[s/x]
-- ---------------------------------
--  (s, t) ⤋ (s', t') ⇐ (x : A) × B
eraseTerm ctx ty (Pair fst snd loc) = do
  (a, b) <- wrapExpect `(expectSig) ctx loc ty
  let b = sub1 b (Ann fst a a.loc)
  fst <- eraseTerm ctx a fst
  snd <- eraseTerm ctx b snd
  pure $ Pair fst snd loc

eraseTerm ctx _ s@(Enum {}) =
  throw $ CompileTimeOnly ctx s

-- '𝐚 ⤋ '𝐚 ⇐ {⋯}
eraseTerm ctx _ (Tag tag loc) =
  pure $ Tag tag loc

eraseTerm ctx ty s@(Eq {}) =
  throw $ CompileTimeOnly ctx s

--        𝑖 ⊢ s ⤋ s' ⇐ A
-- ---------------------------------
--  (δ 𝑖 ⇒ s) ⤋ s' ⇐ Eq (𝑖 ⇒ A) l r
eraseTerm ctx ty (DLam body loc) = do
  a <- fst <$> wrapExpect `(expectEq) ctx loc ty
  eraseTerm (extendDim body.name ctx) a.term body.term

eraseTerm ctx _ s@(NAT {}) =
  throw $ CompileTimeOnly ctx s

-- n ⤋ n ⇐ ℕ
eraseTerm _ _ (Nat n loc) =
  pure $ Nat n loc

--       s ⤋ s' ⇐ ℕ
-- -----------------------
--  succ s ⤋ succ s' ⇐ ℕ
eraseTerm ctx ty (Succ p loc) = do
  p <- eraseTerm ctx ty p
  pure $ Succ p loc

eraseTerm ctx ty s@(STRING {}) =
  throw $ CompileTimeOnly ctx s

-- s ⤋ s ⇐ String
eraseTerm _ _ (Str s loc) =
  pure $ Str s loc

eraseTerm ctx ty s@(BOX {}) =
  throw $ CompileTimeOnly ctx s

-- [s] ⤋ ⌷ ⇐ [0.A]
--
--  π ≠ 0   s ⤋ s' ⇐ A
-- --------------------
--   [s] ⤋ s' ⇐ [π.A]
eraseTerm ctx ty (Box val loc) = do
  (qty, a) <- wrapExpect `(expectBOX) ctx loc ty
  case isErased qty of
    Erased => pure $ Erased loc
    Kept   => eraseTerm ctx a val

--  e ⤋ e' ⇒ B
-- ------------
--  e ⤋ e' ⇐ A
eraseTerm ctx ty (E e) =
  term <$> eraseElim ctx e

eraseTerm ctx ty (CloT (Sub term th)) =
  eraseTerm ctx ty $ pushSubstsWith' id th term

eraseTerm ctx ty (DCloT (Sub term th)) =
  eraseTerm ctx ty $ pushSubstsWith' th id term

--  defω x : A = s
-- ----------------
--    x ⤋ x ⇒ A
eraseElim ctx e@(F x u loc) = do
  Just def <- asksAt DEFS $ lookup x
    | Nothing => throw $ notInScope loc x
  case isErased def.qty.qty of
    Erased => throw $ CompileTimeOnly ctx $ E e
    Kept   => pure $ EraRes (def.typeWith ctx.dimLen ctx.termLen) $ F x loc

--  π.x ∈ Σ   π ≠ 0
-- -----------------
--     x ⤋ x ⇒ A
eraseElim ctx e@(B i loc) = do
  case isErased $ ctx.qtys !!! i of
    Erased => throw $ CompileTimeOnly ctx $ E e
    Kept   => pure $ EraRes (ctx.tctx !! i) $ B i loc

--  f ⤋ f' ⇒ π.(x : A) → B   s ⤋ s' ⇒ A   π ≠ 0
-- ---------------------------------------------
--             f s ⤋ f' s' ⇒ B[s/x]
--
--  f ⤋ f' ⇒ 0.(x : A) → B
-- -------------------------
--    f s ⤋ f' ⇒ B[s/x]
eraseElim ctx (App fun arg loc) = do
  efun              <- eraseElim ctx fun
  (qty, targ, tres) <- wrapExpect `(expectPi) ctx loc efun.type
  let ty = sub1 tres (Ann arg targ arg.loc)
  case isErased qty of
    Erased => pure $ EraRes ty efun.term
    Kept   => do arg <- eraseTerm ctx targ arg
                 pure $ EraRes ty $ App efun.term arg loc

--                        e ⇒ (x : A) × B
--    x : A, y : B | ρ.x, ρ.y ⊢ s ⤋ s' ⇐ R[((x,y) ∷ (x : A) × B)/z]
-- -------------------------------------------------------------------
--   (case0 e return z ⇒ R of {(x, y) ⇒ s}) ⤋ s'[⌷/x, ⌷/y] ⇒ R[e/z]
--
--                     e ⤋ e' ⇒ (x : A) × B   ρ ≠ 0
--    x : A, y : B | ρ.x, ρ.y ⊢ s ⤋ s' ⇐ R[((x,y) ∷ (x : A) × B)/z]
-- ----------------------------------------------------------------------------
--   (caseρ e return z ⇒ R of {(x, y) ⇒ s}) ⤋
--     ⤋
--   let xy = e' in let x = fst xy in let y = snd xy in s' ⇒ R[e/z]
eraseElim ctx (CasePair qty pair ret body loc) = do
  let [< x, y] = body.names
  case isErased qty of
    Kept => do
      EraRes ety eterm <- eraseElim ctx pair
      let ty = sub1 (ret // shift 2) $
               Ann (Pair (BVT 0 loc) (BVT 1 loc) loc) (weakT 2 ety) loc
      (tfst, tsnd) <- wrapExpect `(expectSig) ctx loc ety
      let ctx' = extendTyN [< (qty, x, tfst), (qty, y, tsnd.term)] ctx
      body' <- eraseTerm ctx' ty body.term
      p <- mnb "p" loc
      pure $ EraRes (sub1 ret pair) $
        Let p eterm
          (Let x (Fst (B VZ loc) loc)
            (Let y (Snd (B (VS VZ) loc) loc)
              (body' // (B VZ loc ::: B (VS VZ) loc ::: shift 3))
              loc) loc) loc
    Erased => do
      ety <- computeElimType ctx SOne pair
      let ty = sub1 (ret // shift 2) $
               Ann (Pair (BVT 0 loc) (BVT 1 loc) loc) (weakT 2 ety) loc
      (tfst, tsnd) <- wrapExpect `(expectSig) ctx loc ety
      let ctx' = extendTyN0 [< (x, tfst), (y, tsnd.term)] ctx
      body' <- eraseTerm ctx' ty body.term
      pure $ EraRes (sub1 ret pair) $ subN [< Erased loc, Erased loc] body'

--  e ⤋ e' ⇒ (x : A) × B
-- ----------------------
--   fst e ⤋ fst e' ⇒ A
eraseElim ctx (Fst pair loc) = do
  epair      <- eraseElim ctx pair
  a          <- fst <$> wrapExpect `(expectSig) ctx loc epair.type
  pure $ EraRes a $ Fst epair.term loc

--     e ⤋ e' ⇒ (x : A) × B
-- -----------------------------
--  snd e ⤋ snd e' ⇒ B[fst e/x]
eraseElim ctx (Snd pair loc) = do
  epair      <- eraseElim ctx pair
  b          <- snd <$> wrapExpect `(expectSig) ctx loc epair.type
  pure $ EraRes (sub1 b (Fst pair loc)) $ Snd epair.term loc

-- caseρ e return z ⇒ R of {} ⤋ absurd ⇒ R[e/z]
--
--               s ⤋ s' ⇐ R[𝐚∷{𝐚}/z]
-- -----------------------------------------------
--  case0 e return z ⇒ R of {𝐚 ⇒ s} ⤋ s' ⇒ R[e/z]
--
--         e ⤋ e' ⇒ A   sᵢ ⤋ s'ᵢ ⇐ R[𝐚ᵢ/z]   ρ ≠ 0   i ≠ 0
-- -------------------------------------------------------------------
--  caseρ e return z ⇒ R of {𝐚ᵢ ⇒ sᵢ} ⤋ case e of {𝐚ᵢ ⇒ s'ᵢ} ⇒ R[e/z]
eraseElim ctx e@(CaseEnum qty tag ret arms loc) = do
  let ty = sub1 ret tag
  case isErased qty of
    Erased => case SortedMap.toList arms of
      []         => pure $ EraRes ty $ Absurd loc
      [(t, rhs)] => do
        let ty' = sub1 ret (Ann (Tag t loc) (enum [t] loc) loc)
        rhs' <- eraseTerm ctx ty' rhs
        pure $ EraRes ty rhs'
      _ => throw $ CompileTimeOnly ctx $ E e
    Kept => case List1.fromList $ SortedMap.toList arms of
      Nothing   => pure $ EraRes ty $ Absurd loc
      Just arms => do
        etag <- eraseElim ctx tag
        arms <- for arms $ \(t, rhs) => do
          let ty' = sub1 ret (Ann (Tag t loc) etag.type loc)
          rhs' <- eraseTerm ctx ty' rhs
          pure (t, rhs')
        pure $ EraRes ty $ CaseEnum etag.term arms loc

--        n ⤋ n' ⇒ ℕ   z ⤋ z' ⇐ R[zero∷ℕ/z]   ς ≠ 0
--  m : ℕ, ih : R[m/z] | ρ.m, ς.ih ⊢ s ⤋ s' ⇐ R[(succ m)∷ℕ/z]
-- -----------------------------------------------------------
--      caseρ n return z ⇒ R of {0 ⇒ z; succ m, ς.ih ⇒ s}
--        ⤋
--      case n' of {0 ⇒ z'; succ m, ih ⇒ s'} ⇒ R[n/z]
--
--             n ⤋ n' ⇒ ℕ   z ⤋ z' ⇐ R[zero∷ℕ/z]
--  m : ℕ, ih : R[m/z] | ρ.m, 0.ih ⊢ s ⤋ s' ⇐ R[(succ m)∷ℕ/z]
-- -----------------------------------------------------------
--      caseρ n return z ⇒ R of {0 ⇒ z; succ m, 0.ih ⇒ s}
--        ⤋
--      case n' of {0 ⇒ z'; succ m ⇒ s'[⌷/ih]} ⇒ R[n/z]
eraseElim ctx (CaseNat qty qtyIH nat ret zero succ loc) = do
  let ty = sub1 ret nat
  enat <- eraseElim ctx nat
  zero <- eraseTerm ctx (sub1 ret (Ann (Zero loc) (NAT loc) loc)) zero
  let [< p, ih] = succ.names
  succ' <- eraseTerm
    (extendTyN [< (qty, p, NAT loc),
                  (qtyIH, ih, sub1 (ret // shift 1) (BV 0 loc))] ctx)
    (sub1 (ret // shift 2) (Ann (Succ (BVT 1 loc) loc) (NAT loc) loc))
    succ.term
  let succ = case isErased qtyIH of
               Kept   => NSRec    p ih succ'
               Erased => NSNonrec p    (sub1 (Erased loc) succ')
  pure $ EraRes ty $ CaseNat enat.term zero succ loc

--  b ⤋ b' ⇒ [π.A]   πρ ≠ 0   x : A | πρ.x ⊢ s ⤋ s' ⇐ R[[x]∷[π.A]/z]
-- ------------------------------------------------------------------
--   caseρ b return z ⇒ R of {[x] ⇒ s} ⤋ let x = b' in s' ⇒ R[b/z]
--
--   b ⇒ [π.A]   x : A | 0.x ⊢ s ⤋ s' ⇐ R[[x]∷[0.A]/z]   πρ = 0
-- -------------------------------------------------------------
--  caseρ b return z ⇒ R of {[x] ⇒ s} ⤋ s'[⌷/x] ⇒ R[b/z]
eraseElim ctx (CaseBox qty box ret body loc) = do
  tbox <- computeElimType ctx SOne box -- [fixme] is there any way to avoid this?
  (pi, tinner) <- wrapExpect `(expectBOX) ctx loc tbox
  let ctx' = extendTy (pi * qty) body.name tinner ctx
      bty  = sub1 (ret // shift 1) $
             Ann (Box (BVT 0 loc) loc) (weakT 1 tbox) loc
  case isErased $ qty * pi of
    Kept => do
      ebox  <- eraseElim ctx box
      ebody <- eraseTerm ctx' bty body.term
      pure $ EraRes (sub1 ret box) $ Let body.name ebox.term ebody loc
    Erased => do
      body' <- eraseTerm ctx' bty body.term
      pure $ EraRes (sub1 ret box) $ body' // one (Erased loc)

--    f ⤋ f' ⇒ Eq (𝑖 ⇒ A) l r
-- ------------------------------
--      f @r ⤋ f' ⇒ A‹r/𝑖›
eraseElim ctx (DApp fun arg loc) = do
  efun <- eraseElim ctx fun
  a    <- fst <$> wrapExpect `(expectEq) ctx loc efun.type
  pure $ EraRes (dsub1 a arg) efun.term

--    s ⤋ s' ⇐ A
-- ----------------
--  s ∷ A ⤋ s' ⇒ A
eraseElim ctx (Ann tm ty loc) =
  EraRes ty <$> eraseTerm ctx ty tm

--        s ⤋ s' ⇐ A‹p/𝑖›
-- -----------------------------------
--  coe (𝑖 ⇒ A) @p @q s ⤋ s' ⇒ A‹q/𝑖›
eraseElim ctx (Coe ty p q val loc) = do
  val <- eraseTerm ctx (dsub1 ty p) val
  pure $ EraRes (dsub1 ty q) val

--          s ⤋ s' ⇐ A
-- --------------------------------
--  comp A @p @q s @r {⋯} ⤋ s' ⇒ A
eraseElim ctx (Comp ty p q val r zero one loc) =
  EraRes ty <$> eraseTerm ctx ty val

eraseElim ctx t@(TypeCase ty ret arms def loc) =
  throw $ CompileTimeOnly ctx $ E t

eraseElim ctx (CloE (Sub term th)) =
  eraseElim ctx $ pushSubstsWith' id th term

eraseElim ctx (DCloE (Sub term th)) =
  eraseElim ctx $ pushSubstsWith' th id term


export
uses : Var n -> Term n -> Nat
uses i (F {}) = 0
uses i (B j _) = if i == j then 1 else 0
uses i (Lam x body _) = uses (VS i) body
uses i (App fun arg _) = uses i fun + uses i arg
uses i (Pair fst snd _) = uses i fst + uses i snd
uses i (Fst pair _) = uses i pair
uses i (Snd pair _) = uses i pair
uses i (Tag tag _) = 0
uses i (CaseEnum tag cases _) =
  uses i tag + foldl max 0 (map (assert_total uses i . snd) cases)
uses i (Absurd {}) = 0
uses i (Nat {}) = 0
uses i (Succ nat _) = uses i nat
uses i (CaseNat nat zer suc _) = uses i nat + max (uses i zer) (uses' suc)
  where uses' : CaseNatSuc n -> Nat
        uses' (NSRec  _ _ s) = uses (VS (VS i)) s
        uses' (NSNonrec _ s) = uses (VS i) s
uses i (Str {}) = 0
uses i (Let x rhs body _) = uses i rhs + uses (VS i) body
uses i (Erased {}) = 0

export
inlineable : U.Term n -> Bool
inlineable (F {})      = True
inlineable (B {})      = True
inlineable (Tag {})    = True
inlineable (Nat {})    = True
inlineable (Str {})    = True
inlineable (Absurd {}) = True
inlineable (Erased {}) = True
inlineable _           = False

export
droppable : U.Term n -> Bool
droppable (F {})      = True
droppable (B {})      = True
droppable (Fst e _)   = droppable e
droppable (Snd e _)   = droppable e
droppable (Tag {})    = True
droppable (Nat {})    = True
droppable (Str {})    = True
droppable (Absurd {}) = True
droppable (Erased {}) = True
droppable _           = False

export
trimLets : U.Term n -> U.Term n
trimLets (F x loc) = F x loc
trimLets (B i loc) = B i loc
trimLets (Lam x body loc) = Lam x (trimLets body) loc
trimLets (App fun arg loc) = App (trimLets fun) (trimLets arg) loc
trimLets (Pair fst snd loc) = Pair (trimLets fst) (trimLets snd) loc
trimLets (Fst pair loc) = Fst (trimLets pair) loc
trimLets (Snd pair loc) = Snd (trimLets pair) loc
trimLets (Tag tag loc) = Tag tag loc
trimLets (CaseEnum tag cases loc) =
  let tag   = trimLets tag
      cases = map (map $ \c => trimLets $ assert_smaller cases c) cases in
  if droppable tag && length cases == 1
     then snd cases.head
     else CaseEnum tag cases loc
trimLets (Absurd loc) = Absurd loc
trimLets (Nat n loc) = Nat n loc
trimLets (Succ nat loc) = Succ (trimLets nat) loc
trimLets (CaseNat nat zer suc loc) =
  CaseNat (trimLets nat) (trimLets zer) (trimLets' suc) loc
  where trimLets' : CaseNatSuc n -> CaseNatSuc n
        trimLets' (NSRec x ih s) = NSRec    x ih $ trimLets s
        trimLets' (NSNonrec x s) = NSNonrec x    $ trimLets s
trimLets (Str s loc) = Str s loc
trimLets (Let x rhs body loc) =
  let rhs'  = trimLets rhs
      body' = trimLets body
      uses  = uses VZ body in
  if inlineable rhs' || uses == 1 || (droppable rhs' && uses == 0)
     then sub1 rhs' body'
     else Let x rhs' body' loc
trimLets (Erased loc) = Erased loc


export covering
eraseDef : Name -> Q.Definition -> Eff Erase U.Definition
eraseDef name def@(MkDef qty type body scheme isMain loc) =
  wrapErr (WhileErasing name def) $
  case isErased qty.qty of
    Erased => do
      when isMain $ throw $ MainIsErased loc name
      pure ErasedDef
    Kept   =>
      case scheme of
        Just str => pure $ SchemeDef isMain str
        Nothing  => case body of
          Postulate     => throw $ Postulate loc name
          Concrete body => KeptDef isMain . trimLets <$>
                           eraseTerm empty type body