Trees that explode

5 files changed, 480 insertions, 131 deletions

diff --git a/hs-visinter.cabal b/hs-visinter.cabal
index 276daa2..0b21e36 100644
--- a/hs-visinter.cabal
+++ b/hs-visinter.cabal
@@ -11,6 +11,7 @@ library
   exposed-modules:
     Control.FAlternative
     Data.Bag
+    Data.List.NonEmpty.Util
     HSVIS.AST
     HSVIS.Diagnostic
     HSVIS.Parser
diff --git a/src/Data/List/NonEmpty/Util.hs b/src/Data/List/NonEmpty/Util.hs
new file mode 100644
index 0000000..3a0e7f6
--- /dev/null
+++ b/src/Data/List/NonEmpty/Util.hs
@@ -0,0 +1,12 @@
+module Data.List.NonEmpty.Util where
+
+import Data.List.NonEmpty (NonEmpty(..))
+
+
+sconcatne :: Semigroup s => NonEmpty s -> s
+sconcatne = \(x :| xs) -> go x xs
+  where go a [] = a
+        go a (x : xs) = go (a <> x) xs
+
+foldMapne :: Semigroup s => (a -> s) -> NonEmpty a -> s
+foldMapne f = sconcatne . fmap f
diff --git a/src/HSVIS/AST.hs b/src/HSVIS/AST.hs
index f606d22..7b05e19 100644
--- a/src/HSVIS/AST.hs
+++ b/src/HSVIS/AST.hs
@@ -1,75 +1,162 @@
+{-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE LambdaCase #-}
 module HSVIS.AST where
 
+import Data.Bifunctor (bimap, second)
+import qualified Data.Kind as DK
 import Data.List.NonEmpty (NonEmpty)
+import Data.Proxy
 
 import HSVIS.Pretty
 
 
+-- | An AST type has a single stage argument.
+type ASTType = DK.Type -> DK.Type
+
+-- | Trees that Grow extension type family
+type family X (f :: ASTType) s
+
+-- | Additional constructors
+data family E (f :: ASTType) s
+
+class HasExt f where
+  type HasXField f :: Bool
+  type HasECon f :: Bool
+  type ContainsData f :: [ASTType]
+  extOf :: f s -> X f s
+  extMap :: Proxy f -> Proxy s1 -> Proxy s2
+         -> XMapperTelescope (TransitiveContainsData '[f]) s1 s2
+              (EMapperTelescope (TransitiveContainsData '[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 =
+    If (HasXField f)
+       ((X f s1 -> X f s2) -> XMapperTelescope fs s1 s2 a)
+       (XMapperTelescope fs s1 s2 a)
+
+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)
+       (EMapperTelescope fs s1 s2 a)
+
+
 newtype Name = Name String
   deriving (Show, Eq)
 
-data Program t = Program [DataDef] [FunDef t]
-  deriving (Show)
+data Program s = Program [DataDef s] [FunDef s]
+deriving instance (Show (DataDef s), Show (FunDef s)) => Show (Program s)
 
-data DataDef = DataDef Name [Name] [(Name, [Type])]
-  deriving (Show)
+data DataDef s = DataDef (X DataDef s) Name [Name] [(Name, [Type s])]
+deriving instance ASTShowPrereqs DataDef s => Show (DataDef s)
 
-data FunDef t = FunDef Name (Maybe Type) (NonEmpty (FunEq t))
-  deriving (Show)
+data FunDef s = FunDef (X FunDef s) Name (Maybe (Type s)) (NonEmpty (FunEq s))
+deriving instance ASTShowPrereqs FunDef s => Show (FunDef s)
 
-data FunEq t = FunEq t Name [Pattern t] (RHS t)
-  deriving (Show)
+data FunEq s = FunEq (X FunEq s) Name [Pattern s] (RHS s)
+deriving instance ASTShowPrereqs FunEq s => Show (FunEq s)
 
-data TypeStage = TSTC | TSNormal
-  deriving (Show)
+data Kind s
+  = KType (X Kind s)
+  | KFun (X Kind s) (Kind s) (Kind s)
+deriving instance ASTShowPrereqs Kind s => Show (Kind s)
 
-data GType (stage :: TypeStage) where
-  TApp :: GType s -> [GType s] -> GType s
-  TTup :: [GType s] -> GType s
-  TList :: GType s -> GType s
-  TFun :: GType s -> GType s -> GType s
-  TCon :: Name -> GType s
-  TVar :: Name -> GType s
-
-  -- Type constructor used only while type checking
-  TUniVar :: Int -> GType 'TSTC
-deriving instance Show (GType stage)
-
-type TCType = GType 'TSTC
-type Type = GType 'TSNormal
-
-data Pattern t
-  = PWildcard t
-  | PVar t Name
-  | PAs t Name (Pattern t)
-  | PCon t Name [Pattern t]
-  | POp t (Pattern t) Operator (Pattern t)
-  | PList t [Pattern t]
-  | PTup t [Pattern t]
-  deriving (Show)
+data Type s
+  = TApp (X Type s) (Type s) [Type s]
+  | TTup (X Type s) [Type s]
+  | TList (X Type s) (Type s)
+  | TFun (X Type s) (Type s) (Type s)
+  | TCon (X Type s) Name
+  | TVar (X Type s) Name
 
-data RHS t
-  = Guarded [(Expr t, Expr t)]  -- currently not parsed
-  | Plain (Expr t)
-  deriving (Show)
+  -- extension point
+  | TExt (X Type s) !(E Type s)
+deriving instance ASTShowPrereqs Type s => Show (Type s)
 
-data Expr t
-  = ELit t Literal
-  | EVar t Name
-  | ECon t Name
-  | EList t [Expr t]
-  | ETup t [Expr t]
-  | EApp t (Expr t) [Expr t]
-  | EOp t (Expr t) Operator (Expr t)
-  | EIf t (Expr t) (Expr t) (Expr t)
-  | ECase t (Expr t) [(Pattern t, RHS t)]
-  | ELet t [FunDef t] (Expr t)
-  | EParseError
-  deriving (Show)
+data Pattern s
+  = PWildcard (X Pattern s)
+  | PVar      (X Pattern s) Name
+  | PAs       (X Pattern s) Name (Pattern s)
+  | PCon      (X Pattern s) Name [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)
+
+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)
+
+data Expr s
+  = ELit  (X Expr s) Literal
+  | EVar  (X Expr s) Name
+  | ECon  (X Expr s) Name
+  | EList (X Expr s) [Expr s]
+  | ETup  (X Expr s) [Expr s]
+  | EApp  (X Expr s) (Expr s) [Expr s]
+  | EOp   (X Expr s) (Expr s) Operator (Expr s)
+  | EIf   (X Expr s) (Expr s) (Expr s) (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)
 
 data Literal = LInt Integer | LFloat Rational | LChar Char | LString String
   deriving (Show)
@@ -80,3 +167,124 @@ data Operator = OAdd | OSub | OMul | ODiv | OMod | OEqu | OPow
 
 instance Pretty Name where
   prettysPrec _ (Name n) = showString ("\"" ++ n ++ "\"")
+
+instance HasExt DataDef where
+  type HasXField DataDef = 'True
+  type HasECon DataDef = 'False
+  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)
+
+instance HasExt FunDef where
+  type HasXField FunDef = 'True
+  type HasECon FunDef = 'False
+  type ContainsData FunDef = '[Type, FunEq]
+  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)
+
+instance HasExt FunEq where
+  type HasXField FunEq = 'True
+  type HasECon FunEq = 'False
+  type ContainsData FunEq = '[Pattern, RHS]
+  extOf (FunEq x _ _ _) = x
+  extMap _ ps1 ps2 frh ffd fpa fex fty ffe 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)
+
+instance HasExt Kind where
+  type HasXField Kind = 'True
+  type HasECon Kind = 'False
+  type ContainsData Kind = '[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)
+
+instance HasExt Type where
+  type HasXField Type = 'True
+  type HasECon Type = 'True
+  type ContainsData Type = '[Type]
+  extOf (TApp x _ _) = x
+  extOf (TTup x _) = x
+  extOf (TList x _) = x
+  extOf (TFun x _ _) = x
+  extOf (TCon x _) = x
+  extOf (TVar x _) = x
+  extOf (TExt x _) = x
+
+  extMap p ps1 ps2 f g (TApp x a b) = TApp (f x) (extMap p ps1 ps2 f g a) (map (extMap p ps1 ps2 f g) b)
+  extMap p ps1 ps2 f g (TTup x as) = TTup (f x) (map (extMap p ps1 ps2 f g) as)
+  extMap p ps1 ps2 f g (TList x a) = TList (f x) (extMap p ps1 ps2 f g a)
+  extMap p ps1 ps2 f g (TFun x a b) = TFun (f x) (extMap p ps1 ps2 f g a) (extMap p ps1 ps2 f g b)
+  extMap _ _   _   f _ (TCon x n) = TCon (f x) n
+  extMap _ _   _   f _ (TVar x n) = TVar (f x) n
+  extMap _ _   _   _ g (TExt x e) = g x e
+
+instance HasExt Pattern where
+  type HasXField Pattern = 'True
+  type HasECon Pattern = 'False
+  type ContainsData Pattern = '[Pattern]
+  extOf (PWildcard x) = x
+  extOf (PVar x _) = x
+  extOf (PAs x _ _) = x
+  extOf (PCon x _ _) = x
+  extOf (POp x _ _ _) = x
+  extOf (PList x _) = x
+  extOf (PTup x _) = x
+
+  extMap _ _   _   f (PWildcard x) = PWildcard (f x)
+  extMap _ _   _   f (PVar x n) = PVar (f x) n
+  extMap p ps1 ps2 f (PAs x n a) = PAs (f x) n (extMap p ps1 ps2 f a)
+  extMap p ps1 ps2 f (PCon x n as) = PCon (f x) n (map (extMap p ps1 ps2 f) as)
+  extMap p ps1 ps2 f (POp x a n b) = POp (f x) (extMap p ps1 ps2 f a) n (extMap p ps1 ps2 f b)
+  extMap p ps1 ps2 f (PList x as) = PList (f x) (map (extMap p ps1 ps2 f) as)
+  extMap p ps1 ps2 f (PTup x as) = PTup (f x) (map (extMap p ps1 ps2 f) as)
+
+instance HasExt RHS where
+  type HasXField RHS = 'True
+  type HasECon RHS = 'False
+  type ContainsData RHS = '[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)
+
+instance HasExt Expr where
+  type HasXField Expr = 'True
+  type HasECon Expr = 'False
+  type ContainsData Expr = '[Expr, Pattern, RHS, FunDef]
+  extOf (ELit x _) = x
+  extOf (EVar x _) = x
+  extOf (ECon x _) = x
+  extOf (EList x _) = x
+  extOf (ETup x _) = x
+  extOf (EApp x _ _) = x
+  extOf (EOp x _ _ _) = x
+  extOf (EIf x _ _ _) = x
+  extOf (ECase x _ _) = x
+  extOf (ELet x _ _) = x
+  extOf (EError x) = x
+
+  extMap p ps1 ps2 frh ffd fpa fex fty ffe 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)
+                         (map (bimap (extMap (Proxy @Pattern) ps1 ps2 fpa)
+                                     (extMap (Proxy @RHS) ps1 ps2 frh ffd fpa fex fty ffe 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)
+    EError x -> EError (fex x)
diff --git a/src/HSVIS/Parser.hs b/src/HSVIS/Parser.hs
index 04e6a63..eba78d8 100644
--- a/src/HSVIS/Parser.hs
+++ b/src/HSVIS/Parser.hs
@@ -12,11 +12,15 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE EmptyDataDeriving #-}
 -- I don't want a warning for 'head' and 'tail' in this file. But I also don't
 -- want GHCs before 9.8 to complain that they don't know the x-partial warning.
 {-# OPTIONS_GHC -Wno-unrecognised-warning-flags -Wno-x-partial #-}
 module HSVIS.Parser (
+  StageParsed,
   parse,
+  -- * Staged AST synonyms
+  PProgram, PDataDef, PFunDef, PFunEq, PType, PPattern, PRHS, PExpr,
 ) where
 
 -- import Control.Applicative
@@ -34,11 +38,36 @@ import Data.These
 
 import Control.FAlternative
 import Data.Bag
+import Data.List.NonEmpty.Util
 import HSVIS.AST
 import HSVIS.Diagnostic
 import HSVIS.Pretty
 
 
+-- | The stage for the parser
+data StageParsed
+
+type instance X DataDef StageParsed = Range
+type instance X FunDef  StageParsed = Range
+type instance X FunEq   StageParsed = Range
+type instance X Type    StageParsed = Range
+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)
+
+type PProgram = Program StageParsed
+type PDataDef = DataDef StageParsed
+type PFunDef  = FunDef  StageParsed
+type PFunEq   = FunEq   StageParsed
+type PType    = Type    StageParsed
+type PPattern = Pattern StageParsed
+type PRHS     = RHS     StageParsed
+type PExpr    = Expr    StageParsed
+
+
 -- Positions are zero-based in both dimensions.
 -- See 'isInsideBlock' and 'isAtBlockLeft' for the two relevant "inside the
 -- block" conditions.
@@ -145,7 +174,7 @@ instance KnownFallible fail => MonadChronicle (Bag Diagnostic) (Parser fail) whe
     That res -> Parser (\_ ps kok _ _ -> kok ps mempty res)
     These errs res -> Parser (\_ ps kok _ _ -> kok ps errs res)
 
-parse :: FilePath -> String -> ([Diagnostic], Maybe (Program Range))
+parse :: FilePath -> String -> ([Diagnostic], Maybe PProgram)
 parse fp source =
   runParser pProgram (Context fp (lines source) []) (PS (Pos 0 0) (Pos 0 0) source)
     (\_ errs res -> case errs of
@@ -154,7 +183,7 @@ parse fp source =
     (\errs -> (toList errs, Nothing))
     ()  -- the program parser cannot fail! :D
 
-pProgram :: IParser (Program Range)
+pProgram :: IParser PProgram
 pProgram = do
   defs <- pTopDefs
   let (datadefs, fundefs) = partitionEithers defs
@@ -162,7 +191,7 @@ pProgram = do
   assertEOF Error
   return (Program datadefs fundefs)
 
-pTopDefs :: IParser [Either DataDef (FunDef Range)]
+pTopDefs :: IParser [Either PDataDef PFunDef]
 pTopDefs = do
   faoptional pTopDef >>= \case
     Nothing -> do
@@ -177,7 +206,7 @@ pTopDefs = do
       defs2 <- pTopDefs
       return (defs ++ defs2)
 
-pTopDef :: FParser [Either DataDef (FunDef Range)]
+pTopDef :: FParser [Either PDataDef PFunDef]
 pTopDef = do
   noFail skipWhiteComment
   noFail isAtBlockLeft >>= \case
@@ -187,8 +216,9 @@ pTopDef = do
       noFail $ readWhileInline (const True)
       pTopDef
 
-pDataDef0 :: FParser [DataDef]
+pDataDef0 :: FParser [PDataDef]
 pDataDef0 = do
+  pos1 <- gets psCur
   pKeyword "data"
   noFail $ do
     inlineWhite
@@ -199,9 +229,10 @@ pDataDef0 = do
       Just name -> do
         params <- famany (inlineWhite >> pIdentifier0 InBlock Lowercase WCBacktrack)
         cons <- pDatacons "="
-        return [DataDef name params cons]
+        pos2 <- gets psCur
+        return [DataDef (Range pos1 pos2) name params cons]
   where
-    pDatacons :: String -> IParser [(Name, [Type])]
+    pDatacons :: String -> IParser [(Name, [PType])]
     pDatacons leader = do
       inlineWhite
       facatch (return []) $ do
@@ -218,31 +249,34 @@ data FunEqContext
   | Continue Name
   deriving (Show)
 
-pFunDef0 :: FParser [FunDef Range]
+pFunDef0 :: FParser [PFunDef]
 pFunDef0 =
   faasum'
-    [do (name, typ) <- pStandaloneTypesig0
+    [do (sigrange, (name, typ)) <- ranged pStandaloneTypesig0
         noFail $ do
           faoptional (pFunEq (TypeSig name)) >>= \case
-            Nothing -> do
-              raise Error $ "Expected function equation for " ++ pretty name ++
-                            " after type signature"
-              return []
-            Just [] -> do
-              pos <- gets psCur
-              return [FunDef name (Just typ)
-                             (FunEq (Range pos pos) name [] (Plain EParseError) :| [])]
             Just (clause1 : clauses1) -> do
               clauses <- concat <$> famany (pFunEq (Continue name))
-              return [FunDef name (Just typ) (clause1 :| clauses1 ++ clauses)]
+              let clauses' = clause1 :| clauses1 ++ clauses
+              let rng = sigrange <> foldMapne extOf clauses'
+              return [FunDef rng name (Just 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)
+                             (FunEq rng name [] (Plain rng (EError rng)) :| [])]
     ,do pFunEq FirstLine >>= \case
           clause1@(FunEq _ name _ _) : clauses1 -> noFail $ do
             clauses <- concat <$> famany (pFunEq (Continue name))
-            return [FunDef name Nothing (clause1 :| clauses1 ++ clauses)]
+            let clauses' = clause1 :| clauses1 ++ clauses
+            let rng = foldMapne extOf clauses'
+            return [FunDef rng name Nothing clauses']
           [] -> faempty]
 
 -- | Given the name from the type signature or a previous clause, if any.
-pFunEq :: FunEqContext -> FParser [FunEq Range]
+pFunEq :: FunEqContext -> FParser [PFunEq]
 pFunEq fectx = do
   noFail skipWhiteComment
   faguardM isAtBlockLeft
@@ -275,24 +309,24 @@ pFunEq fectx = do
       else return []
 
 -- | Pass "=" for function definitions and "->" for case clauses.
-pRHS :: String -> IParser (RHS Range)
+pRHS :: String -> IParser PRHS
 pRHS sepsym = do
   -- TODO: parse guards
   inlineWhite
   pKeySym0 sepsym <|>> raise Error ("Expected " ++ show sepsym)
   expr <- pExpr <|>> expectedExpression
-  return (Plain expr)
+  return (Plain (extOf expr) expr)
 
-pPattern :: Int -> FParser (Pattern Range)
+pPattern :: Int -> FParser PPattern
 pPattern d = inlineWhite >> pPattern0 d
 
-pPattern0 :: Int -> FParser (Pattern Range)
+pPattern0 :: Int -> FParser PPattern
 pPattern0 d = do
   pos1 <- gets psCur
   p0 <- pPatExprAtom0 (max 10 d)
   climbRight pPattern (pInfixOp Uppercase) POp d pos1 p0 Nothing
 
-pExpr :: FParser (Expr Range)
+pExpr :: FParser PExpr
 pExpr = do
   inlineWhite
   -- basics: lit, list, var, con, tup
@@ -305,7 +339,7 @@ pExpr = do
             ,pEIf0
             ,pExprOpExpr0 0]
 
-pPatExprAtom0 :: Int -> FParser (Pattern Range)
+pPatExprAtom0 :: Int -> FParser PPattern
 pPatExprAtom0 d =
   faasum' [pPatWildcard0
           ,pPatVarOrAs0
@@ -362,7 +396,7 @@ pPatExprAtom0 d =
                   pos2 <- gets psCur
                   return (PTup (Range pos1 pos2) (p : ps))]]
 
-pELet0 :: FParser (Expr Range)
+pELet0 :: FParser PExpr
 pELet0 = do
   pos1 <- gets psCur
   pKeyword "let"
@@ -388,7 +422,7 @@ pELet0 = do
     inlineWhite
     facatch (do pos2 <- gets psCur
                 raise Error "Expected 'in' after 'let'"
-                return (ELet (Range pos1 pos2) defs EParseError)) $ do
+                return (ELet (Range pos1 pos2) defs (EError (Range pos2 pos2)))) $ do
       pKeyword "in"
       noFail $ do
         inlineWhite
@@ -396,7 +430,7 @@ pELet0 = do
         pos2 <- gets psCur
         return (ELet (Range pos1 pos2) defs body)
 
-pECase0 :: FParser (Expr Range)
+pECase0 :: FParser PExpr
 pECase0 = do
   pos1 <- gets psCur
   pKeyword "case"
@@ -423,7 +457,7 @@ pECase0 = do
           pos2 <- gets psCur
           return (ECase (Range pos1 pos2) e clauses)
 
-pEIf0 :: FParser (Expr Range)
+pEIf0 :: FParser PExpr
 pEIf0 = do
   pos1 <- gets psCur
   pKeyword "if"
@@ -432,24 +466,26 @@ pEIf0 = do
     inlineWhite
     facatch (do pos2 <- gets psCur
                 raise Error "Expected 'then' after 'if'"
-                return (EIf (Range pos1 pos2) e1 EParseError EParseError)) $ do
+                let rng = Range pos2 pos2
+                return (EIf (Range pos1 pos2) e1 (EError rng) (EError rng))) $ do
       pKeyword "then"
       noFail $ do
         e2 <- pExpr <|>> expectedExpression
         inlineWhite
         facatch (do pos2 <- gets psCur
                     raise Error "Expected else after 'then'"
-                    return (EIf (Range pos1 pos2) e1 e2 EParseError)) $ do
+                    let rng = Range pos2 pos2
+                    return (EIf (Range pos1 pos2) e1 e2 (EError rng))) $ do
           pKeyword "else"
           noFail $ do
             e3 <- pExpr <|>> expectedExpression
             pos2 <- gets psCur
             return (EIf (Range pos1 pos2) e1 e2 e3)
 
