quox/lib/Quox/Parser/Syntax.idr

module Quox.Parser.Syntax

import public Quox.Syntax
import public Quox.Syntax.Qty.Three
import public Quox.Definition

import public Control.Monad.Either

import Derive.Prelude
%hide TT.Name

%default total
%language ElabReflection


public export
0 BName : Type
BName = Maybe BaseName

public export
0 PUniverse : Type
PUniverse = Nat

public export
0 PQty : Type
PQty = Three

namespace PDim
  public export
  data PDim = K DimConst | V BaseName
  %name PDim p, q
  %runElab derive "PDim" [Eq, Ord, Show]

namespace PTerm
  mutual
    ||| terms out of the parser with BVs and bidirectionality still tangled up
    public export
    data PTerm =
        TYPE Nat

      | Pi PQty BName PTerm PTerm
      | Lam BName PTerm
      | (:@) PTerm PTerm

      | Sig BName PTerm PTerm
      | Pair PTerm PTerm
      | Case PQty PTerm (BName, PTerm) (PCaseBody)

      | Enum (List TagVal)
      | Tag TagVal

      | Eq (BName, PTerm) PTerm PTerm
      | DLam BName PTerm
      | (:%) PTerm PDim

      | V Name
      | (:#) PTerm PTerm
    %name PTerm s, t

    public export
    data PCaseBody =
        CasePair (BName, BName) PTerm
      | CaseEnum (List (TagVal, PTerm))
    %name PCaseBody body

%runElab deriveMutual ["PTerm", "PCaseBody"] [Eq, Ord, Show]


public export
record PDefinition where
  constructor MkPDef
  qty  : PQty
  name : Name
  type : PTerm
  term : PTerm
%name PDefinition def
%runElab derive "PDefinition" [Eq, Ord, Show]


export
toPDimWith : Context' BaseName d -> Dim d -> PDim
toPDimWith ds (K e) = K e
toPDimWith ds (B i) = V $ ds !!! i

export
toPDim : Dim 0 -> PDim
toPDim = toPDimWith [<]


mutual
  namespace Term
    export
    toPTermWith : Context' BaseName d -> Context' BaseName n ->
                  Term Three d n -> PTerm
    toPTermWith ds ns t =
      let Element t _ = pushSubsts t in
      toPTermWith' ds ns t

    private
    toPTermWith' : Context' BaseName d -> Context' BaseName n ->
                   (t : Term Three d n) -> (0 _ : NotClo t) =>
                   PTerm
    toPTermWith' ds ns s = case s of
      TYPE l =>
        TYPE l
      Pi qty arg (S [x] res) =>
        Pi qty (Just x) (toPTermWith ds ns arg)
                        (toPTermWith ds (ns :< x) res.term)
      Lam (S [x] body) =>
        Lam (Just x) $ toPTermWith ds (ns :< x) body.term
      Sig fst (S [x] snd) =>
        Sig (Just x) (toPTermWith ds ns fst)
                     (toPTermWith ds (ns :< x) snd.term)
      Pair fst snd =>
        Pair (toPTermWith ds ns fst) (toPTermWith ds ns snd)
      Enum cases =>
        Enum $ SortedSet.toList cases
      Tag tag =>
        Tag tag
      Eq (S [i] ty) l r =>
        Eq (Just i, toPTermWith (ds :< i) ns ty.term)
           (toPTermWith ds ns l) (toPTermWith ds ns r)
      DLam (S [i] body) =>
        DLam (Just i) $ toPTermWith (ds :< i) ns body.term
      E e =>
        toPTermWith ds ns e

  namespace Elim
    export
    toPTermWith : Context' BaseName d -> Context' BaseName n ->
                  Elim Three d n -> PTerm
    toPTermWith ds ns e =
      let Element e _ = pushSubsts e in
      toPTermWith' ds ns e

    private
    toPTermWith' : Context' BaseName d -> Context' BaseName n ->
                   (e : Elim Three d n) -> (0 _ : NotClo e) =>
                   PTerm
    toPTermWith' ds ns e = case e of
      F x =>
        V x
      B i =>
        V $ unq $ ns !!! i
      fun :@ arg =>
        toPTermWith ds ns fun :@ toPTermWith ds ns arg
      CasePair qty pair (S [r] ret) (S [x, y] body) =>
        Case qty (toPTermWith ds ns pair)
             (Just r, toPTermWith ds (ns :< r) ret.term)
             (CasePair (Just x, Just y) $
                toPTermWith ds (ns :< x :< y) body.term)
      CaseEnum qty tag (S [r] ret) arms =>
        Case qty (toPTermWith ds ns tag)
             (Just r, toPTermWith ds (ns :< r) ret.term)
             (CaseEnum $ mapSnd (toPTermWith ds ns) <$> SortedMap.toList arms)
      fun :% arg =>
        toPTermWith ds ns fun :% toPDimWith ds arg
      tm :# ty =>
        toPTermWith ds ns tm :# toPTermWith ds ns ty

namespace Term
  export
  toPTerm : Term Three 0 0 -> PTerm
  toPTerm = toPTermWith [<] [<]

namespace Elim
  export
  toPTerm : Elim Three 0 0 -> PTerm
  toPTerm = toPTermWith [<] [<]



public export
data FromPTermError =
    AnnotationNeeded PTerm
  | DuplicatesInEnum (List TagVal)
  | DimNotInScope Name
  | QtyNotGlobal PQty

public export
FromPTerm : (Type -> Type) -> Type
FromPTerm m = MonadError FromPTermError m

private
extendVect : Telescope' a from to -> Vect n a -> Telescope' a from (n + to)
extendVect tel [] = tel
extendVect tel (x :: xs) {n = S n} =
  rewrite plusSuccRightSucc n to in
  extendVect (tel :< x) xs


parameters {auto _ : FromPTerm m} (b : Var n -> m a) (f : Name -> m a)
           (xs : Context' BName n)
  private
  fromBaseName : BaseName -> m a
  fromBaseName x = maybe (f $ unq x) b $ Context.find (== Just x) xs

  fromName : Name -> m a
  fromName x = if null x.mods then fromBaseName x.base else f x

export
fromPDimWith : FromPTerm m =>
               Context' BName d -> PDim -> m (Dim d)
fromPDimWith ds (K e) = pure $ K e
fromPDimWith ds (V i) = fromBaseName (pure . B) (throwError . DimNotInScope) ds i


mutual
  export
  fromPTermWith : FromPTerm m =>
                  Context' BName d -> Context' BName n ->
                  PTerm -> m (Term Three d n)
  fromPTermWith ds ns t0 = case t0 of
    TYPE k =>
      pure $ TYPE $ k

    Pi pi x s t =>
      Pi pi <$> fromPTermWith ds ns s
            <*> fromPTermTScope ds ns [x] t

    Lam x s =>
      Lam <$> fromPTermTScope ds ns [x] s

    s :@ t =>
      map E $ (:@) <$> fromPTermElim ds ns s <*> fromPTermWith ds ns t

    Sig x s t =>
      Sig <$> fromPTermWith ds ns s <*> fromPTermTScope ds ns [x] t

    Pair s t =>
      Pair <$> fromPTermWith ds ns s <*> fromPTermWith ds ns t

    Case pi pair (r, ret) (CasePair (x, y) body) =>
      map E $ Base.CasePair pi
        <$> fromPTermElim ds ns pair
        <*> fromPTermTScope ds ns [r] ret
        <*> fromPTermTScope ds ns [x, y] body

    Case pi tag (r, ret) (CaseEnum arms) =>
      map E $ Base.CaseEnum pi
        <$> fromPTermElim ds ns tag
        <*> fromPTermTScope ds ns [r] ret
        <*> assert_total fromPTermEnumArms ds ns arms

    Enum strs =>
      let set = SortedSet.fromList strs in
      if length strs == length (SortedSet.toList set) then
        pure $ Enum set
      else
        throwError $ DuplicatesInEnum strs

    Tag str => pure $ Tag str

    Eq (i, ty) s t =>
      Eq <$> fromPTermDScope ds ns [i] ty
         <*> fromPTermWith ds ns s
         <*> fromPTermWith ds ns t

    DLam i s =>
      DLam <$> fromPTermDScope ds ns [i] s

    s :% p =>
      map E $ (:%) <$> fromPTermElim ds ns s <*> fromPDimWith ds p

    V x =>
      fromName (pure . E . B) (pure . FT) ns x

    s :# a =>
      map E $ (:#) <$> fromPTermWith ds ns s <*> fromPTermWith ds ns a

  private
  fromPTermEnumArms : FromPTerm m =>
                      Context' BName d -> Context' BName n ->
                      List (TagVal, PTerm) -> m (CaseEnumArms Three d n)
  fromPTermEnumArms ds ns =
    map SortedMap.fromList . traverse (traverse $ fromPTermWith ds ns)

  private
  fromPTermElim : FromPTerm m =>
                  Context' BName d -> Context' BName n ->
                  PTerm -> m (Elim Three d n)
  fromPTermElim ds ns e =
    case !(fromPTermWith ds ns e) of
      E e => pure e
      _   => throwError $ AnnotationNeeded e

  -- [todo] use SN if the var is named but still unused
  private
  fromPTermTScope : {s : Nat} -> FromPTerm m =>
                    Context' BName d -> Context' BName n ->
                    Vect s BName ->
                    PTerm -> m (ScopeTermN s Three d n)
  fromPTermTScope ds ns xs t =
    if all isNothing xs then
      SN <$> fromPTermWith ds ns t
    else
      SY (map (fromMaybe "_") xs) <$> fromPTermWith ds (extendVect ns xs) t

  private
  fromPTermDScope : {s : Nat} -> FromPTerm m =>
                    Context' BName d -> Context' BName n ->
                    Vect s BName ->
                    PTerm -> m (DScopeTermN s Three d n)
  fromPTermDScope ds ns xs t =
    if all isNothing xs then
      SN <$> fromPTermWith ds ns t
    else
      SY (map (fromMaybe "_") xs) <$> fromPTermWith (extendVect ds xs) ns t


export %inline
fromPTerm : FromPTerm m => PTerm -> m (Term Three 0 0)
fromPTerm = fromPTermWith [<] [<]


export
globalPQty : FromPTerm m => (q : PQty) -> m (IsGlobal q)
globalPQty Zero = pure GZero
globalPQty One  = throwError $ QtyNotGlobal One
globalPQty Any  = pure GAny


-- [todo] extend substitutions so they can do this injection. that's the sort of
-- thing they are for.
export
fromPDefinition : FromPTerm m => PDefinition -> m (Name, Definition Three)
fromPDefinition (MkPDef {name, qty, type, term}) =
  pure (name, MkDef' {
    qty, qtyGlobal = !(globalPQty qty),
    type = T $ inject !(fromPTerm type),
    term = Just $ T $ inject !(fromPTerm term)
  })