4 files changed, 282 insertions, 145 deletions
diff --git a/src/Data/Bag.hs b/src/Data/Bag.hs
index 9cdce89..f1173e4 100644
--- a/src/Data/Bag.hs
+++ b/src/Data/Bag.hs
@@ -3,12 +3,16 @@ module Data.Bag where
 
 
 data Bag a
-  = BTwo (Bag a) (Bag a)
+  = BZero
   | BOne a
-  | BZero
+  | BTwo (Bag a) (Bag a)
   deriving (Functor, Foldable, Traversable)
 
-instance Semigroup (Bag a) where (<>) = BTwo
+instance Semigroup (Bag a) where
+  BZero <> b = b
+  b <> BZero = b
+  b1 <> b2 = BTwo b1 b2
+
 instance Monoid (Bag a) where mempty = BZero
 
 instance Applicative Bag where
@@ -18,3 +22,10 @@ instance Applicative Bag where
   _ <*> BZero = BZero
   BOne f <*> b = f <$> b
   BTwo b1 b2 <*> b = BTwo (b1 <*> b) (b2 <*> b)
+
+bagPartition :: (a -> Maybe b) -> Bag a -> (Bag b, Bag a)
+bagPartition _ BZero = (BZero, BZero)
+bagPartition f (BOne x)
+  | Just y <- f x = (BOne y, BZero)
+  | otherwise = (BZero, BOne x)
+bagPartition f (BTwo b1 b2) = bagPartition f b1 <> bagPartition f b2
diff --git a/src/HSVIS/AST.hs b/src/HSVIS/AST.hs
index 7b05e19..a789c82 100644
--- a/src/HSVIS/AST.hs
+++ b/src/HSVIS/AST.hs
@@ -11,9 +11,10 @@
 {-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ConstrainedClassMethods #-}
 module HSVIS.AST where
 
-import Data.Bifunctor (bimap, second)
+import Data.Bifunctor (bimap, first, second)
 import qualified Data.Kind as DK
 import Data.List.NonEmpty (NonEmpty)
 import Data.Proxy
@@ -34,54 +35,16 @@ class HasExt f where
   type HasXField f :: Bool
   type HasECon f :: Bool
   type ContainsData f :: [ASTType]
-  extOf :: f s -> X f s
+  extOf :: HasXField f ~ 'True => f s -> X f s
   extMap :: Proxy f -> Proxy s1 -> Proxy s2
-         -> XMapperTelescope (TransitiveContainsData '[f]) s1 s2
-              (EMapperTelescope (TransitiveContainsData '[f]) s1 s2
+         -> XMapperTelescope (f ': ContainsData f) s1 s2
+              (EMapperTelescope (f ': ContainsData f) s1 s2
                 (f s1 -> f s2))
 
 type family If c a b where
   If 'True a b = a
   If 'False a b = b
 
-type family ASTShowPrereqs f s :: DK.Constraint where
-  ASTShowPrereqs f s =
-    (If (HasXField f) (Show (X f s)) (() :: DK.Constraint)
-    ,If (HasECon f) (Show (E f s)) (() :: DK.Constraint)
-    ,ASTShowPrereqsC (ContainsData f) s)
-
-type family ASTShowPrereqsC l s :: DK.Constraint where
-  ASTShowPrereqsC '[] s = ()
-  ASTShowPrereqsC (f ': fs) s = (Show (f s), ASTShowPrereqsC fs s)
-
-type family a && b where
-  'False && b = 'False
-  'True && b = b
-
-type family In x l where
-  In x '[] = 'False
-  In x (x ': ys) = 'True
-  In x (_ ': ys) = In x ys
-
-type family Sublist l1 l2 where
-  Sublist '[] l = 'True
-  Sublist (x ': xs) l = In x l && Sublist xs l
-
-type family Union l1 l2 where
-  Union '[] l = l
-  Union (x ': xs) l = If (In x l) (Union xs l) (x ': Union xs l)
-
-type family MapContainsData l where
-  MapContainsData '[] = '[]
-  MapContainsData (f ': fs) = Union (ContainsData f) (MapContainsData fs)
-
-type family TransitiveContainsData fs where
-  TransitiveContainsData fs = TransitiveContainsData' fs (Sublist (Union fs (MapContainsData fs)) fs)
-
-type family TransitiveContainsData' fs isfull where
-  TransitiveContainsData' fs 'True = fs
-  TransitiveContainsData' fs 'False = TransitiveContainsData (Union fs (MapContainsData fs))
-
 type family XMapperTelescope fs s1 s2 a where
   XMapperTelescope '[] s1 s2 a = a
   XMapperTelescope (f ': fs) s1 s2 a =
@@ -93,7 +56,9 @@ type family EMapperTelescope fs s1 s2 a where
   EMapperTelescope '[] s1 s2 a = a
   EMapperTelescope (f ': fs) s1 s2 a =
     If (HasECon f)
-       ((X f s1 -> E f s1 -> f s2) -> EMapperTelescope fs s1 s2 a)
+       (If (HasXField f) (X f s1 -> E f s1 -> f s2)
+                         (E f s1 -> f s2)
+        -> EMapperTelescope fs s1 s2 a)
        (EMapperTelescope fs s1 s2 a)
 
 
@@ -103,19 +68,27 @@ newtype Name = Name String
 data Program s = Program [DataDef s] [FunDef s]
 deriving instance (Show (DataDef s), Show (FunDef s)) => Show (Program s)
 
-data DataDef s = DataDef (X DataDef s) Name [Name] [(Name, [Type s])]
-deriving instance ASTShowPrereqs DataDef s => Show (DataDef s)
+data DataDef s = DataDef (X DataDef s) Name [(X Type s, Name)] [(Name, [Type s])]
+deriving instance (Show (X DataDef s), Show (X Type s), Show (E Type s)) => Show (DataDef s)
 
-data FunDef s = FunDef (X FunDef s) Name (Maybe (Type s)) (NonEmpty (FunEq s))
-deriving instance ASTShowPrereqs FunDef s => Show (FunDef s)
+data FunDef s = FunDef (X FunDef s) Name (TypeSig s) (NonEmpty (FunEq s))
+deriving instance (Show (X FunDef s), Show (TypeSig s), Show (FunEq s)) => Show (FunDef s)
+
+data TypeSig s
+  = TypeSig (Type s)
+  | TypeSigExt !(E TypeSig s)
+deriving instance (Show (E TypeSig s), Show (Type s)) => Show (TypeSig s)
 
 data FunEq s = FunEq (X FunEq s) Name [Pattern s] (RHS s)
-deriving instance ASTShowPrereqs FunEq s => Show (FunEq s)
+deriving instance (Show (X FunEq s), Show (Pattern s), Show (RHS s)) => Show (FunEq s)
 
 data Kind s
   = KType (X Kind s)
   | KFun (X Kind s) (Kind s) (Kind s)
-deriving instance ASTShowPrereqs Kind s => Show (Kind s)
+
+  -- extension point
+  | KExt (X Kind s) !(E Kind s)
+deriving instance (Show (X Kind s), Show (E Kind s)) => Show (Kind s)
 
 data Type s
   = TApp (X Type s) (Type s) [Type s]
@@ -127,7 +100,7 @@ data Type s
 
   -- extension point
   | TExt (X Type s) !(E Type s)
-deriving instance ASTShowPrereqs Type s => Show (Type s)
+deriving instance (Show (X Type s), Show (E Type s)) => Show (Type s)
 
 data Pattern s
   = PWildcard (X Pattern s)
@@ -137,12 +110,12 @@ data Pattern s
   | POp       (X Pattern s) (Pattern s) Operator (Pattern s)
   | PList     (X Pattern s) [Pattern s]
   | PTup      (X Pattern s) [Pattern s]
-deriving instance ASTShowPrereqs Pattern s => Show (Pattern s)
+deriving instance (Show (X Pattern s)) => Show (Pattern s)
 
 data RHS s
   = Guarded (X RHS s) [(Expr s, Expr s)]  -- currently not parsed
   | Plain   (X RHS s) (Expr s)
-deriving instance ASTShowPrereqs RHS s => Show (RHS s)
+deriving instance (Show (X RHS s), Show (Expr s)) => Show (RHS s)
 
 data Expr s
   = ELit  (X Expr s) Literal
@@ -156,7 +129,7 @@ data Expr s
   | ECase (X Expr s) (Expr s) [(Pattern s, RHS s)]
   | ELet  (X Expr s) [FunDef s] (Expr s)
   | EError (X Expr s)
-deriving instance ASTShowPrereqs Expr s => Show (Expr s)
+deriving instance (Show (X Expr s), Show (Pattern s), Show (RHS s), Show (FunDef s)) => Show (Expr s)
 
 data Literal = LInt Integer | LFloat Rational | LChar Char | LString String
   deriving (Show)
@@ -174,39 +147,51 @@ instance HasExt DataDef where
   type ContainsData DataDef = '[Type]
   extOf (DataDef x _ _ _) = x
   extMap _ ps1 ps2 fdd fty gty (DataDef x n ps cs) =
-    DataDef (fdd x) n ps (map (second (map (extMap (Proxy @Type) ps1 ps2 fty gty))) cs)
+    DataDef (fdd x) n (map (first fty) ps) (map (second (map (extMap (Proxy @Type) ps1 ps2 fty gty))) cs)
 
 instance HasExt FunDef where
   type HasXField FunDef = 'True
   type HasECon FunDef = 'False
-  type ContainsData FunDef = '[Type, FunEq]
+  type ContainsData FunDef = '[TypeSig, FunEq, Type, Pattern, RHS, Expr]
   extOf (FunDef x _ _ _) = x
-  extMap _ ps1 ps2 ffd ffe frh fty fpa fex gty (FunDef x n msig es) =
-    FunDef (ffd x) n (extMap (Proxy @Type) ps1 ps2 fty gty <$> msig)
-           (fmap (extMap (Proxy @FunEq) ps1 ps2 frh ffd fpa fex fty ffe gty) es)
+  extMap _ ps1 ps2 ffd ffe fty fpa frh fex gts gty (FunDef x n sig es) =
+    FunDef (ffd x) n
+           (extMap (Proxy @TypeSig) ps1 ps2 fty gts gty sig)
+           (fmap (extMap (Proxy @FunEq) ps1 ps2 ffe ffd fty fpa frh fex gts gty) es)
+
+instance HasExt TypeSig where
+  type HasXField TypeSig = 'False
+  type HasECon TypeSig = 'True
+  type ContainsData TypeSig = '[Type]
+  extOf _ = undefined
+  extMap _ ps1 ps2 fty _   gty (TypeSig ty) =
+    TypeSig (extMap (Proxy @Type) ps1 ps2 fty gty ty)
+  extMap _ _   _   _   gts _   (TypeSigExt e) = gts e
 
 instance HasExt FunEq where
   type HasXField FunEq = 'True
   type HasECon FunEq = 'False
-  type ContainsData FunEq = '[Pattern, RHS]
+  type ContainsData FunEq = '[FunDef, TypeSig, Type, Pattern, RHS, Expr]
   extOf (FunEq x _ _ _) = x
-  extMap _ ps1 ps2 frh ffd fpa fex fty ffe gty (FunEq x n ps rhs) =
+  extMap _ ps1 ps2 ffe ffd fty fpa frh fex gts gty (FunEq x n ps rhs) =
     FunEq (ffe x) n (map (extMap (Proxy @Pattern) ps1 ps2 fpa) ps)
-          (extMap (Proxy @RHS) ps1 ps2 frh ffd fpa fex fty ffe gty rhs)
+          (extMap (Proxy @RHS) ps1 ps2 frh ffd ffe fty fpa fex gts gty rhs)
 
 instance HasExt Kind where
   type HasXField Kind = 'True
-  type HasECon Kind = 'False
-  type ContainsData Kind = '[Kind]
+  type HasECon Kind = 'True
+  type ContainsData Kind = '[]
   extOf (KType x) = x
   extOf (KFun x _ _) = x
-  extMap _ _   _   f (KType x) = KType (f x)
-  extMap p ps1 ps2 f (KFun x a b) = KFun (f x) (extMap p ps1 ps2 f a) (extMap p ps1 ps2 f b)
+  extOf (KExt x _) = x
+  extMap _ _   _   f _ (KType x) = KType (f x)
+  extMap p ps1 ps2 f g (KFun x a b) = KFun (f x) (extMap p ps1 ps2 f g a) (extMap p ps1 ps2 f g b)
+  extMap _ _   _   _ g (KExt x a) = g x a
 
 instance HasExt Type where
   type HasXField Type = 'True
   type HasECon Type = 'True
-  type ContainsData Type = '[Type]
+  type ContainsData Type = '[]
   extOf (TApp x _ _) = x
   extOf (TTup x _) = x
   extOf (TList x _) = x
@@ -226,7 +211,7 @@ instance HasExt Type where
 instance HasExt Pattern where
   type HasXField Pattern = 'True
   type HasECon Pattern = 'False
-  type ContainsData Pattern = '[Pattern]
+  type ContainsData Pattern = '[]
   extOf (PWildcard x) = x
   extOf (PVar x _) = x
   extOf (PAs x _ _) = x
@@ -246,21 +231,22 @@ instance HasExt Pattern where
 instance HasExt RHS where
   type HasXField RHS = 'True
   type HasECon RHS = 'False
-  type ContainsData RHS = '[Expr]
+  type ContainsData RHS = '[FunDef, TypeSig, FunEq, Type, Pattern, Expr]
   extOf (Guarded x _) = x
   extOf (Plain x _) = x
 
-  extMap _ ps1 ps2 frh ffd fpa fex fty ffe gty (Guarded x l) =
-    Guarded (frh x) (map (bimap (extMap (Proxy @Expr) ps1 ps2 frh ffd fpa fex fty ffe gty)
-                                (extMap (Proxy @Expr) ps1 ps2 frh ffd fpa fex fty ffe gty))
-                         l)
-  extMap _ ps1 ps2 frh ffd fpa fex fty ffe gty (Plain x e) =
-    Plain (frh x) (extMap (Proxy @Expr) ps1 ps2 frh ffd fpa fex fty ffe gty e)
+  extMap _ ps1 ps2 frh ffd ffe fty fpa fex gts gty (Guarded x l) =
+    Guarded (frh x)
+            (map (bimap (extMap (Proxy @Expr) ps1 ps2 fex ffd ffe fty fpa frh gts gty)
+                        (extMap (Proxy @Expr) ps1 ps2 fex ffd ffe fty fpa frh gts gty))
+                 l)
+  extMap _ ps1 ps2 frh ffd ffe fty fpa fex gts gty (Plain x e) =
+    Plain (frh x) (extMap (Proxy @Expr) ps1 ps2 fex ffd ffe fty fpa frh gts gty e)
 
 instance HasExt Expr where
   type HasXField Expr = 'True
   type HasECon Expr = 'False
-  type ContainsData Expr = '[Expr, Pattern, RHS, FunDef]
+  type ContainsData Expr = '[FunDef, TypeSig, FunEq, Type, Pattern, RHS]
   extOf (ELit x _) = x
   extOf (EVar x _) = x
   extOf (ECon x _) = x
@@ -273,18 +259,18 @@ instance HasExt Expr where
   extOf (ELet x _ _) = x
   extOf (EError x) = x
 
-  extMap p ps1 ps2 frh ffd fpa fex fty ffe gty = \case
+  extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty = \case
     ELit x l -> ELit (fex x) l
     EVar x n -> EVar (fex x) n
     ECon x n -> ECon (fex x) n
-    EList x es -> EList (fex x) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
-    ETup x es -> ETup (fex x) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
-    EApp x a es -> EApp (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
-    EOp x a n b -> EOp (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) n (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b)
-    EIf x a b c -> EIf (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty c)
-    ECase x a l -> ECase (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a)
+    EList x es -> EList (fex x) (map (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty) es)
+    ETup x es -> ETup (fex x) (map (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty) es)
+    EApp x a es -> EApp (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a) (map (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty) es)
+    EOp x a n b -> EOp (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a) n (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty b)
+    EIf x a b c -> EIf (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty b) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty c)
+    ECase x a l -> ECase (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a)
                          (map (bimap (extMap (Proxy @Pattern) ps1 ps2 fpa)
-                                     (extMap (Proxy @RHS) ps1 ps2 frh ffd fpa fex fty ffe gty))
+                                     (extMap (Proxy @RHS) ps1 ps2 frh ffd ffe fty fpa fex gts gty))
                               l)
-    ELet x as b -> ELet (fex x) (map (extMap (Proxy @FunDef) ps1 ps2 ffd ffe frh fty fpa fex gty) as) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b)
+    ELet x as b -> ELet (fex x) (map (extMap (Proxy @FunDef) ps1 ps2 ffd ffe fty fpa frh fex gts gty) as) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty b)
     EError x -> EError (fex x)
diff --git a/src/HSVIS/Parser.hs b/src/HSVIS/Parser.hs
index eba78d8..3251989 100644
--- a/src/HSVIS/Parser.hs
+++ b/src/HSVIS/Parser.hs
@@ -21,6 +21,7 @@ module HSVIS.Parser (
   parse,
   -- * Staged AST synonyms
   PProgram, PDataDef, PFunDef, PFunEq, PType, PPattern, PRHS, PExpr,
+  E(NoTypeSig),
 ) where
 
 -- import Control.Applicative
@@ -55,12 +56,14 @@ type instance X Pattern StageParsed = Range
 type instance X RHS     StageParsed = Range
 type instance X Expr    StageParsed = Range
 
-data instance E Type StageParsed
-  deriving (Show)
+data instance E TypeSig StageParsed = NoTypeSig deriving (Show)
+data instance E Type    StageParsed deriving (Show)
+data instance E Kind    StageParsed deriving (Show)
 
 type PProgram = Program StageParsed
 type PDataDef = DataDef StageParsed
 type PFunDef  = FunDef  StageParsed
+-- type PTypeSig = TypeSig StageParsed
 type PFunEq   = FunEq   StageParsed
 type PType    = Type    StageParsed
 type PPattern = Pattern StageParsed
@@ -227,7 +230,7 @@ pDataDef0 = do
         raise Error "Expected data declaration after 'data'"
         return []
       Just name -> do
-        params <- famany (inlineWhite >> pIdentifier0 InBlock Lowercase WCBacktrack)
+        params <- famany (inlineWhite >> ranged (pIdentifier0 InBlock Lowercase WCBacktrack))
         cons <- pDatacons "="
         pos2 <- gets psCur
         return [DataDef (Range pos1 pos2) name params cons]
@@ -245,7 +248,7 @@ pDataDef0 = do
 
 data FunEqContext
   = FirstLine
-  | TypeSig Name
+  | WithTypeSig Name
   | Continue Name
   deriving (Show)
 
@@ -254,25 +257,25 @@ pFunDef0 =
   faasum'
     [do (sigrange, (name, typ)) <- ranged pStandaloneTypesig0
         noFail $ do
-          faoptional (pFunEq (TypeSig name)) >>= \case
+          faoptional (pFunEq (WithTypeSig name)) >>= \case
             Just (clause1 : clauses1) -> do
               clauses <- concat <$> famany (pFunEq (Continue name))
               let clauses' = clause1 :| clauses1 ++ clauses
               let rng = sigrange <> foldMapne extOf clauses'
-              return [FunDef rng name (Just typ) clauses']
+              return [FunDef rng name (TypeSig typ) clauses']
             _ -> do
               pos <- gets psCur
               raise Error $ "Expected function equation for " ++ pretty name ++
                             " after type signature"
               let rng = Range pos pos
-              return [FunDef sigrange name (Just typ)
+              return [FunDef sigrange name (TypeSig typ)
                              (FunEq rng name [] (Plain rng (EError rng)) :| [])]
     ,do pFunEq FirstLine >>= \case
           clause1@(FunEq _ name _ _) : clauses1 -> noFail $ do
             clauses <- concat <$> famany (pFunEq (Continue name))
             let clauses' = clause1 :| clauses1 ++ clauses
             let rng = foldMapne extOf clauses'
-            return [FunDef rng name Nothing clauses']
+            return [FunDef rng name (TypeSigExt NoTypeSig) clauses']
           [] -> faempty]
 
 -- | Given the name from the type signature or a previous clause, if any.
@@ -290,7 +293,7 @@ pFunEq fectx = do
     -- avoid code duplication or an early exit monad, we use a boolean here.
     success <- case fectx of
       FirstLine -> return True
-      TypeSig checkName
+      WithTypeSig checkName
         | name == checkName -> return True
         | otherwise -> noFail $ do
             raise Error $ "Name of function clause does not correspond with type signature: " ++
diff --git a/src/HSVIS/Typecheck.hs b/src/HSVIS/Typecheck.hs
index 23fb961..2b05793 100644
--- a/src/HSVIS/Typecheck.hs
+++ b/src/HSVIS/Typecheck.hs
@@ -3,25 +3,27 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TupleSections #-}
 module HSVIS.Typecheck where
 
 import Control.Monad
-import Control.Monad.Reader
-import Control.Monad.State
-import Control.Monad.Writer.CPS
+import Data.Bifunctor (first)
 import Data.Foldable (toList)
+import Data.List (partition)
 import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
 
 import Data.Bag
 import Data.List.NonEmpty.Util
 import HSVIS.AST
 import HSVIS.Parser
 import HSVIS.Diagnostic
+import HSVIS.Pretty
 
 
 data StageTC
 
-type instance X DataDef StageTC = CType
+type instance X DataDef StageTC = ()
 type instance X FunDef  StageTC = CType
 type instance X FunEq   StageTC = CType
 type instance X Kind    StageTC = ()
@@ -30,8 +32,8 @@ 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)
+data instance E Type StageTC = TUniVar Int deriving (Show)
+data instance E Kind StageTC = KUniVar Int deriving (Show)
 
 type CProgram = Program StageTC
 type CDataDef = DataDef StageTC
@@ -54,8 +56,8 @@ 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)
+data instance E Type StageTyped deriving (Show)
+data instance E Kind StageTyped deriving (Show)
 
 type TProgram = Program StageTyped
 type TDataDef = DataDef StageTyped
@@ -68,77 +70,203 @@ type TRHS     = RHS     StageTyped
 type TExpr    = Expr    StageTyped
 
 
-typecheck :: FilePath -> String -> Program Range -> ([Diagnostic], Program TType)
+typecheck :: FilePath -> String -> PProgram -> ([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
-        $ tcProgram prog
+  let (ds1, cs, _, _, progtc) =
+        runTCM (tcProgram prog) (fp, source) 1 (Env mempty mempty)
       (ds2, sub) = solveConstrs cs
   in (toList (ds1 <> ds2), substProg sub progtc)
 
-data Constr =
-  CEq CType CType Range  -- ^ These types must be equal because of the expression here
+data Constr
+  = CEq CType CType Range  -- ^ These types must be equal because of the expression here
+  | CEqK CKind CKind Range  -- ^ These kinds must be equal because of the type here
   deriving (Show)
 
 data Env = Env (Map Name CKind) (Map Name CType)
   deriving (Show)
 
-newtype TCM a = TCM { runTCM ::
-    ReaderT (FilePath, String)
-    (StateT Int
-    (StateT Env
-    (Writer (Bag Diagnostic, Bag Constr))))
-    a }
-  deriving newtype (Functor, Applicative, Monad)
+newtype TCM a = TCM {
+  runTCM :: (FilePath, String)  -- ^ reader context: file and file contents
+         -> Int  -- ^ state: next id to generate
+         -> Env  -- ^ state: type and value environment
+         -> (Bag Diagnostic  -- ^ writer: diagnostics
+            ,Bag Constr      -- ^ writer: constraints
+            ,Int, Env, a)
+  }
+
+instance Functor TCM where
+  fmap f (TCM g) = TCM $ \ctx i env ->
+    let (ds, cs, i', env', x) = g ctx i env
+    in (ds, cs, i', env', f x)
+
+instance Applicative TCM where
+  pure x = TCM $ \_ i env -> (mempty, mempty, i, env, x)
+  (<*>) = ap
+
+instance Monad TCM where
+  TCM f >>= g = TCM $ \ctx i1 env1 ->
+    let (ds2, cs2, i2, env2, x) = f ctx i1 env1
+        (ds3, cs3, i3, env3, y) = runTCM (g x) ctx i2 env2
+    in (ds2 <> ds3, cs2 <> cs3, i3, env3, y)
 
 raise :: Range -> String -> TCM ()
-raise rng@(Range (Pos y _) _) msg = do
-  (fp, source) <- ask
-  TCM $ lift $ lift $ tell (pure (Diagnostic fp rng [] (source !! y) msg), mempty)
+raise rng@(Range (Pos y _) _) msg = TCM $ \(fp, source) i env ->
+  (pure (Diagnostic fp rng [] (lines source !! y) msg), mempty, i, env, ())
+
+emit :: Constr -> TCM ()
+emit c = TCM $ \_ i env -> (mempty, pure c, i, env, ())
+
+collectConstraints :: (Constr -> Maybe b) -> TCM a -> TCM (Bag b, a)
+collectConstraints predicate (TCM f) = TCM $ \ctx i env ->
+  let (ds, cs, i', env', x) = f ctx i env
+      (yes, no) = bagPartition predicate cs
+  in (ds, no, i', env', (yes, x))
+
+getFullEnv :: TCM Env
+getFullEnv = TCM $ \_ i env -> (mempty, mempty, i, env, env)
+
+putFullEnv :: Env -> TCM ()
+putFullEnv env = TCM $ \_ i _ -> (mempty, mempty, i, env, ())
+
+genId :: TCM Int
+genId = TCM $ \_ i env -> (mempty, mempty, i, env, i)
+
+getKind :: Name -> TCM (Maybe CKind)
+getKind name = do
+  Env tenv _ <- getFullEnv
+  return (Map.lookup name tenv)
+
+getType :: Name -> TCM (Maybe CType)
+getType name = do
+  Env _ venv <- getFullEnv
+  return (Map.lookup name venv)
 
 modifyTEnv :: (Map Name CKind -> Map Name CKind) -> TCM ()
-modifyTEnv f = TCM $ lift $ lift $ modify (\(Env a b) -> Env (f a) b)
+modifyTEnv f = do
+  Env tenv venv <- getFullEnv
+  putFullEnv (Env (f tenv) venv)
 
 modifyVEnv :: (Map Name CType -> Map Name CType) -> TCM ()
-modifyVEnv f = TCM $ lift $ lift $ modify (\(Env a b) -> Env a (f b))
+modifyVEnv f = do
+  Env tenv venv <- getFullEnv
+  putFullEnv (Env tenv (f venv))
+
+scopeTEnv :: TCM a -> TCM a
+scopeTEnv m = do
+  Env origtenv _ <- getFullEnv
+  res <- m
+  modifyTEnv (\_ -> origtenv)
+  return res
+
+scopeVEnv :: TCM a -> TCM a
+scopeVEnv m = do
+  Env _ origvenv <- getFullEnv
+  res <- m
+  modifyVEnv (\_ -> origvenv)
+  return res
+
+genKUniVar :: TCM CKind
+genKUniVar = KExt () . KUniVar <$> genId
 
 genUniVar :: CKind -> TCM CType
-genUniVar k = TCM $ lift $ state (\i -> (TExt k (TUniVar i), i + 1))
+genUniVar k = TExt k . TUniVar <$> genId
+
+getKind' :: Range -> Name -> TCM CKind
+getKind' rng name = getKind name >>= \case
+  Nothing -> do
+    raise rng $ "Type not in scope: " ++ pretty name
+    genKUniVar
+  Just k -> return k
+
+getType' :: Range -> Name -> TCM CType
+getType' rng name = getType name >>= \case
+  Nothing -> do
+    raise rng $ "Variable not in scope: " ++ pretty name
+    genUniVar (KType ())
+  Just k -> return k
 
 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
+  (kconstrs, ddefs') <- collectConstraints isCEqK $ do
+    mapM_ prepareDataDef ddefs
+    mapM tcDataDef ddefs
+
+  solveKindVars kconstrs
+
+  fdefs' <- mapM tcFunDef fdefs
+
+  return (Program ddefs' fdefs')
+
+prepareDataDef :: PDataDef -> TCM ()
+prepareDataDef (DataDef _ name params _) = do
+  parkinds <- mapM (\_ -> genKUniVar) params
+  let k = foldr (KFun ()) (KType ()) parkinds
+  modifyTEnv (Map.insert name k)
+
+-- Assumes that the kind of the name itself has already been registered with
+-- the correct arity (this is doen by prepareDataDef).
+tcDataDef :: PDataDef -> TCM CDataDef
+tcDataDef (DataDef rng name params cons) = do
+  kd <- getKind' rng name
+  let (pkinds, kret) = splitKind kd
+
+  -- sanity checking; would be nicer to store these in prepareDataDef already
+  when (length pkinds /= length params) $ error "tcDataDef: Invalid param kind list length"
+  case kret of Right () -> return ()
+               _ -> error "tcDataDef: Invalid ret kind"
+
+  cons' <- scopeTEnv $ do
+    modifyTEnv (Map.fromList (zip (map snd params) pkinds) <>)
+    mapM (\(cname, ty) -> (cname,) <$> mapM kcType ty) cons
+  return (DataDef () name (zip pkinds (map snd params)) cons')
 
-  Program <$> traverse checkDataDef ddefs <*> traverse tcFunDef fdefs
+kcType :: PType -> TCM CType
+kcType = \case
+  TApp rng t ts -> do
+    t' <- kcType t
+    ts' <- mapM kcType ts
+    retk <- genKUniVar
+    let expected = foldr (KFun ()) retk (map extOf ts')
+    emit $ CEqK (extOf t') expected rng
+    return (TApp retk t' ts')
 
-checkDataDef :: PDataDef -> TCM CDataDef
-checkDataDef (DataDef rng name params cons) = do
-  _
+  TTup _ ts -> do
+    ts' <- mapM kcType ts
+    forM_ (zip (map extOf ts) ts') $ \(trng, ct) ->
+      emit $ CEqK (extOf ct) (KType ()) trng
+    return (TTup (KType ()) ts')
+
+  TList _ t -> do
+    t' <- kcType t
+    emit $ CEqK (extOf t') (KType ()) (extOf t)
+    return (TList (KType ()) t')
+
+  TFun _ t1 t2 -> do
+    t1' <- kcType t1
+    t2' <- kcType t2
+    emit $ CEqK (extOf t1') (KType ()) (extOf t1)
+    emit $ CEqK (extOf t2') (KType ()) (extOf t2)
+    return (TFun (KType ()) t1' t2')
+
+  TCon rng n -> TCon <$> getKind' rng n <*> pure n
+
+  TVar rng n -> TVar <$> getKind' rng n <*> pure n
 
 tcFunDef :: PFunDef -> TCM CFunDef
-tcFunDef (FunDef rng name msig eqs) = do
+tcFunDef (FunDef _ name msig eqs) = do
   when (not $ allEq (fmap (length . funeqPats) eqs)) $
     raise (sconcatne (fmap extOf eqs)) "Function equations have differing numbers of arguments"
 
   typ <- case msig of
-    Just sig -> checkType sig
-    Nothing -> genUniVar (KType ())
+    TypeSig sig -> kcType sig
+    TypeSigExt NoTypeSig -> genUniVar (KType ())
 
-  _
+  eqs' <- mapM (tcFunEq typ) eqs
 
-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 -> _
+  return (FunDef typ name (TypeSig typ) eqs')
+
+tcFunEq :: CType -> PFunEq -> TCM CFunEq
+tcFunEq = _
 
 allEq :: (Eq a, Foldable t) => t a -> Bool
 allEq l = case toList l of
@@ -147,3 +275,12 @@ allEq l = case toList l of
 
 funeqPats :: FunEq t -> [Pattern t]
 funeqPats (FunEq _ _ pats _) = pats
+
+splitKind :: Kind s -> ([Kind s], Either (E Kind s) (X Kind s))
+splitKind (KType x) = ([], Right x)
+splitKind (KFun _ k1 k2) = first (k1:) (splitKind k2)
+splitKind (KExt _ e) = ([], Left e)
+
+isCEqK :: Constr -> Maybe (CKind, CKind, Range)
+isCEqK (CEqK k1 k2 rng) = Just (k1, k2, rng)
+isCEqK _ = Nothing