-pExprOpExpr :: Int -> FParser (Expr Range)
+pExprOpExpr :: Int -> FParser PExpr
 pExprOpExpr d = inlineWhite >> pExprOpExpr0 d
 
-pExprOpExpr0 :: Int -> FParser (Expr Range)
+pExprOpExpr0 :: Int -> FParser PExpr
 pExprOpExpr0 d = do
   pos1 <- gets psCur
   e0 <- pEApp0
@@ -478,7 +514,7 @@ climbRight pExpr' pOper makeOp d lhspos lhs mlhsop =
     pos2 <- gets psCur
     climbRight pExpr' pOper makeOp d lhspos (makeOp (Range lhspos pos2) lhs op rhs) (Just paop)
 
-pEApp0 :: FParser (Expr Range)
+pEApp0 :: FParser PExpr
 pEApp0 = do
   pos1 <- gets psCur
   e1 <- pEAtom0
@@ -488,7 +524,7 @@ pEApp0 = do
     [] -> return e1
     _ -> return (EApp (Range pos1 pos2) e1 es)
 
-pEAtom0 :: FParser (Expr Range)
+pEAtom0 :: FParser PExpr
 pEAtom0 = faasum'
   [uncurry ELit <$> ranged pLiteral0
   ,pEList0
@@ -550,7 +586,7 @@ pStringChar = faasum'
                 return '?'
   ,do satisfy (\c -> c `notElem` "\n\r\\\'")]
 
