Can typecheck universe-polymorphic id

3 files changed, 195 insertions, 86 deletions

diff --git a/src/AST.hs b/src/AST.hs
index ba018f5..d828e26 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -1,15 +1,15 @@
 module AST where
 
-import Data.Map.Strict (Map)
+import Numeric.Natural
 
+import Data.Map.Strict (Map)
 
-data Nat = Z | S Nat
-  deriving (Show)
 
 type Name = String
 
 data Term
   = TSet Term            -- ^ The n-th universe (n : Level)
+  | TSetw Natural        -- ^ The n-th ω-universe
   | TVar Name            -- ^ variable
   | TPi Name Term Term   -- ^ Pi: (x : A) -> B
   | TLam Name Term Term  -- ^ λ(x : A). B
@@ -36,7 +36,7 @@ data Term
 data Definition = Definition Name (OfType Term Term)
   deriving (Show)
 
-type Env = Map Name Term
+type Env = Map Name (OfType (Maybe Term) Term)
 
 data OfType a b = a :| b
   deriving (Show)
diff --git a/src/Parser.hs b/src/Parser.hs
index 7464d85..fcab060 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -4,7 +4,6 @@ import Control.Monad
 import Data.Char
 import Data.List (foldl')
 import Text.Parsec hiding (token)
-import Text.Parsec.Expr
 
 import AST
 
@@ -28,7 +27,7 @@ pDef = do
   whitespace
   token ":"
   whitespace
-  typ <- pTerm
+  typ <- pTerm 0
   whitespaceTo0
   name2 <- pName
   when (name /= name2) $
@@ -36,27 +35,56 @@ pDef = do
   whitespace
   token "="
   whitespace
-  rhs <- pTerm
+  rhs <- pTerm 0
   return (Definition name (rhs :| typ))
 
-pTerm :: Parser Term
-pTerm = buildExpressionParser optable pTermEAtom
-  where
-    optable =
-      [[Infix (token "->" >> return makePi) AssocRight
-       ,Infix (token "→" >> return makePi) AssocRight]
-      ,[Infix (token ":" >> return (\a b -> TAnnot (a :| b))) AssocNone]]
+-- operator precedences:
+--   0 _:_
+--   1 (_ : _) -> _
+--   2 _,_
+
+pTerm :: Int -> Parser Term
+pTerm p | p <= 0 = do
+  e1 <- pTerm 1
+  whitespace
+  choice
+    [do token ":"
+        whitespace
+        e2 <- pTerm 1
+        return (TAnnot (e1 :| e2))
+    ,return e1]
+
+pTerm 1 = do
+  e1 <- pTerm 2
+  whitespace
+  choice
+    [do (token "->" <|> token "→")
+        (x :| ty) <- case e1 of
+          TAnnot (TVar x :| ty) -> return (x :| ty)
+          _ -> fail $ "Expected (x : type) as left-hand side of (->)"
+        whitespace
+        e2 <- pTerm 1
+        return (TPi x ty e2)
+    ,return e1]
+
+pTerm 2 = do
+  e1 <- pTerm 3
+  whitespace
+  choice
+    [do token ","
+        whitespace
+        e2 <- pTerm 2
+        return (TPair e1 e2)
+    ,return e1]
 
-makePi :: Term -> Term -> Term
-makePi (TAnnot (TVar x :| ty1)) ty2 = TPi x ty1 ty2
-makePi ty1 ty2 = TPi ";_" ty1 ty2
+pTerm _ = pTermEAtom
 
 pTermEAtom :: Parser Term
-pTermEAtom = pLambda <|> pApp
+pTermEAtom = pLambda <|> pSpecialApp <|> pApp
 
 pLambda :: Parser Term
 pLambda = do
-  token "\\" <|> token "λ"
+  void (char '\\') <|> token "λ"
   whitespace
   _ <- char '('
   whitespace
@@ -64,63 +92,103 @@ pLambda = do
   whitespace
   token ":"
   whitespace
-  ty <- pTerm
+  ty <- pTerm 0
   whitespace
   _ <- char ')'
   whitespace
   token "->"
-  rhs <- pTerm
+  whitespace
+  rhs <- pTerm 0
   return (TLam x ty rhs)
 
+pSpecialApp :: Parser Term
+pSpecialApp = pSet <|> pSetw <|> pSigma
+  where
+    pSet :: Parser Term
+    pSet = do
+      token "Set"
+      whitespace
+      e <- pTermEAtom
+      return (TSet e)
+
+    pSetw :: Parser Term
+    pSetw = do
+      token "Setw"
+      whitespace
+      n <- read <$> many1 digit
+      return (TSetw n)
+
+    pSigma :: Parser Term
+    pSigma = do
+      token "Sigma"
+      whitespace
+      _ <- char '('
+      whitespace
+      x <- pName
+      whitespace
+      token ":"
+      whitespace
+      ty <- pTerm 0
+      whitespace
+      _ <- char ')'
+      whitespace
+      rhs <- pTermEAtom
+      return (TSigma x ty rhs)
+
 pApp :: Parser Term
 pApp = do
   e1 <- pAtom
-  args <- many (try (whitespace >> pTerm))
+  args <- many (try (whitespace >> pAtom))
   return (foldl' TApp e1 args)
 
 pAtom :: Parser Term
 pAtom = choice
   [do _ <- char '('
       whitespace
-      e <- pTerm
+      e <- pTerm 0
       whitespace
       _ <- char ')'
       return e
   ,TVar <$> pName]
 
-TODO TPair, TSigma ("%"). The rest are environment entries.
+-- TODO TPair, TSigma ("%"). The rest are environment entries.
 
 pName :: Parser Name
 pName = do
   name <- many1 pNameChar
   guard (name `notElem` keywords)
+  guard (head name /= '\\')  -- lambda syntax
   return name
   where
     keywords = ["=", "->", "→", ":", "?", "\\", "λ", "∀", "..", "..."
-               ,"%"  -- for Sigma; hopefully temporary
-               ,"forall", "data", "import", "in", "let"]
+               ,"forall", "data", "import", "in", "let"
+               ,"Sigma", ","  -- temporary
+               ]
 
 pNameChar :: Parser Char
 pNameChar = satisfy (\c -> not (isSpace c) && c `notElem` ".;{}()@\"")
 
 inlineWhite :: Parser ()
 inlineWhite = do
-  spaces
+  skipMany inlineSpace
   choice
     [inlineComment
     -- TODO: nested {- -} comments
     ,return ()]
 
 inlineComment :: Parser ()
-inlineComment = token "--" >> manyTill anyChar newline >> return ()
+inlineComment = token "--" >> manyTill anyChar (lookAhead newline) >> return ()
 
 -- | Skips to the furthest point from here that can be reached without
 -- consuming any tokens, and without ending on column 0.
 whitespace :: Parser ()
 whitespace = do
   inlineWhite
-  optional $
-    many newline >> (void space <|> inlineComment) >> whitespace
+  optional $ try $
+    many newline >> (inlineSpace <|> inlineComment) >> whitespace
+
+inlineSpace :: Parser ()
+inlineSpace = void (satisfy (\c -> isSpace c && c /= '\n')) <?> "inline whitespace"
 
 -- | Skips as much whitespace as possible, and throws an error if that doesn't
 -- leave us on column 0.
diff --git a/src/TypeCheck.hs b/src/TypeCheck.hs
index 1406a67..6777ec8 100644
--- a/src/TypeCheck.hs
+++ b/src/TypeCheck.hs
@@ -3,7 +3,8 @@
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE DeriveFoldable #-}
-module TypeCheck (typeCheck, typeInfer) where
+{-# LANGUAGE ViewPatterns #-}
+module TypeCheck (checkProgram, typeCheck, typeInfer) where
 
 import Control.Monad
 -- import Control.Monad.Trans.Class
@@ -12,15 +13,24 @@ import Control.Monad
 -- import Data.Foldable (toList)
 -- import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
--- import Numeric.Natural
+import Numeric.Natural
 
 import AST
 
 
-typeCheck :: Env -> Term -> Term -> Either String Term
-typeCheck env typ term =
-  runTM (check env term typ)
-  -- runTM (check env term typ) >>= \case
+checkProgram :: Env -> [Definition] -> Either String Env
+checkProgram = foldM $ \env (Definition name jud@(_ :| ty)) -> do
+  ty' <- runTM (inferW env ty) >>= \case
+    ty' :| TSet _ -> return ty'
+    ty' :| TSetw _ -> return ty'
+    _ :| kind -> Left $ "Kind of declared type is not Set i or Setw i, but: " ++ show kind
+  term <- typeCheck env jud
+  return (Map.insert name (Just term :| ty') env)
+
+typeCheck :: Env -> OfType Term Term -> Either String Term
+typeCheck env jud =
+  runTM (check env jud)
+  -- runTM (check env jud) >>= \case
   --   ([], term') -> return term'
   --   (_, _) -> error "Don't know how to solve constraints yet"
 
@@ -64,13 +74,13 @@ throw err = TM (Left err)
 -- emit :: Constr -> TM ()
 -- emit c = TM (tell (BOne c))
 
-check :: Env -> Term -> Term -> TM Term
-check env topTerm typ = case topTerm of
+check :: Env -> OfType Term Term -> TM Term
+check env (topTerm :| typ) = case topTerm of
   TPair a b -> do
     case whnf env typ of
       TSigma name t1 t2 -> do
-        a' <- check env a t1
-        b' <- check (Map.insert name t1 env) b t2
+        a' <- check env (a :| t1)
+        b' <- check (Map.insert name (Nothing :| t1) env) (b :| t2)
         return (TPair a' b')
       t -> throw $ "Pair expression cannot have type " ++ show t
 
@@ -90,35 +100,33 @@ infer env = \case
   TSet i -> do
     return (TSet i :| TSet (TISucc i))
 
+  TSetw i -> do
+    return (TSetw i :| TSetw (succ i))
+
   TVar n -> do
     case Map.lookup n env of
-      Just ty -> return (TVar n :| ty)
+      Just (_ :| ty) -> return (TVar n :| ty)
       Nothing -> throw $ "Variable out of scope: " ++ n
 
   TPi x a b -> do
     inferW env a >>= \case
-      a' :| TSet lvlA -> do
-        inferW (Map.insert x a' env) b >>= \case
-          b' :| TSet lvlB -> do
-            when (x `freeIn` lvlB) $
-              throw $ "Variable " ++ show x ++ " escapes Pi"
-            return (TPi x a' b' :| TSet (TIMax lvlA lvlB))
-          _ :| tb -> throw $ "RHS of a Pi not of type Set i, but: " ++ show tb
-      _ :| ta -> throw $ "LHS type of a Pi not of type Set i, but: " ++ show ta
+      a' :| (argumentKind -> Just kindA) -> do
+        inferW (Map.insert x (Nothing :| a') env) b >>= \case
+          b' :| (argumentKind -> Just kindB) -> do
+            return (TPi x a' b' :| akLower (kindA <> kindB))
+          _ :| tb -> throw $ "RHS of a Pi not of acceptable type, but: " ++ show tb
+      _ :| ta -> throw $ "LHS type of a Pi not of acceptable type, but: " ++ show ta
 
   TLam x t e -> do
     inferW env t >>= \case
-      t' :| TSet{} -> do
-        e' :| te <- inferW (Map.insert x t' env) e
-        when (x `freeIn` te) $
-          throw $ "Variable " ++ show x ++ " escape lambda"
+      t' :| _ -> do
+        e' :| te <- inferW (Map.insert x (Nothing :| t') env) e
         return (TLam x t' e' :| TPi x t' te)
-      _ :| tt -> throw $ "Lambda variable type not of type Set i, but: " ++ show tt
 
   TApp a b -> do
     inferW env a >>= \case
       a' :| TPi name t1 t2 -> do
-        b' <- check env b t1
+        b' <- check env (b :| t1)
         return (TApp a' b' :| subst name b' t2)
       _ :| ta -> throw $ "LHS of application not of Pi type, but: " ++ show ta
 
@@ -152,6 +160,12 @@ infer env = \case
         return (TISucc a' :| TLevel)
       _ :| ta -> throw $ "Argument of isucc not of type Level, but: " ++ show ta
 
+  TAnnot (a :| b) -> do
+    inferW env b >>= \case
+      b' :| _ -> do
+        a' <- check env (a :| b')
+        return (a' :| b')
+
   TOne -> do
     return (TOne :| TSet TIZero)
 
@@ -160,14 +174,12 @@ infer env = \case
 
   TSigma x a b -> do
     inferW env a >>= \case
-      a' :| TSet lvlA -> do
-        inferW (Map.insert x a' env) b >>= \case
-          b' :| TSet lvlB -> do
-            when (x `freeIn` lvlB) $
-              throw $ "Variable " ++ show x ++ " escapes Sigma"
-            return (TSigma x a' b' :| TSet (TIMax lvlA lvlB))
-          _ :| tb -> throw $ "RHS of a Sigma not of type Set i, but: " ++ show tb
-      _ :| ta -> throw $ "LHS type of a Sigma not of type Set i, but: " ++ show ta
+      a' :| (argumentKind -> Just kindA) -> do
+        inferW (Map.insert x (Nothing :| a') env) b >>= \case
+          b' :| (argumentKind -> Just kindB) -> do
+            return (TSigma x a' b' :| akLower (kindA <> kindB))
+          _ :| tb -> throw $ "RHS of a Sigma not of acceptable type, but: " ++ show tb
+      _ :| ta -> throw $ "LHS type of a Sigma not of acceptable type, but: " ++ show ta
 
   TPair{} -> do
     throw "Dependent pair occurring in non-checking position"
@@ -184,33 +196,61 @@ infer env = \case
         return (TProj2 e' :| subst name (TProj1 e') t2)
       _ :| t -> throw $ "Argument of proj2 not of Sigma type, but: " ++ show t
 
-freeIn :: Name -> Term -> Bool
-freeIn target = \case
-  TSet a -> rec a
-  TVar n -> n == target
-  TPi n a b -> rec a && (if n == target then True else rec b)
-  TLam n a b -> rec a && (if n == target then True else rec b)
-  TApp a b -> rec a && rec b
-  TLift a -> rec a
-  TLevel -> False
-  TLevelUniv -> False
-  TIZero -> False
-  TIMax a b -> rec a && rec b
-  TISucc a -> rec a
-  TOne -> False
-  TUnit -> False
-  TSigma n a b -> rec a && (if n == target then True else rec b)
-  TPair a b -> rec a && rec b
-  TProj1 a -> rec a
-  TProj2 a -> rec a
-  where
-    rec = freeIn target
+data ArgKind = AKSet Term | AKSetw Natural | AKLevelUniv
+  deriving (Show)
+
+instance Semigroup ArgKind where
+  AKSet n <> AKSet m = AKSet (TIMax n m)
+  AKSet _ <> ak = AKSetw 0 <> ak
+  ak <> AKSet _ = ak <> AKSetw 0
+
+  AKLevelUniv <> AKLevelUniv = AKLevelUniv
+  AKLevelUniv <> ak = AKSetw 0 <> ak
+  ak <> AKLevelUniv = ak <> AKSetw 0
+
+  AKSetw i <> AKSetw j = AKSetw (max i j)
+
+argumentKind :: Term -> Maybe ArgKind
+argumentKind (TSet t) = Just (AKSet t)
+argumentKind (TSetw i) = Just (AKSetw i)
+argumentKind TLevelUniv = Just AKLevelUniv
+argumentKind _ = Nothing
+
+akLower :: ArgKind -> Term
+akLower (AKSet t) = TSet t
+akLower (AKSetw i) = TSetw i
+akLower AKLevelUniv = TLevelUniv
+
+-- freeIn :: Name -> Term -> Bool
+-- freeIn target = \case
+--   TSet a -> rec a
+--   TSetw _ -> False
+--   TVar n -> n == target
+--   TPi n a b -> rec a || (if n == target then True else rec b)
+--   TLam n a b -> rec a || (if n == target then True else rec b)
+--   TApp a b -> rec a || rec b
+--   TLift a -> rec a
+--   TLevel -> False
+--   TLevelUniv -> False
+--   TIZero -> False
+--   TIMax a b -> rec a || rec b
+--   TISucc a -> rec a
+--   TAnnot (a :| b) -> rec a || rec b
+--   TOne -> False
+--   TUnit -> False
+--   TSigma n a b -> rec a || (if n == target then True else rec b)
+--   TPair a b -> rec a || rec b
+--   TProj1 a -> rec a
+--   TProj2 a -> rec a
+--   where
+--     rec = freeIn target
 
 subst :: Name -> Term -> Term -> Term
 subst target repl = \case
   TVar n | n == target -> repl
 
   TSet a -> TSet (rec a)
+  TSetw i -> TSetw i
   TVar n -> TVar n
   TPi n a b -> TPi n (rec a) (if n == target then b else rec b)
   TLam n a b -> TLam n (rec a) (if n == target then b else rec b)
@@ -221,6 +261,7 @@ subst target repl = \case
   TIZero -> TIZero
   TIMax a b -> TIMax (rec a) (rec b)
   TISucc a -> TISucc (rec a)
+  TAnnot (a :| b) -> TAnnot (rec a :| rec b)
   TOne -> TOne
   TUnit -> TUnit
   TSigma n a b -> TSigma n (rec a) (if n == target then b else rec b)
@@ -232,6 +273,7 @@ subst target repl = \case
 
 unify :: Term -> Term -> TM ()
 unify (TSet a) (TSet b) = unify a b
+unify (TSetw i) (TSetw j) | i == j = return ()
 unify (TVar n) (TVar m) | n == m = return ()
 unify (TPi n a b) (TPi m c d) = unify a c >> unify b (subst m (TVar n) d)
 unify (TLam n a b) (TLam m c d) = unify a c >> unify b (subst m (TVar n) d)
@@ -251,10 +293,9 @@ unify a b = throw $ "Cannot unify:\n- " ++ show a ++ "\n- " ++ show b
 
 whnf :: Env -> Term -> Term
 whnf env = \case
-  TVar n | Just t <- Map.lookup n env -> whnf env t
   TApp (TLam n _ a) b -> whnf env (subst n b a)
   TIMax a b -> merge (whnf env a) (whnf env b)
-    where
+    where -- TODO all of the properties from https://agda.readthedocs.io/en/v2.6.3/language/universe-levels.html#intrinsic-level-properties
       merge TIZero l = l
       merge l TIZero = l
       merge (TISucc l) (TISucc m) = TISucc (merge l m)