content: {-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}
  module CTT where
  
  
  import Prelude hiding ((<>))
  import Control.Applicative
  import Data.List
  import Data.Maybe
  import qualified Data.Map as Map
  import Text.PrettyPrint as PP
  import qualified Data.Set as Set
  
  import Cartesian
  
  --------------------------------------------------------------------------------
  -- | Terms
  
  data Loc = Loc { locFile :: String
                 , locPos  :: (Int,Int) }
    deriving Eq
  
  type Ident  = String
  type LIdent = String
  
  -- Telescope (x1 : A1) .. (xn : An)
  type Tele   = [(Ident,Ter)]
  
  data Label = OLabel LIdent Tele -- Object label
             | PLabel LIdent Tele [Name] (System Ter) -- Path label
    deriving (Eq,Show)
  
  -- OBranch of the form: c x1 .. xn -> e
  -- PBranch of the form: c x1 .. xn i1 .. im -> e
  data Branch = OBranch LIdent [Ident] Ter
              | PBranch LIdent [Ident] [Name] Ter
    deriving (Eq,Show)
  
  -- Declarations: x : A = e
  -- A group of mutual declarations is identified by its location. It is used to
  -- speed up the Eq instance for Ctxt.
  type Decl  = (Ident,(Ter,Ter))
  data Decls = MutualDecls Loc [Decl]
             | OpaqueDecl Ident
             | TransparentDecl Ident
             | TransparentAllDecl
             deriving Eq
  
  declIdents :: [Decl] -> [Ident]
  declIdents decls = [ x | (x,_) <- decls ]
  
  declTers :: [Decl] -> [Ter]
  declTers decls = [ d | (_,(_,d)) <- decls ]
  
  declTele :: [Decl] -> Tele
  declTele decls = [ (x,t) | (x,(t,_)) <- decls ]
  
  declDefs :: [Decl] -> [(Ident,Ter)]
  declDefs decls = [ (x,d) | (x,(_,d)) <- decls ]
  
  labelTele :: Label -> (LIdent,Tele)
  labelTele (OLabel c ts) = (c,ts)
  labelTele (PLabel c ts _ _) = (c,ts)
  
  labelName :: Label -> LIdent
  labelName = fst . labelTele
  
  labelTeles :: [Label] -> [(LIdent,Tele)]
  labelTeles = map labelTele
  
  lookupLabel :: LIdent -> [Label] -> Maybe Tele
  lookupLabel x xs = lookup x (labelTeles xs)
  
  lookupPLabel :: LIdent -> [Label] -> Maybe (Tele,[Name],System Ter)
  lookupPLabel x xs = listToMaybe [ (ts,is,es) | PLabel y ts is es <- xs, x == y ]
  
  branchName :: Branch -> LIdent
  branchName (OBranch c _ _) = c
  branchName (PBranch c _ _ _) = c
  
  lookupBranch :: LIdent -> [Branch] -> Maybe Branch
  lookupBranch _ []      = Nothing
  lookupBranch x (b:brs) = case b of
    OBranch c _ _   | x == c    -> Just b
                    | otherwise -> lookupBranch x brs
    PBranch c _ _ _ | x == c    -> Just b
                    | otherwise -> lookupBranch x brs
  
  -- Terms
  data Ter = Pi Ter
           | App Ter Ter
           | Lam Ident Ter Ter
           | Where Ter Decls
           | Var Ident
           | U
             -- Sigma types:
           | Sigma Ter
           | Pair Ter Ter
           | Fst Ter
           | Snd Ter
             -- constructor c Ms
           | Con LIdent [Ter]
           | PCon LIdent Ter [Ter] [II] -- c A ts phis (A is the data type)
             -- branches c1 xs1  -> M1,..., cn xsn -> Mn
           | Split Ident Loc Ter [Branch]
             -- labelled sum c1 A1s,..., cn Ans (assumes terms are constructors)
           | Sum Loc Ident [Label] -- TODO: should only contain OLabels
           | HSum Loc Ident [Label]
             -- undefined and holes
           | Undef Loc Ter -- Location and type
           | Hole Loc
             -- Path and line types
           | PathP Ter Ter Ter
           | LineP Ter
           | PLam Name Ter
           | AppII Ter II
             -- Coe
           | Coe II II Ter Ter
             -- Homogeneous Kan composition
           | HCom II II Ter (System Ter) Ter
             -- Heterogeneous Kan composition (derived)
           | Com II II Ter (System Ter) Ter
  
             -- V-types
           | V II Ter Ter Ter -- V r A B E    (where E : A ~= B)
           | Vin II Ter Ter -- Vin r M N      (where M : A and N : B)
           | Vproj II Ter Ter Ter Ter -- Vproj r O A B E    (where O : V r A B E)
  
             -- Universes
           | Box II II (System Ter) Ter
           | Cap II II (System Ter) Ter
  
             -- Glue
           -- | Glue Ter (System Ter)
           -- | GlueElem Ter (System Ter)
           -- | UnGlueElem Ter Ter (System Ter)
    deriving Eq
  
  -- For an expression t, returns (u,ts) where u is no application and t = u ts
  unApps :: Ter -> (Ter,[Ter])
  unApps = aux []
    where aux :: [Ter] -> Ter -> (Ter,[Ter])
          aux acc (App r s) = aux (s:acc) r
          aux acc t         = (t,acc)
  
  mkApps :: Ter -> [Ter] -> Ter
  mkApps (Con l us) vs = Con l (us ++ vs)
  mkApps t ts          = foldl App t ts
  
  mkWheres :: [Decls] -> Ter -> Ter
  mkWheres []     e = e
  mkWheres (d:ds) e = Where (mkWheres ds e) d
  
  --------------------------------------------------------------------------------
  -- | Values
  
  data Val = VU
           | Ter Ter Env
           | VPi Val Val
           | VSigma Val Val
           | VPair Val Val
           | VCon LIdent [Val]
           | VPCon LIdent Val [Val] [II]
  
             -- Path values
           | VPathP Val Val Val
           | VLineP Val
           | VPLam Name Val
  
             -- Homogeneous composition; the type is constant
           | VHCom II II Val (System Val) Val
  
             -- Coe
           | VCoe II II Val Val
  
             -- V-types values
           | VV Name Val Val Val -- V i A B E    (where E : A ~= B)
           | VVin Name Val Val -- Vin i M N      (where M : A and N : B)
           | VVproj Name Val Val Val Val -- Vproj i O A B E    (where O : V i A B E)
  
             -- Glue values
           -- | VGlue Val (System Val)
           -- | VGlueElem Val (System Val)
           -- | VUnGlueElem Val Val (System Val)  -- unglue u A [phi -> (T,w)]
  
           -- Universe values
           | VHComU II II (System Val) Val -- r s bs a
           | VBox II II (System Val) Val -- r s ns m
           | VCap II II (System Val) Val -- r s bs m
  
             -- Neutral values:
           | VVar Ident Val
           | VOpaque Ident Val
           | VFst Val
           | VSnd Val
           | VSplit Val Val
           | VApp Val Val
           | VAppII Val II
           | VLam Ident Val Val
           -- | VUnGlueElemU Val Val (System Val)
    deriving Eq
  
  isNeutral :: Val -> Bool
  isNeutral v = case v of
    Ter Undef{} _     -> True
    Ter Hole{} _      -> True
    VVar{}            -> True
    VOpaque{}         -> True
    VHCom{}           -> True
    VCoe{}            -> True
    VFst{}            -> True
    VSnd{}            -> True
    VSplit{}          -> True
    VApp{}            -> True
    VAppII{}          -> True
    -- VUnGlueElemU{} -> True
    -- VUnGlueElem{}  -> True
    VCap{}            -> True
    VVproj{}          -> True
    _                 -> False
  
  isNeutralSystem :: System Val -> Bool
  isNeutralSystem (Sys xs) = any isNeutral (Map.elems xs)
  isNeutralSystem (Triv a) = isNeutral a
  
  -- isNeutralPath :: Val -> Bool
  -- isNeutralPath (VPath _ v) = isNeutral v
  -- isNeutralPath _ = True
  
  mkVar :: Int -> String -> Val -> Val
  mkVar k x = VVar (x ++ show k)
  
  mkVarNice :: [String] -> String -> Val -> Val
  mkVarNice xs x = VVar (head (ys \\ xs))
    where ys = x:map (\n -> x ++ show n) [0..]
  
  unCon :: Val -> [Val]
  unCon (VCon _ vs) = vs
  unCon v           = error $ "unCon: not a constructor: " ++ show v
  
  isCon :: Val -> Bool
  isCon VCon{} = True
  isCon _      = False
  
  -- Constant path: <_> v
  constPath :: Val -> Val
  constPath = VPLam (N "_")
  
  
  --------------------------------------------------------------------------------
  -- | Environments
  
  data Ctxt = Empty
            | Upd Ident Ctxt
            | Sub Name Ctxt
            | Def Loc [Decl] Ctxt
    deriving (Show)
  
  instance Eq Ctxt where
      c == d = case (c, d) of
          (Empty, Empty)              -> True
          (Upd x c', Upd y d')        -> x == y && c' == d'
          (Sub i c', Sub j d')        -> i == j && c' == d'
          (Def m xs c', Def n ys d')  -> (m == n || xs == ys) && c' == d'
              -- Invariant: if two declaration groups come from the same
              -- location, they are equal and their contents are not compared.
          _                           -> False
  
  -- The Idents and Names in the Ctxt refer to the elements in the two
  -- lists. This is more efficient because acting on an environment now
  -- only need to affect the lists and not the whole context.
  -- The last list is the list of opaque names
  newtype Env = Env (Ctxt,[Val],[II],Nameless (Set.Set Ident))
    deriving (Eq)
  
  emptyEnv :: Env
  emptyEnv = Env (Empty,[],[],Nameless Set.empty)
  
  def :: Decls -> Env -> Env
  def (MutualDecls m ds) (Env (rho,vs,fs,Nameless os)) =
    Env (Def m ds rho,vs,fs,Nameless (os Set.\\ Set.fromList (declIdents ds)))
  def (OpaqueDecl n) (Env (rho,vs,fs,Nameless os)) = Env (rho,vs,fs,Nameless (Set.insert n os))
  def (TransparentDecl n) (Env (rho,vs,fs,Nameless os)) = Env (rho,vs,fs,Nameless (Set.delete n os))
  def TransparentAllDecl (Env (rho,vs,fs,Nameless os)) = Env (rho,vs,fs,Nameless Set.empty)
  
  defWhere :: Decls -> Env -> Env
  defWhere (MutualDecls m ds) (Env (rho,vs,fs,Nameless os)) =
    Env (Def m ds rho,vs,fs,Nameless (os Set.\\ Set.fromList (declIdents ds)))
  defWhere (OpaqueDecl _) rho = rho
  defWhere (TransparentDecl _) rho = rho
  defWhere TransparentAllDecl rho = rho
  
  sub :: (Name,II) -> Env -> Env
  sub (i,phi) (Env (rho,vs,fs,os)) = Env (Sub i rho,vs,phi:fs,os)
  
  upd :: (Ident,Val) -> Env -> Env
  upd (x,v) (Env (rho,vs,fs,Nameless os)) = Env (Upd x rho,v:vs,fs,Nameless (Set.delete x os))
  
  upds :: [(Ident,Val)] -> Env -> Env
  upds xus rho = foldl (flip upd) rho xus
  
  updsTele :: Tele -> [Val] -> Env -> Env
  updsTele tele vs = upds (zip (map fst tele) vs)
  
  subs :: [(Name,II)] -> Env -> Env
  subs iphis rho = foldl (flip sub) rho iphis
  
  mapEnv :: (Val -> Val) -> (II -> II) -> Env -> Env
  mapEnv f g (Env (rho,vs,fs,os)) = Env (rho,map f vs,map g fs,os)
  
  valAndIIOfEnv :: Env -> ([Val],[II])
  valAndIIOfEnv (Env (_,vs,fs,_)) = (vs,fs)
  
  valOfEnv :: Env -> [Val]
  valOfEnv = fst . valAndIIOfEnv
  
  formulaOfEnv :: Env -> [II]
  formulaOfEnv = snd . valAndIIOfEnv
  
  domainEnv :: Env -> [Name]
  domainEnv (Env (rho,_,_,_)) = domCtxt rho
    where domCtxt rho = case rho of
            Empty      -> []
            Upd _ e    -> domCtxt e
            Def _ ts e -> domCtxt e
            Sub i e    -> i : domCtxt e
  
  -- Extract the context from the environment, used when printing holes
  contextOfEnv :: Env -> [String]
  contextOfEnv rho = case rho of
    Env (Empty,_,_,_)               -> []
    Env (Upd x e,VVar n t:vs,fs,os) -> (n ++ " : " ++ show t) : contextOfEnv (Env (e,vs,fs,os))
    Env (Upd x e,v:vs,fs,os)        -> (x ++ " = " ++ show v) : contextOfEnv (Env (e,vs,fs,os))
    Env (Def _ _ e,vs,fs,os)        -> contextOfEnv (Env (e,vs,fs,os))
    Env (Sub i e,vs,phi:fs,os)      -> (show i ++ " = " ++ show phi) : contextOfEnv (Env (e,vs,fs,os))
  
  --------------------------------------------------------------------------------
  -- | Pretty printing
  
  instance Show Env where
    show = render . showEnv True
  
  showEnv :: Bool -> Env -> Doc
  showEnv b e =
    let -- This decides if we should print "x = " or not
        names x = if b then text x <+> equals else PP.empty
        par   x = if b then parens x else x
        com     = if b then comma else PP.empty
        showEnv1 e = case e of
          Env (Upd x env,u:us,fs,os)   ->
            showEnv1 (Env (env,us,fs,os)) <+> names x <+> showVal1 u <> com
          Env (Sub i env,us,phi:fs,os) ->
            showEnv1 (Env (env,us,fs,os)) <+> names (show i) <+> text (show phi) <> com
          Env (Def _ _ env,vs,fs,os)   -> showEnv1 (Env (env,vs,fs,os))
          _                            -> showEnv b e
    in case e of
      Env (Empty,_,_,_)            -> PP.empty
      Env (Def _ _ env,vs,fs,os)   -> showEnv b (Env (env,vs,fs,os))
      Env (Upd x env,u:us,fs,os)   ->
        par $ showEnv1 (Env (env,us,fs,os)) <+> names x <+> showVal1 u
      Env (Sub i env,us,phi:fs,os) ->
        par $ showEnv1 (Env (env,us,fs,os)) <+> names (show i) <+> text (show phi)
  
  instance Show Loc where
    show = render . showLoc
  
  showLoc :: Loc -> Doc
  showLoc (Loc name (i,j)) = text (show (i,j) ++ " in " ++ name)
  
  showII :: II -> Doc
  showII = text . show
  
  instance Show Ter where
    show = render . showTer
  
  showTer :: Ter -> Doc
  showTer v = case v of
    U                    -> char 'U'
    App e0 e1            -> showTer e0 <+> showTer1 e1
    Pi e0                -> text "Pi" <+> showTer e0
    Lam x t e            ->
      char '\\' <> parens (text x <+> colon <+> showTer t) <+> text " ->" <+> showTer e
    Fst e                -> showTer1 e <> text ".1"
    Snd e                -> showTer1 e <> text ".2"
    Sigma e0             -> text "Sigma" <+> showTer1 e0
    Pair e0 e1           -> parens (showTer e0 <> comma <> showTer e1)
    Where e d            -> showTer e <+> text "where" <+> showDecls d
    Var x                -> text x
    Con c es             -> text c <+> showTers es
    PCon c a es phis     ->
      text c <+> braces (showTer a) <+> showTers es <+> hsep (map ((char '@' <+>) . showII) phis)
    Split f _ _ _        -> text f
    Sum _ n _            -> text n
    HSum _ n _           -> text n
    Undef{}              -> text "undefined"
    Hole{}               -> text "?"
    PathP e0 e1 e2       -> text "PathP" <+> showTers [e0,e1,e2]
    LineP e              -> text "LineP" <+> showTer e
    PLam i e             -> char '<' <> text (show i) <> char '>' <+> showTer e
    AppII e phi          -> showTer1 e <+> char '@' <+> showII phi
    Com r s a ts t      ->
      text "com" <+> showII r <> text "->" <> showII s <+> showTer1 a <+> text (show ts) <+> showTer1 t
    HCom r s a ts t      ->
      text "hcom" <+> showII r <> text "->" <> showII s <+> showTer1 a <+> text (show ts) <+> showTer1 t
    Coe r s e t0         ->
      text "coe" <+> showII r <> text "->" <> showII s <+> showTer1 e <+> showTer1 t0
    -- Comp e t ts       -> text "comp" <+> showTers [e,t] <+> text (show ts)
    V r a b e            -> text "V" <+> showII r <+> showTers [a,b,e]
    Vin r m n            -> text "Vin" <+> showII r <+> showTers [m,n]
    Vproj r o a b e      -> text "Vproj" <+> showII r <+> showTers [o,a,b,e]
    Box r s ts t ->
      text "box" <+> showII r <> text "->" <> showII s <+> text (show ts) <+> showTer1 t
    Cap r s ts t ->
      text "cap" <+> showII r <> text "<-" <> showII s <+> text (show ts) <+> showTer1 t
    -- Glue a ts         -> text "Glue" <+> showTer1 a <+> text (show ts)
    -- GlueElem a ts     -> text "glue" <+> showTer1 a <+> text (show ts)
    -- UnGlueElem a b ts -> text "unglue" <+> showTers [a,b] <+> text (show ts)
  
  showTers :: [Ter] -> Doc
  showTers = hsep . map showTer1
  
  showTer1 :: Ter -> Doc
  showTer1 t = case t of
    U        -> char 'U'
    Con c [] -> text c
    Var{}    -> showTer t
    Undef{}  -> showTer t
    Hole{}   -> showTer t
    Split{}  -> showTer t
    Sum{}    -> showTer t
    HSum{}   -> showTer t
    Fst{}    -> showTer t
    Snd{}    -> showTer t
    _        -> parens (showTer t)
  
  showDecls :: Decls -> Doc
  showDecls (MutualDecls _ defs) =
    hsep $ punctuate comma
    [ text x <+> equals <+> showTer d | (x,(_,d)) <- defs ]
  showDecls (OpaqueDecl i) = text "opaque" <+> text i
  showDecls (TransparentDecl i) = text "transparent" <+> text i
  showDecls TransparentAllDecl = text "transparent_all"
  
  instance Show Val where
    show = render . showVal
  
  showVal :: Val -> Doc
  showVal v = case v of
    VU                     -> char 'U'
    Ter t@Sum{} rho        -> showTer t <+> showEnv False rho
    Ter t@HSum{} rho       -> showTer t <+> showEnv False rho
    Ter t@Split{} rho      -> showTer t <+> showEnv False rho
    Ter t rho              -> showTer1 t <+> showEnv True rho
    VCon c us              -> text c <+> showVals us
    VPCon c a us phis      -> text c <+> braces (showVal a) <+> showVals us <+> hsep (map ((char '@' <+>) . showII) phis)
    VHCom r s v0 vs v1     -> text "hcom" <+> showII r <> text "->" <> showII s <+> showVal1 v0 <+> text (show vs) <+> showVal1 v1
    VCoe r s u v0          -> text "coe" <+> showII r <> text "->" <> showII s <+> showVal1 u <+> showVal1 v0
    VPi a l@(VLam x t b)
      | "_" `isPrefixOf` x -> showVal1 a <+> text "->" <+> showVal1 b
      | otherwise          -> char '(' <> showLam v
    VPi a b                -> text "Pi" <+> showVals [a,b]
    VPair u v              -> parens (showVal u <> comma <> showVal v)
    VSigma u v             -> text "Sigma" <+> showVals [u,v]
    VApp u v               -> showVal u <+> showVal1 v
    VLam{}                 -> text "\\(" <> showLam v
    VPLam{}                -> char '<' <> showPLam v
    VSplit u v             -> showVal u <+> showVal1 v
    VVar x _               -> text x
    VOpaque x _            -> text ('#':x)
    VFst u                 -> showVal1 u <> text ".1"
    VSnd u                 -> showVal1 u <> text ".2"
    VPathP v0 v1 v2        -> text "PathP" <+> showVals [v0,v1,v2]
    VLineP v               -> text "LineP" <+> showVal v
    VAppII v phi           -> showVal v <+> char '@' <+> showII phi
    VV i a b e             -> text "V" <+> text (show i) <+> showVals [a,b,e]
    VVin i m n             -> text "Vin" <+> text (show i) <+> showVals [m,n]
    VVproj i o a b e       -> text "Vproj" <+> text (show i) <+> showVals [o,a,b,e]
    VBox r s ts t          ->
      text "box" <+> showII r <> text "->" <> showII s <+> text (show ts) <+> showVal1 t
    VCap r s ts t          ->
      text "cap" <+> showII r <> text "<-" <> showII s <+> text (show ts) <+> showVal1 t
    VHComU r s ts t       ->
      text "hcomp U" <+> showII r <> text "->" <> showII s <+> text (show ts) <+> showVal1 t
    -- VGlue a ts          -> text "Glue" <+> showVal1 a <+> text (show ts)
    -- VGlueElem a ts      -> text "glue" <+> showVal1 a <+> text (show ts)
    -- VUnGlueElem v a ts  -> text "unglue" <+> showVals [v,a] <+> text (show ts)
    -- VUnGlueElemU v b es -> text "unglue U" <+> showVals [v,b] <+> text (show es)
  
  showPLam :: Val -> Doc
  showPLam e = case e of
    VPLam i a@VPLam{} -> text (show i) <+> showPLam a
    VPLam i a         -> text (show i) <> char '>' <+> showVal a
    _                 -> showVal e
  
  -- Merge lambdas of the same type
  showLam :: Val -> Doc
  showLam e = case e of
    VLam x t a@(VLam _ t' _)
      | t == t'   -> text x <+> showLam a
      | otherwise ->
        text x <+> colon <+> showVal t <> char ')' <+> text "->" <+> showVal a
    VPi _ (VLam x t a@(VPi _ (VLam _ t' _)))
      | t == t'   -> text x <+> showLam a
      | otherwise ->
        text x <+> colon <+> showVal t <> char ')' <+> text "->" <+> showVal a
    VLam x t e         ->
      text x <+> colon <+> showVal t <> char ')' <+> text "->" <+> showVal e
    VPi _ (VLam x t e) ->
      text x <+> colon <+> showVal t <> char ')' <+> text "->" <+> showVal e
    _ -> showVal e
  
  showVal1 :: Val -> Doc
  showVal1 v = case v of
    VU                -> showVal v
    VCon c []         -> showVal v
    VVar{}            -> showVal v
    VFst{}            -> showVal v
    VSnd{}            -> showVal v
    Ter t rho | isEmpty (showEnv False rho) -> showTer1 t
    _                 -> parens (showVal v)
  
  showVals :: [Val] -> Doc
  showVals = hsep . map showVal1
  