-pEList0 :: FParser (Expr Range)
+pEList0 :: FParser PExpr
 pEList0 = do
   pos1 <- gets psCur
   char '['  -- special syntax, no need for pKeySym
@@ -561,13 +597,13 @@ pEList0 = do
     pos2 <- gets psCur
     return (EList (Range pos1 pos2) es)
 
-pEVarOrCon0 :: FParser (Expr Range)
+pEVarOrCon0 :: FParser PExpr
 pEVarOrCon0 =
   ranged (pIdentifier0 InBlock Don'tCare ()) >>= \case
     (rng, (Lowercase, name)) -> return (EVar rng name)
     (rng, (Uppercase, name)) -> return (ECon rng name)
 
-pEParens0 :: FParser (Expr Range)
+pEParens0 :: FParser PExpr
 pEParens0 = do
   char '('
   noFail $ do
@@ -576,8 +612,11 @@ pEParens0 = do
     char ')' <|>> raise Error "Expected closing ')'"
     return e
 
-expectedExpression :: IParser (Expr Range)
-expectedExpression = raise Error "Expected expression" >> return EParseError
+expectedExpression :: IParser PExpr
+expectedExpression = do
+  pos <- gets psCur
+  raise Error "Expected expression"
+  return (EError (Range pos pos))
 
 data Associativity = AssocLeft | AssocRight | AssocNone
   deriving (Show, Eq)
@@ -609,65 +648,72 @@ pUpperInfixOp0 :: FParser ParsedOperator
 pUpperInfixOp0 =
   faasum' [PaOp OCons 5 AssocRight <$> ranged' (pKeySym0 ":")]
 
