quox/lib/Quox/Equal.idr

module Quox.Equal

import public Quox.Syntax
import public Quox.Definition
import public Quox.Typing
import public Control.Monad.Either
import public Control.Monad.Reader
import Data.Maybe


private %inline
ClashE : EqMode -> Term q d n -> Elim q d n -> Elim q d n -> Error q
ClashE mode ty = ClashT mode ty `on` E


public export
record Env where
  constructor MakeEnv
  mode : EqMode


public export
0 HasEnv : (Type -> Type) -> Type
HasEnv = MonadReader Env

public export
0 CanEqual : (q : Type) -> (Type -> Type) -> Type
CanEqual q m = (HasErr q m, HasEnv m)



private %inline
mode : HasEnv m => m EqMode
mode = asks mode

private %inline
clashT : CanEqual q m => Term q d n -> Term q d n -> Term q d n -> m a
clashT ty s t = throwError $ ClashT !mode ty s t

private %inline
clashE : CanEqual q m => Elim q d n -> Elim q d n -> m a
clashE e f = throwError $ ClashE !mode e f


public export %inline
isTyCon : (t : Term {}) -> Bool
isTyCon (TYPE  {}) = True
isTyCon (Pi    {}) = True
isTyCon (Lam   {}) = False
isTyCon (Sig   {}) = True
isTyCon (Pair  {}) = False
isTyCon (Eq    {}) = True
isTyCon (DLam  {}) = False
isTyCon (E     {}) = True
isTyCon (CloT  {}) = False
isTyCon (DCloT {}) = False

public export %inline
sameTyCon : (s, t : Term q d n) ->
            (0 ts : So (isTyCon s)) => (0 tt : So (isTyCon t)) =>
            Bool
sameTyCon (TYPE {}) (TYPE {}) = True
sameTyCon (TYPE {}) _         = False
sameTyCon (Pi   {}) (Pi   {}) = True
sameTyCon (Pi   {}) _         = False
sameTyCon (Sig  {}) (Sig  {}) = True
sameTyCon (Sig  {}) _         = False
sameTyCon (Eq   {}) (Eq   {}) = True
sameTyCon (Eq   {}) _         = False
sameTyCon (E    {}) (E    {}) = True
sameTyCon (E    {}) _         = False


