1 files changed, 103 insertions, 17 deletions

diff --git a/src/HSVIS/Typecheck.hs b/src/HSVIS/Typecheck.hs
index b1ffbb9..23fb961 100644
--- a/src/HSVIS/Typecheck.hs
+++ b/src/HSVIS/Typecheck.hs
@@ -1,5 +1,8 @@
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE EmptyDataDeriving #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE TypeFamilies #-}
 module HSVIS.Typecheck where
 
 import Control.Monad
@@ -7,17 +10,69 @@ import Control.Monad.Reader
 import Control.Monad.State
 import Control.Monad.Writer.CPS
 import Data.Foldable (toList)
-import qualified Data.List.NonEmpty as NE
+import Data.Map.Strict (Map)
 
 import Data.Bag
+import Data.List.NonEmpty.Util
 import HSVIS.AST
+import HSVIS.Parser
 import HSVIS.Diagnostic
 
 
-typecheck :: FilePath -> String -> Program Range -> ([Diagnostic], Program Type)
+data StageTC
+
+type instance X DataDef StageTC = CType
+type instance X FunDef  StageTC = CType
+type instance X FunEq   StageTC = CType
+type instance X Kind    StageTC = ()
+type instance X Type    StageTC = CKind
+type instance X Pattern StageTC = CType
+type instance X RHS     StageTC = CType
+type instance X Expr    StageTC = CType
+
+data instance E Type StageTC = TUniVar Int
+  deriving (Show)
+
+type CProgram = Program StageTC
+type CDataDef = DataDef StageTC
+type CFunDef  = FunDef  StageTC
+type CFunEq   = FunEq   StageTC
+type CKind    = Kind    StageTC
+type CType    = Type    StageTC
+type CPattern = Pattern StageTC
+type CRHS     = RHS     StageTC
+type CExpr    = Expr    StageTC
+
+data StageTyped
+
+type instance X DataDef StageTyped = TType
+type instance X FunDef  StageTyped = TType
+type instance X FunEq   StageTyped = TType
+type instance X Kind    StageTyped = ()
+type instance X Type    StageTyped = TKind
+type instance X Pattern StageTyped = TType
+type instance X RHS     StageTyped = TType
+type instance X Expr    StageTyped = TType
+
+data instance E t StageTyped
+  deriving (Show)
+
+type TProgram = Program StageTyped
+type TDataDef = DataDef StageTyped
+type TFunDef  = FunDef  StageTyped
+type TFunEq   = FunEq   StageTyped
+type TKind    = Kind    StageTyped
+type TType    = Type    StageTyped
+type TPattern = Pattern StageTyped
+type TRHS     = RHS     StageTyped
+type TExpr    = Expr    StageTyped
+
+
+typecheck :: FilePath -> String -> Program Range -> ([Diagnostic], Program TType)
 typecheck fp source prog =
   let (progtc, (ds1, cs)) =
         runWriter
+        . flip evalStateT (Env mempty mempty)
         . flip evalStateT 1
         . flip runReaderT (fp, source)
         . runTCM
@@ -26,9 +81,18 @@ typecheck fp source prog =
   in (toList (ds1 <> ds2), substProg sub progtc)
 
 data Constr =
-  CEq Type Type Range  -- ^ These types must be equal because of the expression here
+  CEq CType CType Range  -- ^ These types must be equal because of the expression here
+  deriving (Show)
+
+data Env = Env (Map Name CKind) (Map Name CType)
+  deriving (Show)
 
-newtype TCM a = TCM { runTCM :: ReaderT (FilePath, String) (StateT Int (Writer (Bag Diagnostic, Bag Constr))) a }
+newtype TCM a = TCM { runTCM ::
+    ReaderT (FilePath, String)
+    (StateT Int
+    (StateT Env
+    (Writer (Bag Diagnostic, Bag Constr))))
+    a }
   deriving newtype (Functor, Applicative, Monad)
 
 raise :: Range -> String -> TCM ()
@@ -36,28 +100,50 @@ raise rng@(Range (Pos y _) _) msg = do
   (fp, source) <- ask
   TCM $ lift $ lift $ tell (pure (Diagnostic fp rng [] (source !! y) msg), mempty)
 
-tcProgram :: Program Range -> TCM (Program TCType)
-tcProgram (Program ddefs fdefs) = Program ddefs <$> traverse tcFunDef fdefs
+modifyTEnv :: (Map Name CKind -> Map Name CKind) -> TCM ()
+modifyTEnv f = TCM $ lift $ lift $ modify (\(Env a b) -> Env (f a) b)
 
-tcFunDef :: FunDef Range -> TCM (FunDef TCType)
-tcFunDef (FunDef name msig eqs) = do
+modifyVEnv :: (Map Name CType -> Map Name CType) -> TCM ()
+modifyVEnv f = TCM $ lift $ lift $ modify (\(Env a b) -> Env a (f b))
+
+genUniVar :: CKind -> TCM CType
+genUniVar k = TCM $ lift $ state (\i -> (TExt k (TUniVar i), i + 1))
+
+tcProgram :: PProgram -> TCM CProgram
+tcProgram (Program ddefs fdefs) = do
+  -- TODO: add preliminary kinds for the data definitions to the environment,
+  -- then solve for the kind variables in checkDataDef
+
+  Program <$> traverse checkDataDef ddefs <*> traverse tcFunDef fdefs
+
+checkDataDef :: PDataDef -> TCM CDataDef
+checkDataDef (DataDef rng name params cons) = do
+  _
+
+tcFunDef :: PFunDef -> TCM CFunDef
+tcFunDef (FunDef rng name msig eqs) = do
   when (not $ allEq (fmap (length . funeqPats) eqs)) $
-    raise (sconcatne (fmap funeqRange eqs)) "Function equations have differing numbers of arguments"
+    raise (sconcatne (fmap extOf eqs)) "Function equations have differing numbers of arguments"
+
+  typ <- case msig of
+    Just sig -> checkType sig
+    Nothing -> genUniVar (KType ())
 
   _
 
+checkType :: PType -> TCM CType
+checkType = \case
+  TApp r t ts -> _
+  TTup r ts -> _
+  TList r t -> _
+  TFun r s t -> _
+  TCon r n -> _
+  TVar r n -> _
+
 allEq :: (Eq a, Foldable t) => t a -> Bool
 allEq l = case toList l of
             [] -> True
             x : xs -> all (== x) xs
 
-funeqRange :: FunEq t -> t
-funeqRange (FunEq rng _ _ _) = rng
-
 funeqPats :: FunEq t -> [Pattern t]
 funeqPats (FunEq _ _ pats _) = pats
-
-sconcatne :: Semigroup a => NE.NonEmpty a -> a
-sconcatne = \(x NE.:| xs) -> go x xs
-  where go a [] = a
-        go a (x : xs) = go (a <> x) xs