Terms.hs

{-# LANGUAGE GADTs, KindSignatures, OverloadedStrings, TypeSynonymInstances, RankNTypes, PatternGuards #-}
module Terms where

import Prelude hiding (length,elem,foldl,all,concatMap,and,drop,concat)
import Basics
import Display
import Data.Foldable
import Control.Arrow (first, second)
import Data.Sequence 
import Options

type NF = Term
type Neutral = Term

-- | Representation of terms.
data Term :: * where
     Star :: Position -> Sort -> NF
     
     Pi  :: Ident -> NF -> NF -> NF
     Lam :: Ident -> NF -> NF -> NF 
     App :: Neutral -> NF -> Neutral
     
     Sigma  :: Position -> [(String,NF)] -> NF 
     Pair   :: Position -> [(String,NF)] -> NF -- Note: Unlike Sigma, Pairs do not bind variables
     Proj   :: Neutral -> String -> Neutral      
              
     Fin :: Position -> [String] -> NF
     Tag :: Position -> String -> NF
     Cas :: Position -> [(String, NF)] -> NF

     V :: Position -> Int ->  -- ^ deBruijn index 
          Neutral
     Hole :: Position -> String -> Neutral
     
     This :: NF -- ^ reference to the module currently type-checked
  
     Ann :: Neutral -> NF -> NF -- ^ type annotation     
  deriving (Show)

termPosition t = case t of
   (Hole p _) -> p
   (Star p _) -> p
   (V p _) -> p
   (Pi i _ _) -> identPosition i
   (Sigma p _) -> p
   (Lam i _ _) -> identPosition i
   (Pair p _ ) -> p
   (App x y) -> s x
   (Proj x _) -> s x
   (Ann x _) -> s x 
   (Fin x _) -> x
   (Tag x _) -> x
   (Cas x _) -> x
   _ -> dummyPosition
  where s = termPosition

type Subst = [NF]

var = V dummyPosition
hole = Hole dummyPosition

-- | The identity substitution
identity = map var [0..]

subst0 :: NF -> Subst
subst0 u = u:identity

-- | Weakening
wkn :: Int -> Subst
wkn n = map var [n..]

wk = wkn 1

-- | Hereditary substitution application
apply :: Subst -> Term -> NF
apply f t = case t of
  Star p x -> Star p x
  Lam i ty bo -> Lam i (s ty) (s' bo)
  Pair i fs -> Pair i (map (second s) fs)
  Pi i a b -> Pi i (s a) (s' b)
  Sigma i [] -> Sigma i []
  Sigma i ((f,x):xs) -> let Sigma _ xs' = s' (Sigma i xs) in Sigma i ((f,s x):xs')
  (App a b) -> app (s a) (s b)
  (Proj x k) -> proj (s x) k 
  Hole p x -> Hole p x
  V _ x -> f !! x
  This -> This
  Ann x t -> ann (s x) (s t)
  
  Cas p cs -> Cas p (map (second s) cs)
  Fin p x -> Fin p x
  Tag p x -> Tag p x
 where s' = apply (var 0 : wk ∘ f)
       s  = apply f

(∙) = apply
σ ∘ ρ = map (apply σ) ρ

ann x t = Ann x t

-- | Application that computes
app :: NF -> NF -> NF 
app (Lam i _ bo) u = subst0 u ∙ bo
app (Cas _ cs) (Tag _ t) | Just x <- lookup t cs = x
app n            u = App n u

-- | Projection that computes
proj :: NF -> String -> NF
proj (Pair _ fs) f | Just x <- lookup f fs = x
proj x k = Proj x k 

-----------------------------------
-- Display

dec xs = [ x - 1 | x <- xs, x > 0]

freeVars :: Term -> [Int]
freeVars (Pi _ a b) = freeVars a <> (dec $ freeVars b)
freeVars (Sigma _ []) = []
freeVars (Sigma p ((_,x):xs)) = freeVars x <> (dec $ freeVars (Sigma p xs))
freeVars (V _ x) = [x]
freeVars (App a b) = freeVars a <> freeVars b
freeVars (Lam _ ty b) = freeVars ty <> (dec $ freeVars b)
freeVars (Star _ _) = mempty
freeVars (Hole _ _) = mempty
freeVars (Pair _ xs) = concatMap (freeVars . snd) xs
freeVars (Proj x _) = freeVars x
freeVars This = []
freeVars (Ann x t) = freeVars x <> freeVars t
freeVars (Fin _ _) = []
freeVars (Tag _ _) = []
freeVars (Cas _ cs) = concatMap (freeVars . snd) cs

iOccursIn :: Int -> Term -> Bool
iOccursIn x t = x `elem` (freeVars t)

allocName :: DisplayContext -> Ident -> Ident
allocName g s 
  | s `elem` g = allocName g (modId (++ "'") s)
  | otherwise = s

cPrint :: Int -> DisplayContext -> Term -> Doc
cPrint p ii (Hole _ x) = text x
cPrint p ii (Star _ i) = pretty i
cPrint p ii (V _ k) 
  | k < 0 || k >= length ii  = text "<global " <> pretty (k - length ii) <> ">"
  | otherwise = pretty (ii `index` k) 
cPrint p ii (Proj x f)      = cPrint p ii x <> "." <> text f 
cPrint p ii t@(App _ _)     = let (fct,args) = nestedApp t in 
                                 parensIf (p > 3) (cPrint 3 ii fct <+> sep [ cPrint 4 ii a | a <- args]) 
cPrint p ii t@(Pi _ _ _)    = parensIf (p > 1) (printBinders "→" ii mempty $ nestedPis t)
cPrint p ii t@(Sigma _ _) = parensIf (p > 1) (printBinders "×" ii mempty $ nestedSigmas t)
cPrint p ii (t@(Lam _ _ _)) = parensIf (p > 1) (nestedLams ii mempty t)
cPrint p ii (Pair _ fs) = parensIf (p > (-1)) (sep (punctuate comma [text name <+> text "=" <+> cPrint 0 ii x | (name,x) <- fs ]))
cPrint p ii This = "this" 
cPrint p ii (Ann x t) = parensIf (p > 0) (cPrint 0 ii x <+> ":" <+> cPrint 0 ii t)
cPrint p ii (Fin _ ts) = "[" <> sep (punctuate comma (map text ts)) <> "]"
cPrint p ii (Tag _ t) = "'" <> text t
cPrint p ii (Cas _ cs) = "case {" <> sep (punctuate ";" [text c <> "↦" <> cPrint 0 ii a | (c,a) <- cs]) <> "}"

-- FIXME: should remember the variable names in the substitution
dispEnv :: DisplayContext -> [(Int,NF)] -> [Doc]
dispEnv ii g = [pretty i <> "↦" <> cPrint 0 ii v | (i,v) <- g]

nestedPis  :: NF -> ([(Ident,Bool,NF)], NF)
nestedPis (Pi i a b) = (first ([(i,0 `iOccursIn` b,a)] ++)) (nestedPis b)
nestedPis x = ([],x)

nestedSigmas  :: NF -> ([(Ident,Bool,NF)], NF)
nestedSigmas (Sigma p ((i,x):xs)) = (first ([(synthId i,0 `iOccursIn` Sigma p xs,x)] ++)) (nestedSigmas (Sigma p xs))
nestedSigmas (Sigma p []) = ([],Hole p "⊤")

printBinders :: Doc -> DisplayContext -> Seq Doc -> ([(Ident,Bool,NF)], NF) -> Doc
printBinders sep ii xs (((x,occurs,a):pis),b) = printBinders sep (i <| ii) (xs |> (printBind' ii i occurs a <+> sep)) (pis,b)
        where i = allocName ii x
printBinders _ ii xs ([],b)                 = sep $ toList $ (xs |> cPrint 1 ii b) 


nestedLams :: DisplayContext -> Seq Doc -> Term -> Doc
nestedLams ii xs (Lam x ty c) = nestedLams (i <| ii) (xs |> parens (pretty i <+> ":" <+> cPrint 0 ii ty)) c
                                  where i = allocName ii x
nestedLams ii xs t         = (text "\\ " <> (sep $ toList $ (xs |> "->")) <> " " <> nest 3 (cPrint 0 ii t))

printBind' ii name occurs d = case not (isDummyId name) || occurs of
                  True -> parens (pretty name <+> ":" <+> cPrint 0 ii d)
                  False -> cPrint 2 ii d
                  
nestedApp :: Neutral -> (Neutral,[NF])
nestedApp (App f a) = (second (++ [a])) (nestedApp f)
nestedApp t = (t,[])

prettyTerm = cPrint (-100)

instance Pretty Term where
    pretty = prettyTerm mempty