-pStandaloneTypesig0 :: FParser (Name, Type)
+pStandaloneTypesig0 :: FParser (Name, PType)
 pStandaloneTypesig0 = do
   name <- pIdentifier0 AtLeft Lowercase WCBacktrack
   inlineWhite
   pKeySym0 "::"
-  noFail $ pushContext ("type signature for " ++ pretty name) $ do
-    ty <- pType <|>> (raise Error "Expected type" >> return (TTup []))
-    return (name, ty)
+  ty <- pType <|>> (raise Error "Expected type" >> faempty)
+  return (name, ty)
 
-pType :: FParser Type
+pType :: FParser PType
 pType = do
+  pos1 <- gets psCur
   ty1 <- pTypeApp
   facatch (return ty1) $ do
     inlineWhite
     pKeySym0 "->"
     noFail $ do
-      ty2 <- pType <|>> (raise Error "Expected type" >> return (TTup []))
-      return (TFun ty1 ty2)
+      pos2 <- gets psCur
+      ty2 <- pType <|>> (raise Error "Expected type" >> return (TTup (Range pos2 pos2) []))
+      return (TFun (Range pos1 pos2) ty1 ty2)
 
-pTypeApp :: FParser Type
+pTypeApp :: FParser PType
 pTypeApp = fasome pTypeAtom >>= \case
   t :| [] -> return t
