module Quox.Typechecker

import public Quox.Typing
import public Quox.Equal
import Quox.Displace

import Data.List
import Data.SnocVect
import Data.List1
import Quox.EffExtra

%default total


public export
0 TC : List (Type -> Type)
TC = [ErrorEff, DefsReader, NameGen]


parameters (loc : Loc)
  export
  popQs : Has ErrorEff fs => QContext s -> QOutput (s + n) ->
          Eff fs (QOutput n)
  popQs [<]         qout         = pure qout
  popQs (pis :< pi) (qout :< rh) = do expectCompatQ loc rh pi; popQs pis qout

  export %inline
  popQ : Has ErrorEff fs => Qty -> QOutput (S n) -> Eff fs (QOutput n)
  popQ pi = popQs [< pi]

export
lubs1 : List1 (QOutput n) -> QOutput n
lubs1 ([<] ::: _) = [<]
lubs1 ((qs :< p) ::: pqs) =
  let (qss, ps) = unzip $ map unsnoc pqs in
  lubs1 (qs ::: qss) :< foldl lub p ps

export
lubs : TyContext d n -> List (QOutput n) -> QOutput n
lubs ctx []        = zeroFor ctx
lubs ctx (x :: xs) = lubs1 $ x ::: xs


export
typecaseTel : (k : TyConKind) -> BContext (arity k) -> Universe ->
              CtxExtension d n (arity k + n)
typecaseTel k xs u = case k of
  KTYPE => [<]
  -- A : ★ᵤ, B : 0.A → ★ᵤ
  KPi =>
    let [< a, b] = xs in
    [< (Zero, a, TYPE u a.loc),
       (Zero, b, Arr Zero (BVT 0 b.loc) (TYPE u b.loc) b.loc)]
  KSig =>
    let [< a, b] = xs in
    [< (Zero, a, TYPE u a.loc),
       (Zero, b, Arr Zero (BVT 0 b.loc) (TYPE u b.loc) b.loc)]
  KEnum => [<]
  -- A₀ : ★ᵤ, A₁ : ★ᵤ, A : (A₀ ≡ A₁ : ★ᵤ), L : A₀, R : A₀
  KEq =>
    let [< a0, a1, a, l, r] = xs in
    [< (Zero, a0, TYPE u a0.loc),
       (Zero, a1, TYPE u a1.loc),
       (Zero, a,  Eq0 (TYPE u a.loc) (BVT 1 a.loc) (BVT 0 a.loc) a.loc),
       (Zero, l,  BVT 2 l.loc),
       (Zero, r,  BVT 2 r.loc)]
  KNat => [<]
  -- A : ★ᵤ
  KBOX => let [< a] = xs in [< (Zero, a, TYPE u a.loc)]