parameters (defs : Definitions' q g)
  private
  isSubSing : Term q 0 n -> Bool
  isSubSing ty =
    let Element ty nc = whnfD defs ty in
    case ty of
      TYPE _             => False
      Pi   {res, _}      => isSubSing res.term
      Lam  {}            => False
      Sig  {fst, snd, _} => isSubSing fst && isSubSing snd.term
      Pair {}            => False
      Eq   {}            => True
      DLam {}            => False
      E    (s :# _)      => isSubSing s
      E    _             => False


parameters {auto _ : HasErr q m}
  export %inline
  ensure : (a -> Error q) -> (p : a -> Bool) -> (t : a) -> m (So (p t))
  ensure e p t = case nchoose $ p t of
    Left  y => pure y
    Right _ => throwError $ e t

  export %inline
  ensureTyCon : HasErr q m => (t : Term q d n) -> m (So (isTyCon t))
  ensureTyCon = ensure NotType isTyCon

parameters (defs : Definitions' q _) {auto _ : (CanEqual q m, Eq q)}
  mutual
    namespace Term
      export covering %inline
      compare0 : TContext q 0 n -> (ty, s, t : Term q 0 n) -> m ()
      compare0 ctx ty s t = do
        let Element ty nty = whnfD defs ty
            Element s  ns  = whnfD defs s
            Element t  nt  = whnfD defs t
        tty <- ensureTyCon ty
        compare0' ctx ty s t

      private %inline
      toLamBody : Elim q d n -> Term q d (S n)
      toLamBody e = E $ weakE e :@ BVT 0

      private covering
      compare0' : TContext q 0 n ->
                  (ty, s, t : Term q 0 n) ->
                  (0 nty : NotRedex defs ty) => (0 tty : So (isTyCon ty)) =>
                  (0 ns  : NotRedex defs s)  => (0 nt  : NotRedex defs t) =>
                  m ()
      compare0' ctx (TYPE _) s t = compareType ctx s t

      compare0' ctx ty@(Pi {arg, res, _}) s t {n} = local {mode := Equal} $
        let ctx' = ctx :< arg
            eta : Elim q 0 n -> ScopeTerm q 0 n -> m ()
            eta e (TUsed   b) = compare0 ctx' res.term (toLamBody e) b
            eta e (TUnused _) = clashT ty s t
        in
        case (s, t) of
          (Lam _ b1, Lam _ b2) => compare0 ctx' res.term b1.term b2.term
          (E e,      Lam _ b)  => eta e b
          (Lam _ b,  E e)      => eta e b
          (E e,      E f)      => Elim.compare0 ctx e f
          _ => throwError $ WrongType ty s t

      compare0' ctx ty@(Sig {fst, snd, _}) s t = local {mode := Equal} $
        -- no η (no fst/snd for π ≱ 0), so…
        -- [todo] η for π ≥ 0 maybe
        case (s, t) of
          (Pair sFst sSnd, Pair tFst tSnd) => do
            compare0 ctx fst sFst tFst
            compare0 ctx (sub1 snd (sFst :# fst)) sSnd tSnd
          _ => throwError $ WrongType ty s t

      -- ✨ uip ✨
      compare0' _ (Eq {}) _ _ = pure ()

      compare0' ctx ty@(E _) s t = do
        -- a neutral type can only be inhabited by neutral values
        -- e.g. an abstract value in an abstract type, bound variables, …
        E e <- pure s | _ => throwError $ WrongType ty s t
        E f <- pure t | _ => throwError $ WrongType ty s t
        Elim.compare0 ctx e f

    export covering
    compareType : TContext q 0 n -> (s, t : Term q 0 n) -> m ()
    compareType ctx s t = do
      let Element s ns = whnfD defs s
          Element t nt = whnfD defs t
      ts <- ensureTyCon s
      tt <- ensureTyCon t
      st <- either pure (const $ clashT (TYPE UAny) s t) $
            nchoose $ sameTyCon s t
      compareType' ctx s t

    private covering
    compareType' : TContext q 0 n -> (s, t : Term q 0 n) ->
                   (0 ns : NotRedex defs s) => (0 ts : So (isTyCon s)) =>
                   (0 nt : NotRedex defs t) => (0 tt : So (isTyCon t)) =>
                   (0 st : So (sameTyCon s t)) =>
                   m ()
    compareType' ctx (TYPE k) (TYPE l) =
      expectModeU !mode k l

    compareType' ctx (Pi {qty = sQty, arg = sArg, res = sRes, _})
                     (Pi {qty = tQty, arg = tArg, res = tRes, _}) = do
      expectEqualQ sQty tQty
      local {mode $= flip} $ compareType ctx sArg tArg -- contra
      compareType (ctx :< sArg) sRes.term tRes.term

    compareType' ctx (Sig {fst = sFst, snd = sSnd, _})
                     (Sig {fst = tFst, snd = tSnd, _}) = do
      compareType ctx sFst tFst
      compareType (ctx :< sFst) sSnd.term tSnd.term

    compareType' ctx (Eq {ty = sTy, l = sl, r = sr, _})
                     (Eq {ty = tTy, l = tl, r = tr, _}) = do
      compareType ctx sTy.zero tTy.zero
      compareType ctx sTy.one  tTy.one
      local {mode := Equal} $ do
        Term.compare0 ctx sTy.zero sl tl
        Term.compare0 ctx sTy.one  sr tr

    compareType' ctx (E e) (E f) = do
      -- no fanciness needed here cos anything other than a neutral
      -- has been inlined by whnfD
      Elim.compare0 ctx e f

    ||| assumes the elim is already typechecked! only does the work necessary
    ||| to calculate the overall type
    private covering
    computeElimType : TContext q 0 n -> (e : Elim q 0 n) ->
                      (0 ne : NotRedex defs e) ->
                      m (Term q 0 n)
    computeElimType ctx (F x) _ = do
      defs <- lookupFree' defs x
      pure $ defs.type.get
    computeElimType ctx (B i) _ = do
      pure $ ctx !! i
    computeElimType ctx (f :@ s) ne = do
      (_, arg, res) <- expectPi defs !(computeElimType ctx f (noOr1 ne))
      pure $ sub1 res (s :# arg)
    computeElimType ctx (CasePair {pair, ret, _}) _ = do
      pure $ sub1 ret pair
    computeElimType ctx (f :% p) ne = do
      (ty, _, _) <- expectEq defs !(computeElimType ctx f (noOr1 ne))
      pure $ dsub1 ty p
    computeElimType ctx (_ :# ty) _ = do
      pure ty

    private covering
    replaceEnd : TContext q 0 n ->
                 (e : Elim q 0 n) -> DimConst -> (0 ne : NotRedex defs e) ->
                 m (Elim q 0 n)
    replaceEnd ctx e p ne = do
      (ty, l, r) <- expectEq defs !(computeElimType ctx e ne)
      pure $ ends l r p :# dsub1 ty (K p)

    namespace Elim
      -- [fixme] the following code ends up repeating a lot of work in the
      -- computeElimType calls. the results should be shared better
      export covering %inline
      compare0 : TContext q 0 n -> (e, f : Elim q 0 n) -> m ()
      compare0 ctx e f =
        let Element e ne = whnfD defs e
            Element f nf = whnfD defs f
        in
        -- [fixme] there is a better way to do this "isSubSing" stuff for sure
        unless (isSubSing defs !(computeElimType ctx e ne)) $
          compare0' ctx e f ne nf

      private covering
      compare0' : TContext q 0 n ->
                  (e, f : Elim q 0 n) ->
                  (0 ne : NotRedex defs e) -> (0 nf : NotRedex defs f) ->
                  m ()
      -- replace applied equalities with the appropriate end first
      -- e.g. e : Eq [i ⇒ A] s t ⊢ e 𝟎 = s : A‹𝟎/i›
      compare0' ctx (e :% K p) f ne nf =
        compare0 ctx !(replaceEnd ctx e p $ noOr1 ne) f
      compare0' ctx e (f :% K q) ne nf =
        compare0 ctx e !(replaceEnd ctx f q $ noOr1 nf)

      compare0' _ e@(F x) f@(F y) _ _ = unless (x == y) $ clashE e f
      compare0' _ e@(F _) f _ _ = clashE e f

      compare0' ctx e@(B i) f@(B j) _ _ = unless (i == j) $ clashE e f
      compare0' _ e@(B _) f _ _ = clashE e f

      compare0' ctx (e :@ s) (f :@ t) ne nf =
        local {mode := Equal} $ do
          compare0 ctx e f
          (_, arg, _) <- expectPi defs !(computeElimType ctx e (noOr1 ne))
          Term.compare0 ctx arg s t
      compare0' _ e@(_ :@ _) f _ _ = clashE e f

      compare0' ctx (CasePair epi e _ eret _ _ ebody)
                    (CasePair fpi f _ fret _ _ fbody) ne nf =
        local {mode := Equal} $ do
          compare0 ctx e f
          ety <- computeElimType ctx e (noOr1 ne)
          compareType (ctx :< ety) eret.term fret.term
          (fst, snd) <- expectSig defs ety
          Term.compare0 (ctx :< fst :< snd.term) (substCasePairRet ety eret)
            ebody.term fbody.term
          unless (epi == fpi) $ throwError $ ClashQ epi fpi
      compare0' _ e@(CasePair {}) f _ _ = clashE e f

      compare0' ctx   (s :# a) (t :# _) _ _ = Term.compare0 ctx a s t
      compare0' ctx   (s :# a) f        _ _ = Term.compare0 ctx a s (E f)
      compare0' ctx e          (t :# b) _ _ = Term.compare0 ctx b (E e) t
      compare0' _   e@(_ :# _) f        _ _ = clashE e f


parameters {auto _ : (HasDefs' q _ m, HasErr q m, Eq q)} (ctx : TyContext q d n)
  parameters (mode : EqMode)
    namespace Term
      export covering
      compare : (ty, s, t : Term q d n) -> m ()
      compare ty s t = do
        defs <- ask
        runReaderT {m} (MakeEnv {mode}) $
          for_ (splits ctx.dctx) $ \th =>
            compare0 defs (map (/// th) ctx.tctx)
                     (ty /// th) (s /// th) (t /// th)

      export covering
      compareType : (s, t : Term q d n) -> m ()
      compareType s t = do
        defs <- ask
        runReaderT {m} (MakeEnv {mode}) $
          for_ (splits ctx.dctx) $ \th =>
            compareType defs (map (/// th) ctx.tctx) (s /// th) (t /// th)

    namespace Elim
      ||| you don't have to pass the type in but the arguments must still be
      ||| of the same type!!
      export covering %inline
      compare : (e, f : Elim q d n) -> m ()
      compare e f = do
        defs <- ask
        runReaderT {m} (MakeEnv {mode}) $
          for_ (splits ctx.dctx) $ \th =>
            compare0 defs (map (/// th) ctx.tctx) (e /// th) (f /// th)

  namespace Term
    export covering %inline
    equal, sub : (ty, s, t : Term q d n) -> m ()
    equal = compare Equal
    sub   = compare Sub

    export covering %inline
    equalType, subtype : (s, t : Term q d n) -> m ()
    equalType = compareType Equal
    subtype   = compareType Sub

  namespace Elim
    export covering %inline
    equal, sub : (e, f : Elim q d n) -> m ()
    equal = compare Equal
    sub   = compare Sub