quox/lib/Quox/Untyped/Syntax.idr

module Quox.Untyped.Syntax

import Quox.Var
import Quox.Context
import Quox.Name
import Quox.Pretty
import Quox.Syntax.Subst

import Data.Vect
import Data.DPair
import Data.SortedMap
import Data.SnocVect
import Derive.Prelude
%hide TT.Name

%default total
%language ElabReflection


public export
data Term : Nat -> Type where
  F : (x : Name) -> Loc -> Term n
  B : (i : Var n) -> Loc -> Term n

  Lam : (x : BindName) -> (body : Term (S n)) -> Loc -> Term n
  App : (fun, arg : Term n) -> Loc -> Term n

  Pair : (fst, snd : Term n) -> Loc -> Term n
  Fst : (pair : Term n) -> Loc -> Term n
  Snd : (pair : Term n) -> Loc -> Term n

  Tag : (tag : String) -> Loc -> Term n
  CaseEnum : (tag : Term n) -> (cases : List (String, Term n)) -> Loc -> Term n
  ||| empty match with an erased head
  Absurd : Loc -> Term n

  Zero : Loc -> Term n
  Succ : (nat : Term n) -> Loc -> Term n
  CaseNat : (nat : Term n) ->
            (zer : Term n) ->
            (x, ih : BindName) -> (suc : Term (2 + n)) ->
            Loc ->
            Term n

  Erased : Loc -> Term n
%name Term s, t, u
%runElab deriveIndexed "Term" [Eq, Ord, Show]


export
Located (Term n) where
  (F x loc).loc                      = loc
  (B i loc).loc                      = loc
  (Lam x body loc).loc               = loc
  (App fun arg loc).loc              = loc
  (Pair fst snd loc).loc             = loc
  (Fst pair loc).loc                 = loc
  (Snd pair loc).loc                 = loc
  (Tag tag loc).loc                  = loc
  (CaseEnum tag cases loc).loc       = loc
  (Absurd loc).loc                   = loc
  (Zero loc).loc                     = loc
  (Succ nat loc).loc                 = loc
  (CaseNat nat zer x ih suc loc).loc = loc
  (Erased loc).loc                   = loc


public export
data Definition = ErasedDef | KeptDef (Term 0)

public export
0 Definitions : Type
Definitions = SortedMap Name Definition


parameters {opts : LayoutOpts}
  export
  prettyTerm : BContext n -> Term n -> Eff Pretty (Doc opts)

  export
  prettyArg : BContext n -> Term n -> Eff Pretty (Doc opts)
  prettyArg xs arg = withPrec Arg $ prettyTerm xs arg

  export
  prettyApp' : Context' BindName n -> Doc opts -> Term n -> Eff Pretty (Doc opts)
  prettyApp' xs fun arg =
    parensIfM App =<< do
      arg <- prettyArg xs arg
      pure $ sep [fun, arg]

  export
  prettyApp : Context' BindName n -> Term n -> Term n -> Eff Pretty (Doc opts)
  prettyApp xs fun arg = prettyApp' xs !(prettyArg xs fun) arg

  public export
  PrettyCaseArm : Nat -> Type
  PrettyCaseArm n = Exists $ \s => (Vect s BindName, Term (s + n))

  export %inline
  caseArm : Vect s BindName -> Term (s + n) -> PrettyCaseArm n
  caseArm xs t = Evidence _ (xs, t)

  export
  prettyCase : Context' BindName n ->
               (a -> Eff Pretty (Doc opts)) ->
               Term n -> List (a, PrettyCaseArm n) ->
               Eff Pretty (Doc opts)
  prettyCase xs f head arms =
    parensIfM Outer =<< Prelude.do
      header <- hsep <$> sequence [caseD, prettyTerm xs head, ofD]
      cases  <- for arms $ \(lhs, (Evidence s (ys, rhs))) => do
        lhs <- hsep <$> sequence [f lhs, darrowD]
        rhs <- withPrec Outer $ prettyTerm (xs <>< ys) rhs
        hangDSingle lhs rhs
      body <- braces $ separateLoose !semiD cases
      pure $ sep [header, body]

  prettyTerm _  (F x _) = prettyFree x
  prettyTerm xs (B i _) = prettyTBind $ xs !!! i
  prettyTerm xs (Lam x body _) =
    parensIfM Outer =<< do
      header <- hsep <$> sequence [lamD, prettyTBind x, darrowD]
      body   <- withPrec Outer $ prettyTerm (xs :< x) body
      hangDSingle header body
  prettyTerm xs (App fun arg _) = prettyApp xs fun arg
  prettyTerm xs (Pair fst snd _) =
    parens =<< separateTight !commaD <$>
      sequence {t = List} [prettyTerm xs fst, prettyTerm xs snd]
  prettyTerm xs (Fst pair _) = prettyApp' xs !fstD pair
  prettyTerm xs (Snd pair _) = prettyApp' xs !sndD pair
  prettyTerm xs (Tag tag _) = prettyTag tag
  prettyTerm xs (CaseEnum tag cases _) = assert_total
    prettyCase xs prettyTag tag $ map (mapSnd $ caseArm []) cases
  prettyTerm xs (Absurd _) = hl Syntax "absurd"
  prettyTerm xs (Zero _) = zeroD
  prettyTerm xs (Succ nat _) = prettyApp' xs !succD nat
  prettyTerm xs (CaseNat nat zer x ih suc _) = assert_total
    prettyCase xs pure nat
      [(!zeroD,  caseArm []      zer),
       (!sucPat, caseArm [x, ih] suc)]
    where
      sucPat = separateTight {t = List} !commaD <$>
               sequence [[|succD <++> prettyTBind x|], prettyTBind ih]
  prettyTerm _ (Erased _) =
    hl Syntax =<< ifUnicode "⌷" "[]"

  export
  prettyDef : Name -> Definition -> Eff Pretty (Maybe (Doc opts))
  prettyDef _    ErasedDef     = [|Nothing|]
  prettyDef name (KeptDef rhs) = map Just $ do
    name <- prettyFree name
    eq   <- cstD
    rhs  <- prettyTerm [<] rhs
    hangDSingle (name <++> eq) rhs


public export
USubst : Nat -> Nat -> Type
USubst = Subst Term


public export FromVar Term where fromVarLoc = B


public export
CanSubstSelf Term where
  s // th = case s of
    F x loc =>
      F x loc
    B i loc =>
      getLoc th i loc
    Lam x body loc =>
      Lam x (assert_total $ body // push th) loc
    App fun arg loc =>
      App (fun // th) (arg // th) loc
    Pair fst snd loc =>
      Pair (fst // th) (snd // th) loc
    Fst pair loc =>
      Fst (pair // th) loc
    Snd pair loc =>
      Snd (pair // th) loc
    Tag tag loc =>
      Tag tag loc
    CaseEnum tag cases loc =>
      CaseEnum (tag // th) (map (assert_total mapSnd (// th)) cases) loc
    Absurd loc =>
      Absurd loc
    Zero loc =>
      Zero loc
    Succ nat loc =>
      Succ (nat // th) loc
    CaseNat nat zer x ih suc loc =>
      CaseNat (nat // th) (zer // th)
              x ih (assert_total $ suc // pushN 2 th) loc
    Erased loc =>
      Erased loc

public export
subN : SnocVect s (Term n) -> Term (s + n) -> Term n
subN th t = t // fromSnocVect th

public export
sub1 : Term n -> Term (S n) -> Term n
sub1 e = subN [< e]