mutual
  ||| "Ψ | Γ ⊢ σ · s ⇐ A ⊳ Σ"
  |||
  ||| `check ctx sg subj ty` checks that in the context `ctx`, the term
  ||| `subj` has the type `ty`, with quantity `sg`. if so, returns the
  ||| quantities of all bound variables that it used.
  |||
  ||| if the dimension context is inconsistent, then return `Nothing`, without
  ||| doing any further work.
  export covering %inline
  check : (ctx : TyContext d n) -> SQty -> Term d n -> Term d n ->
          Eff TC (CheckResult ctx.dctx n)
  check ctx sg subj ty = ifConsistent ctx.dctx $ checkC ctx sg subj ty

  ||| "Ψ | Γ ⊢₀ s ⇐ A"
  |||
  ||| `check0 ctx subj ty` checks a term (as `check`) in a zero context.
  export covering %inline
  check0 : TyContext d n -> Term d n -> Term d n -> Eff TC ()
  check0 ctx tm ty = ignore $ check ctx SZero tm ty
    -- the output will always be 𝟎 because the subject quantity is 0

  ||| `check`, assuming the dimension context is consistent
  export covering %inline
  checkC : (ctx : TyContext d n) -> SQty -> Term d n -> Term d n ->
           Eff TC (CheckResult' n)
  checkC ctx sg subj ty =
    wrapErr (WhileChecking ctx sg subj ty) $
    checkCNoWrap ctx sg subj ty

  export covering %inline
  checkCNoWrap : (ctx : TyContext d n) -> SQty -> Term d n -> Term d n ->
                 Eff TC (CheckResult' n)
  checkCNoWrap ctx sg subj ty =
    let Element subj nc = pushSubsts subj in
    check' ctx sg subj ty

  ||| "Ψ | Γ ⊢₀ s ⇐ ★ᵢ"
  |||
  ||| `checkType ctx subj ty` checks a type (in a zero context). sometimes the
  ||| universe doesn't matter, only that a term is _a_ type, so it is optional.
  export covering %inline
  checkType : TyContext d n -> Term d n -> Maybe Universe -> Eff TC ()
  checkType ctx subj l = ignore $ ifConsistent ctx.dctx $ checkTypeC ctx subj l

  export covering %inline
  checkTypeC : TyContext d n -> Term d n -> Maybe Universe -> Eff TC ()
  checkTypeC ctx subj l =
    wrapErr (WhileCheckingTy ctx subj l) $ checkTypeNoWrap ctx subj l

  export covering %inline
  checkTypeNoWrap : TyContext d n -> Term d n -> Maybe Universe -> Eff TC ()
  checkTypeNoWrap ctx subj l =
    let Element subj nc = pushSubsts subj in
    checkType' ctx subj l

  ||| "Ψ | Γ ⊢ σ · e ⇒ A ⊳ Σ"
  |||
  ||| `infer ctx sg subj` infers the type of `subj` in the context `ctx`,
  ||| and returns its type and the bound variables it used.
  |||
  ||| if the dimension context is inconsistent, then return `Nothing`, without
  ||| doing any further work.
  export covering %inline
  infer : (ctx : TyContext d n) -> SQty -> Elim d n ->
          Eff TC (InferResult ctx.dctx d n)
  infer ctx sg subj = ifConsistent ctx.dctx $ inferC ctx sg subj

  ||| `infer`, assuming the dimension context is consistent
  export covering %inline
  inferC : (ctx : TyContext d n) -> SQty -> Elim d n ->
           Eff TC (InferResult' d n)
  inferC ctx sg subj =
    wrapErr (WhileInferring ctx sg subj) $
      let Element subj nc = pushSubsts subj in
      infer' ctx sg subj


  private covering
  toCheckType : TyContext d n -> SQty ->
                (subj : Term d n) -> (0 nc : NotClo subj) => Term d n ->
                Eff TC (CheckResult' n)
  toCheckType ctx sg t ty = do
    u <- expectTYPE !(askAt DEFS) ctx sg ty.loc ty
    expectEqualQ t.loc Zero sg.qty
    checkTypeNoWrap ctx t (Just u)
    pure $ zeroFor ctx

  private covering
  check' : TyContext d n -> SQty ->
           (subj : Term d n) -> (0 nc : NotClo subj) => Term d n ->
           Eff TC (CheckResult' n)

  check' ctx sg t@(TYPE {}) ty = toCheckType ctx sg t ty

  check' ctx sg t@(Pi {}) ty = toCheckType ctx sg t ty

  check' ctx sg (Lam body loc) ty = do
    (qty, arg, res) <- expectPi !(askAt DEFS) ctx SZero ty.loc ty
    -- if Ψ | Γ, x : A ⊢ σ · t ⇐ B ⊳ Σ, ρ·x
    -- with ρ ≤ σπ
    let qty' = sg.qty * qty
    qout <- checkC (extendTy qty' body.name arg ctx) sg body.term res.term
    -- then Ψ | Γ ⊢ σ · (λx ⇒ t) ⇐ (π·x : A) → B ⊳ Σ
    popQ loc qty' qout

  check' ctx sg t@(Sig {}) ty = toCheckType ctx sg t ty

  check' ctx sg (Pair fst snd loc) ty = do
    (tfst, tsnd) <- expectSig !(askAt DEFS) ctx SZero ty.loc ty
    -- if Ψ | Γ ⊢ σ · s ⇐ A ⊳ Σ₁
    qfst <- checkC ctx sg fst tfst
    let tsnd = sub1 tsnd (Ann fst tfst fst.loc)
    -- if Ψ | Γ ⊢ σ · t ⇐ B[s] ⊳ Σ₂
    qsnd <- checkC ctx sg snd tsnd
    -- then Ψ | Γ ⊢ σ · (s, t) ⇐ (x : A) × B ⊳ Σ₁ + Σ₂
    pure $ qfst + qsnd

  check' ctx sg t@(Enum {}) ty = toCheckType ctx sg t ty

  check' ctx sg (Tag t loc) ty = do
    tags <- expectEnum !(askAt DEFS) ctx SZero ty.loc ty
    -- if t ∈ ts
    unless (t `elem` tags) $ throw $ TagNotIn loc t tags
    -- then Ψ | Γ ⊢ σ · t ⇐ {ts} ⊳ 𝟎
    pure $ zeroFor ctx

  check' ctx sg t@(Eq {}) ty = toCheckType ctx sg t ty

  check' ctx sg (DLam body loc) ty = do
    (ty, l, r) <- expectEq !(askAt DEFS) ctx SZero ty.loc ty
    let ctx' = extendDim body.name ctx
        ty   = ty.term
        body = body.term
    -- if Ψ, i | Γ ⊢ σ · t ⇐ A ⊳ Σ
    qout <- checkC ctx' sg body ty
    -- if Ψ, i, i = 0 | Γ ⊢ t = l : A
    let ctx0 = eqDim (B VZ loc) (K Zero loc) ctx'
    lift $ equal loc ctx0 sg ty body $ dweakT 1 l
    -- if Ψ, i, i = 1 | Γ ⊢ t = r : A
    let ctx1 = eqDim (B VZ loc) (K One loc) ctx'
    lift $ equal loc ctx1 sg ty body $ dweakT 1 r
    -- then Ψ | Γ ⊢ σ · (δ i ⇒ t) ⇐ Eq [i ⇒ A] l r ⊳ Σ
    pure qout

  check' ctx sg t@(Nat {}) ty = toCheckType ctx sg t ty

  check' ctx sg (Zero {}) ty = do
    expectNat !(askAt DEFS) ctx SZero ty.loc ty
    pure $ zeroFor ctx

  check' ctx sg (Succ n {}) ty = do
    expectNat !(askAt DEFS) ctx SZero ty.loc ty
    checkC ctx sg n ty

  check' ctx sg t@(BOX {}) ty = toCheckType ctx sg t ty

  check' ctx sg (Box val loc) ty = do
    (q, ty) <- expectBOX !(askAt DEFS) ctx SZero ty.loc ty
    -- if Ψ | Γ ⊢ σ · s ⇐ A ⊳ Σ
    valout <- checkC ctx sg val ty
    -- then Ψ | Γ ⊢ σ · [s] ⇐ [π.A] ⊳ πΣ
    pure $ q * valout

  check' ctx sg (E e) ty = do
    -- if Ψ | Γ ⊢ σ · e ⇒ A' ⊳ Σ
    infres <- inferC ctx sg e
    -- if Ψ | Γ ⊢ A' <: A
    lift $ subtype e.loc ctx infres.type ty
    -- then Ψ | Γ ⊢ σ · e ⇐ A ⊳ Σ
    pure infres.qout

  private covering
  checkType' : TyContext d n ->
               (subj : Term d n) -> (0 nc : NotClo subj) =>
               Maybe Universe -> Eff TC ()

  checkType' ctx (TYPE k loc) u = do
    -- if 𝓀 < ℓ then Ψ | Γ ⊢₀ Type 𝓀 ⇐ Type ℓ
    case u of
      Just l  => unless (k < l) $ throw $ BadUniverse loc k l
      Nothing => pure ()

  checkType' ctx (Pi qty arg res _) u = do
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    checkTypeC ctx arg u
    -- if Ψ | Γ, x : A ⊢₀ B ⇐ Type ℓ
    checkTypeScope ctx arg res u
    -- then Ψ | Γ ⊢₀ (π·x : A) → B ⇐ Type ℓ

  checkType' ctx t@(Lam {}) u =
    throw $ NotType t.loc ctx t

  checkType' ctx (Sig fst snd _) u = do
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    checkTypeC ctx fst u
    -- if Ψ | Γ, x : A ⊢₀ B ⇐ Type ℓ
    checkTypeScope ctx fst snd u
    -- then Ψ | Γ ⊢₀ (x : A) × B ⇐ Type ℓ

  checkType' ctx t@(Pair {}) u =
    throw $ NotType t.loc ctx t

  checkType' ctx (Enum {}) u = pure ()
    -- Ψ | Γ ⊢₀ {ts} ⇐ Type ℓ

  checkType' ctx t@(Tag {}) u =
    throw $ NotType t.loc ctx t

  checkType' ctx (Eq t l r _) u = do
    -- if Ψ, i | Γ ⊢₀ A ⇐ Type ℓ
    case t.body of
      Y t' => checkTypeC (extendDim t.name ctx) t' u
      N t' => checkTypeC ctx                    t' u
    -- if Ψ | Γ ⊢₀ l ⇐ A‹0›
    check0 ctx l t.zero
    -- if Ψ | Γ ⊢₀ r ⇐ A‹1›
    check0 ctx r t.one
    -- then Ψ | Γ ⊢₀ Eq [i ⇒ A] l r ⇐ Type ℓ

  checkType' ctx t@(DLam {}) u =
    throw $ NotType t.loc ctx t

  checkType' ctx (Nat {}) u = pure ()
  checkType' ctx t@(Zero {}) u = throw $ NotType t.loc ctx t
  checkType' ctx t@(Succ {}) u = throw $ NotType t.loc ctx t

  checkType' ctx (BOX q ty _) u = checkType ctx ty u
  checkType' ctx t@(Box {}) u = throw $ NotType t.loc ctx t

  checkType' ctx (E e) u = do
    -- if Ψ | Γ ⊢₀ E ⇒ Type ℓ
    infres <- inferC ctx SZero e
    -- if Ψ | Γ ⊢ Type ℓ <: Type 𝓀
    case u of
      Just u  => lift $ subtype e.loc ctx infres.type (TYPE u noLoc)
      Nothing => ignore $ expectTYPE !(askAt DEFS) ctx SZero e.loc infres.type
    -- then Ψ | Γ ⊢₀ E ⇐ Type 𝓀


  private covering
  checkTypeScope : TyContext d n -> Term d n ->
                   ScopeTerm d n -> Maybe Universe -> Eff TC ()
  checkTypeScope ctx s (S _     (N body)) u = checkType ctx body u
  checkTypeScope ctx s (S [< x] (Y body)) u =
    checkType (extendTy Zero x s ctx) body u


  private covering
  infer' : TyContext d n -> SQty ->
           (subj : Elim d n) -> (0 nc : NotClo subj) =>
           Eff TC (InferResult' d n)

  infer' ctx sg (F x u loc) = do
    -- if π·x : A {≔ s} in global context
    g <- lookupFree x loc !(askAt DEFS)
    -- if σ ≤ π
    expectCompatQ loc sg.qty g.qty.qty
    -- then Ψ | Γ ⊢ σ · x ⇒ A ⊳ 𝟎
    let Val d = ctx.dimLen; Val n = ctx.termLen
    pure $ InfRes {type = g.typeAt u, qout = zeroFor ctx}

  infer' ctx sg (B i _) =
    -- if x : A ∈ Γ
    -- then Ψ | Γ ⊢ σ · x ⇒ A ⊳ (𝟎, σ·x, 𝟎)
    pure $ lookupBound sg.qty i ctx.tctx
  where
    lookupBound : forall n. Qty -> Var n -> TContext d n -> InferResult' d n
    lookupBound pi VZ (ctx :< type) =
      InfRes {type = weakT 1 type, qout = zeroFor ctx :< pi}
    lookupBound pi (VS i) (ctx :< _) =
      let InfRes {type, qout} = lookupBound pi i ctx in
      InfRes {type = weakT 1 type, qout = qout :< Zero}

  infer' ctx sg (App fun arg loc) = do
    -- if Ψ | Γ ⊢ σ · f ⇒ (π·x : A) → B ⊳ Σ₁
    funres <- inferC ctx sg fun
    (qty, argty, res) <- expectPi !(askAt DEFS) ctx SZero fun.loc funres.type
    -- if Ψ | Γ ⊢ σ ⨴ π · s ⇐ A ⊳ Σ₂
    argout <- checkC ctx (subjMult sg qty) arg argty
    -- then Ψ | Γ ⊢ σ · f s ⇒ B[s] ⊳ Σ₁ + πΣ₂
    pure $ InfRes {
      type = sub1 res $ Ann arg argty arg.loc,
      qout = funres.qout + qty * argout
    }

  infer' ctx sg (CasePair pi pair ret body loc) = do
    -- no check for 1 ≤ π, since pairs have a single constructor.
    -- e.g. at 0 the components are also 0 in the body
    --
    -- if Ψ | Γ ⊢ σ · pair ⇒ (x : A) × B ⊳ Σ₁
    pairres <- inferC ctx sg pair
    -- if Ψ | Γ, p : (x : A) × B ⊢₀ ret ⇐ Type
    checkTypeC (extendTy Zero ret.name pairres.type ctx) ret.term Nothing
    (tfst, tsnd) <- expectSig !(askAt DEFS) ctx SZero pair.loc pairres.type
    -- if Ψ | Γ, x : A, y : B ⊢ σ · body ⇐
    --    ret[(x, y) ∷ (x : A) × B/p] ⊳ Σ₂, ρ₁·x, ρ₂·y
    -- with ρ₁, ρ₂ ≤ πσ
    let [< x, y]  = body.names
        pisg    = pi * sg.qty
        bodyctx = extendTyN [< (pisg, x, tfst), (pisg, y, tsnd.term)] ctx
        bodyty  = substCasePairRet body.names pairres.type ret
    bodyout <- checkC bodyctx sg body.term bodyty >>=
               popQs loc [< pisg, pisg]
    -- then Ψ | Γ ⊢ σ · caseπ ⋯ ⇒ ret[pair/p] ⊳ πΣ₁ + Σ₂
    pure $ InfRes {
      type = sub1 ret pair,
      qout = pi * pairres.qout + bodyout
    }

  infer' ctx sg (Fst pair loc) = do
    -- if Ψ | Γ ⊢ σ · e ⇒ (x : A) × B ⊳ Σ
    pairres <- inferC ctx sg pair
    (tfst, _) <- expectSig !(askAt DEFS) ctx SZero pair.loc pairres.type
    -- then Ψ | Γ ⊢ σ · fst e ⇒ A ⊳ ωΣ
    pure $ InfRes {
      type = tfst,
      qout = Any * pairres.qout
    }

  infer' ctx sg (Snd pair loc) = do
    -- if Ψ | Γ ⊢ σ · e ⇒ (x : A) × B ⊳ Σ
    pairres <- inferC ctx sg pair
    (_, tsnd) <- expectSig !(askAt DEFS) ctx SZero pair.loc pairres.type
    -- then Ψ | Γ ⊢ σ · snd e ⇒ B[fst e/x] ⊳ ωΣ
    pure $ InfRes {
      type = sub1 tsnd (Fst pair loc),
      qout = Any * pairres.qout
    }

  infer' ctx sg (CaseEnum pi t ret arms loc) {d, n} = do
    -- if Ψ | Γ ⊢ σ · t ⇒ {ts} ⊳ Σ₁
    tres <- inferC ctx sg t
    ttags <- expectEnum !(askAt DEFS) ctx SZero t.loc tres.type
    -- if 1 ≤ π, OR there is only zero or one option
    unless (length (SortedSet.toList ttags) <= 1) $ expectCompatQ loc One pi
    -- if Ψ | Γ, x : {ts} ⊢₀ A ⇐ Type
    checkTypeC (extendTy Zero ret.name tres.type ctx) ret.term Nothing
    -- if for each "a ⇒ s" in arms,
    --   Ψ | Γ ⊢ σ · s ⇐ A[a ∷ {ts}/x] ⊳ Σᵢ
    -- with Σ₂ = lubs Σᵢ
    let arms = SortedMap.toList arms
    let armTags = SortedSet.fromList $ map fst arms
    unless (ttags == armTags) $ throw $ BadCaseEnum loc ttags armTags
    armres <- for arms $ \(a, s) =>
      checkC ctx sg s $ sub1 ret $ Ann (Tag a s.loc) tres.type s.loc
    pure $ InfRes {
      type = sub1 ret t,
      qout = pi * tres.qout + lubs ctx armres
    }

  infer' ctx sg (CaseNat pi pi' n ret zer suc loc) = do
    -- if 1 ≤ π
    expectCompatQ loc One pi
    -- if Ψ | Γ ⊢ σ · n ⇒ ℕ ⊳ Σn
    nres <- inferC ctx sg n
    let nat = nres.type
    expectNat !(askAt DEFS) ctx SZero n.loc nat
    -- if Ψ | Γ, n : ℕ ⊢₀ A ⇐ Type
    checkTypeC (extendTy Zero ret.name nat ctx) ret.term Nothing
    -- if Ψ | Γ ⊢ σ · zer ⇐ A[0 ∷ ℕ/n] ⊳ Σz
    zerout <- checkC ctx sg zer $ sub1 ret $ Ann (Zero zer.loc) nat zer.loc
    -- if Ψ | Γ, n : ℕ, ih : A ⊢ σ · suc ⇐ A[succ p ∷ ℕ/n] ⊳ Σs, ρ₁.p, ρ₂.ih
    -- with ρ₂ ≤ π'σ, (ρ₁ + ρ₂) ≤ πσ
    let [< p, ih] = suc.names
        pisg      = pi * sg.qty
        sucCtx    = extendTyN [< (pisg, p, Nat p.loc), (pi', ih, ret.term)] ctx
        sucType   = substCaseSuccRet suc.names ret
    sucout :< qp :< qih <- checkC sucCtx sg suc.term sucType
    expectCompatQ loc qih (pi' * sg.qty)
    -- [fixme] better error here
    expectCompatQ loc (qp + qih) pisg
    -- then Ψ | Γ ⊢ caseπ ⋯ ⇒ A[n] ⊳ πΣn + Σz + ωΣs
    pure $ InfRes {
      type = sub1 ret n,
      qout = pi * nres.qout + zerout + Any * sucout
    }

  infer' ctx sg (CaseBox pi box ret body loc) = do
    -- if Ψ | Γ ⊢ σ · b ⇒ [ρ.A] ⊳ Σ₁
    boxres <- inferC ctx sg box
    (q, ty) <- expectBOX !(askAt DEFS) ctx SZero box.loc boxres.type
    -- if Ψ | Γ, x : [ρ.A] ⊢₀ R ⇐ Type
    checkTypeC (extendTy Zero ret.name boxres.type ctx) ret.term Nothing
    -- if Ψ | Γ, x : A ⊢ t ⇐ R[[x] ∷ [ρ.A/x]] ⊳ Σ₂, ς·x
    --   with ς ≤ ρπσ
    let qpisg    = q * pi * sg.qty
        bodyCtx  = extendTy qpisg body.name ty ctx
        bodyType = substCaseBoxRet body.name ty ret
    bodyout <- checkC bodyCtx sg body.term bodyType >>= popQ loc qpisg
    -- then Ψ | Γ ⊢ caseπ ⋯ ⇒ R[b/x] ⊳ Σ₁ + Σ₂
    pure $ InfRes {
      type = sub1 ret box,
      qout = boxres.qout + bodyout
    }

  infer' ctx sg (DApp fun dim loc) = do
    -- if Ψ | Γ ⊢ σ · f ⇒ Eq [𝑖 ⇒ A] l r ⊳ Σ
    InfRes {type, qout} <- inferC ctx sg fun
    ty <- fst <$> expectEq !(askAt DEFS) ctx SZero fun.loc type
    -- then Ψ | Γ ⊢ σ · f p ⇒ A‹p/𝑖› ⊳ Σ
    pure $ InfRes {type = dsub1 ty dim, qout}

  infer' ctx sg (Coe ty p q val loc) = do
    checkType (extendDim ty.name ctx) ty.term Nothing
    qout <- checkC ctx sg val $ dsub1 ty p
    pure $ InfRes {type = dsub1 ty q, qout}

  infer' ctx sg (Comp ty p q val r (S [< j0] val0) (S [< j1] val1) loc) = do
    checkType ctx ty Nothing
    qout <- checkC ctx sg val ty
    let ty'  = dweakT 1 ty; val' = dweakT 1 val; p' = weakD 1 p
        ctx0 = extendDim j0 $ eqDim r (K Zero j0.loc) ctx
        val0 = getTerm val0
    qout0 <- check ctx0 sg val0 ty'
    lift $ equal loc (eqDim (B VZ p.loc) p' ctx0) sg ty' val0 val'
    let ctx1 = extendDim j1 $ eqDim r (K One j1.loc) ctx
        val1  = getTerm val1
    qout1 <- check ctx1 sg val1 ty'
    lift $ equal loc (eqDim (B VZ p.loc) p' ctx1) sg ty' val1 val'
    let qouts  = qout :: catMaybes [toMaybe qout0, toMaybe qout1]
    pure $ InfRes {type = ty, qout = lubs ctx qouts}

  infer' ctx sg (TypeCase ty ret arms def loc) = do
    -- if σ = 0
    expectEqualQ loc Zero sg.qty
    -- if Ψ, Γ ⊢₀ e ⇒ Type u
    u <- inferC ctx SZero ty >>=
         expectTYPE !(askAt DEFS) ctx SZero ty.loc . type
    -- if Ψ, Γ ⊢₀ C ⇐ Type (non-dependent return type)
    checkTypeC ctx ret Nothing
    -- if Ψ, Γ' ⊢₀ A ⇐ C for each rhs A
    for_ allKinds $ \k =>
      for_ (lookupPrecise k arms) $ \(S names t) =>
        check0 (extendTyN (typecaseTel k names u) ctx)
               (getTerm t) (weakT (arity k) ret)
    -- then Ψ, Γ ⊢₀ type-case ⋯ ⇒ C
    pure $ InfRes {type = ret, qout = zeroFor ctx}

  infer' ctx sg (Ann term type loc) = do
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    checkTypeC ctx type Nothing
    -- if Ψ | Γ ⊢ σ · s ⇐ A ⊳ Σ
    qout <- checkC ctx sg term type
    -- then Ψ | Γ ⊢ σ · (s ∷ A) ⇒ A ⊳ Σ
    pure $ InfRes {type, qout}