-  t :| ts -> return (TApp t ts)
+  t :| ts -> return (TApp (foldMapne extOf (t :| ts)) t ts)
 
-pTypeAtom :: FParser Type
+pTypeAtom :: FParser PType
 pTypeAtom = faasum' [pTypeParens, pTypeList, pTypeName]
   where
     pTypeParens = do
       inlineWhite
+      pos1 <- gets psCur
       char '('
       faasum'
         [do inlineWhite
             char ')'
-            return (TTup [])
+            pos2 <- gets psCur
+            return (TTup (Range pos1 pos2) [])
         ,do ty1 <- pType
             noFail $ do
               ty2s <- famany $ do
                 inlineWhite
                 char ','
-                noFail $ pType <|>> (raise Error "Expected type" >> return (TTup []))
+                pos <- gets psCur
+                noFail $ pType <|>> (raise Error "Expected type" >> return (TTup (Range pos pos) []))
               inlineWhite
               char ')' <|>> raise Error "Expected closing ')'"
+              pos2 <- gets psCur
               case ty2s of
                 [] -> return ty1
-                _ -> return (TTup (ty1 : ty2s))]
+                _ -> return (TTup (Range pos1 pos2) (ty1 : ty2s))]
 
     pTypeList = do
       inlineWhite
+      pos1 <- gets psCur
       char '['
       ty <- pType
