From ae603f2423e967c55dfd31b0dec26d19584aa322 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Sat, 23 Dec 2023 23:12:03 +0100
Subject: Can typecheck universe-polymorphic id

---
 src/TypeCheck.hs | 153 +++++++++++++++++++++++++++++++++++--------------------
 1 file changed, 97 insertions(+), 56 deletions(-)

(limited to 'src/TypeCheck.hs')

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)
-- 
cgit v1.2.3-70-g09d2