-      noFail $ char ']' <|>> raise Error "Expecte closing ']'"
-      return (TList ty)
+      noFail $ char ']' <|>> raise Error "Expected closing ']'"
+      pos2 <- gets psCur
+      return (TList (Range pos1 pos2) ty)
 
     pTypeName = do
       inlineWhite
-      (cs, name) <- pIdentifier0 InBlock Don'tCare ()
+      (rng, (cs, name)) <- ranged $ pIdentifier0 InBlock Don'tCare ()
       case cs of
-        Uppercase -> return (TCon name)
-        Lowercase -> return (TVar name)
+        Uppercase -> return (TCon rng name)
+        Lowercase -> return (TVar rng name)
 
 -- | Parse the given name-like keyword, ensuring that it is the entire word.
 pKeyword :: String -> FParser ()
@@ -702,10 +748,6 @@ type family WithCaseOutput care a where
   WithCaseOutput 'True a = a
   WithCaseOutput 'False a = (Case 'True, a)
 
-type family If c a b where
-  If 'True a b = a
-  If 'False a b = b
-
 data WrongCaseBacktrack
   = WCBacktrack  -- ^ If a word was found but it had the wrong case, fail and backtrack.
   | WCAssume     -- ^ Be certain that this case is expected here, and assume incorrect
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