WIP typecheck and other stuff

5 files changed, 261 insertions, 109 deletions

diff --git a/src/Data/Bag.hs b/src/Data/Bag.hs
new file mode 100644
index 0000000..9cdce89
--- /dev/null
+++ b/src/Data/Bag.hs
@@ -0,0 +1,20 @@
+{-# LANGUAGE DeriveTraversable #-}
+module Data.Bag where
+
+
+data Bag a
+  = BTwo (Bag a) (Bag a)
+  | BOne a
+  | BZero
+  deriving (Functor, Foldable, Traversable)
+
+instance Semigroup (Bag a) where (<>) = BTwo
+instance Monoid (Bag a) where mempty = BZero
+
+instance Applicative Bag where
+  pure = BOne
+
+  BZero <*> _ = BZero
+  _ <*> BZero = BZero
+  BOne f <*> b = f <$> b
+  BTwo b1 b2 <*> b = BTwo (b1 <*> b) (b2 <*> b)
diff --git a/src/HSVIS/AST.hs b/src/HSVIS/AST.hs
index 5a90205..f606d22 100644
--- a/src/HSVIS/AST.hs
+++ b/src/HSVIS/AST.hs
@@ -1,3 +1,7 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE StandaloneDeriving #-}
 module HSVIS.AST where
 
 import Data.List.NonEmpty (NonEmpty)
@@ -17,18 +21,27 @@ data DataDef = DataDef Name [Name] [(Name, [Type])]
 data FunDef t = FunDef Name (Maybe Type) (NonEmpty (FunEq t))
   deriving (Show)
 
-data FunEq t = FunEq Name [Pattern t] (RHS t)
+data FunEq t = FunEq t Name [Pattern t] (RHS t)
   deriving (Show)
 
-data Type
-  = TApp Type [Type]
-  | TTup [Type]
-  | TList Type
-  | TFun Type Type
-  | TCon Name
-  | TVar Name
+data TypeStage = TSTC | TSNormal
   deriving (Show)
 
+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
@@ -55,6 +68,7 @@ data 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 Literal = LInt Integer | LFloat Rational | LChar Char | LString String
diff --git a/src/HSVIS/Diagnostic.hs b/src/HSVIS/Diagnostic.hs
index 75a677f..322f9eb 100644
--- a/src/HSVIS/Diagnostic.hs
+++ b/src/HSVIS/Diagnostic.hs
@@ -7,12 +7,15 @@ data Pos = Pos
   { posLine :: Int  -- ^ zero-based
   , posCol :: Int   -- ^ zero-based
   }
-  deriving (Show)
+  deriving (Show, Eq, Ord)
 
 -- | Inclusive-exclusive range of positions in a file.
 data Range = Range Pos Pos
   deriving (Show)
 
+instance Semigroup Range where
+  Range a b <> Range c d = Range (min a c) (max b d)
+
 data Diagnostic = Diagnostic
   { dFile :: FilePath      -- ^ The file for which the diagnostic was raised
   , dRange :: Range        -- ^ Where in the file
diff --git a/src/HSVIS/Parser.hs b/src/HSVIS/Parser.hs
index f153e16..04e6a63 100644
--- a/src/HSVIS/Parser.hs
+++ b/src/HSVIS/Parser.hs
@@ -33,6 +33,7 @@ import Data.These
 -- import Debug.Trace
 
 import Control.FAlternative
+import Data.Bag
 import HSVIS.AST
 import HSVIS.Diagnostic
 import HSVIS.Pretty
@@ -64,8 +65,8 @@ newtype Parser fail a = Parser
       :: forall r.
          Context
       -> PS
-      -> (PS -> [Diagnostic] -> a -> r)  -- ^ OK: some diagnostics, but parsing succeeded
-      -> ([Diagnostic] -> r)  -- ^ Fatal: error that prevented parsing from proceeding
+      -> (PS -> Bag Diagnostic -> a -> r)  -- ^ OK: some diagnostics, but parsing succeeded
+      -> (Bag Diagnostic -> r)  -- ^ Fatal: error that prevented parsing from proceeding
       -> BacktrackPath fail r  -- ^ Backtrack: alternative was exhausted without success
       -> r }
 
@@ -77,7 +78,7 @@ instance Functor (Parser fail) where
     g ctx ps (\ps' errs x -> kok ps' errs (f x)) kfat kbt)
 
 instance Applicative (Parser fail) where
-  pure x = Parser (\_ ps kok _ _ -> kok ps [] x)
+  pure x = Parser (\_ ps kok _ _ -> kok ps mempty x)
   (<*>) = ap
 
 instance Monad (Parser fail) where
@@ -108,29 +109,29 @@ instance FAlternative Parser where
 instance MonadState PS (Parser fail) where
   state f = Parser $ \_ ps kok _ _ ->
     let (x, ps') = f ps
-    in kok ps' [] x
+    in kok ps' mempty x
 
 instance MonadReader Context (Parser fail) where
-  reader f = Parser $ \ctx ps kok _ _ -> kok ps [] (f ctx)
+  reader f = Parser $ \ctx ps kok _ _ -> kok ps mempty (f ctx)
   local f (Parser g) = Parser (\ctx -> g (f ctx))
 
-instance KnownFallible fail => MonadChronicle [Diagnostic] (Parser fail) where
+instance KnownFallible fail => MonadChronicle (Bag Diagnostic) (Parser fail) where
   dictate errs = Parser $ \_ ps kok _ _ -> kok ps errs ()
   confess errs = Parser $ \_ _ _ kfat _ -> kfat errs
   memento (Parser f) = Parser $ \ctx ps kok _ kbt ->
     f ctx ps
       (\ps' errs x -> kok ps' errs (Right x))
-      (\errs -> kok ps [] (Left errs))
+      (\errs -> kok ps mempty (Left errs))
       kbt
   absolve def (toFallible -> Parser f) = Parser $ \ctx ps kok _ _ ->
     f ctx ps
       kok
-      (\_ -> kok ps [] def)
-      (kok ps [] def)
+      (\_ -> kok ps mempty def)
+      (kok ps mempty def)
   condemn (Parser f) = Parser $ \ctx ps kok kfat kbt ->
     f ctx ps
       (\ps' errs x -> case errs of
-                        [] -> kok ps' [] x
+                        BZero -> kok ps' mempty x
                         _ -> kfat errs)
       kfat
       kbt
@@ -141,19 +142,19 @@ instance KnownFallible fail => MonadChronicle [Diagnostic] (Parser fail) where
       kbt
   chronicle th = case th of
     This errs -> Parser (\_ _ _ kfat _ -> kfat errs)
-    That res -> Parser (\_ ps kok _ _ -> kok ps [] res)
+    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 ()))
+parse :: FilePath -> String -> ([Diagnostic], Maybe (Program Range))
 parse fp source =
   runParser pProgram (Context fp (lines source) []) (PS (Pos 0 0) (Pos 0 0) source)
     (\_ errs res -> case errs of
-                      [] -> ([], Just res)
-                      _ -> (errs, Just res))
-    (\errs -> (errs, Nothing))
+                      BZero -> ([], Just res)
+                      _ -> (toList errs, Just res))
+    (\errs -> (toList errs, Nothing))
     ()  -- the program parser cannot fail! :D
 
-pProgram :: IParser (Program ())
+pProgram :: IParser (Program Range)
 pProgram = do
   defs <- pTopDefs
   let (datadefs, fundefs) = partitionEithers defs
@@ -161,7 +162,7 @@ pProgram = do
   assertEOF Error
   return (Program datadefs fundefs)
 
-pTopDefs :: IParser [Either DataDef (FunDef ())]
+pTopDefs :: IParser [Either DataDef (FunDef Range)]
 pTopDefs = do
   faoptional pTopDef >>= \case
     Nothing -> do
@@ -176,7 +177,7 @@ pTopDefs = do
       defs2 <- pTopDefs
       return (defs ++ defs2)
 
-pTopDef :: FParser [Either DataDef (FunDef ())]
+pTopDef :: FParser [Either DataDef (FunDef Range)]
 pTopDef = do
   noFail skipWhiteComment
   noFail isAtBlockLeft >>= \case
@@ -204,7 +205,7 @@ pDataDef0 = do
     pDatacons leader = do
       inlineWhite
       facatch (return []) $ do
-        pKeySym leader
+        pKeySym0 leader
         inlineWhite
         name <- pIdentifier0 InBlock Uppercase WCAssume
         fields <- noFail $ famany pTypeAtom
@@ -217,7 +218,7 @@ data FunEqContext
   | Continue Name
   deriving (Show)
 
-pFunDef0 :: FParser [FunDef ()]
+pFunDef0 :: FParser [FunDef Range]
 pFunDef0 =
   faasum'
     [do (name, typ) <- pStandaloneTypesig0
@@ -227,23 +228,27 @@ pFunDef0 =
               raise Error $ "Expected function equation for " ++ pretty name ++
                             " after type signature"
               return []
-            Just [] -> return [FunDef name (Just typ) (FunEq name [] (Plain (ETup () [])) :| [])]
+            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)]
     ,do pFunEq FirstLine >>= \case
-          clause1@(FunEq name _ _) : clauses1 -> noFail $ do
+          clause1@(FunEq _ name _ _) : clauses1 -> noFail $ do
             clauses <- concat <$> famany (pFunEq (Continue name))
             return [FunDef name Nothing (clause1 :| clauses1 ++ clauses)]
           [] -> faempty]
 
 -- | Given the name from the type signature or a previous clause, if any.
-pFunEq :: FunEqContext -> FParser [FunEq ()]
+pFunEq :: FunEqContext -> FParser [FunEq Range]
 pFunEq fectx = do
   noFail skipWhiteComment
   faguardM isAtBlockLeft
 
   pushLocatedContext "function equation" $ do
+    pos1 <- gets psCur
     name <- pIdentifier0 AtLeft Lowercase WCAssume
 
     -- We want to do various checks with what came before, and there are
@@ -265,27 +270,29 @@ pFunEq fectx = do
       then do
         pats <- famany (pPattern 11)
         rhs <- pRHS "="
-        return [FunEq name pats rhs]
+        pos2 <- gets psCur
+        return [FunEq (Range pos1 pos2) name pats rhs]
       else return []
 
 -- | Pass "=" for function definitions and "->" for case clauses.
-pRHS :: String -> IParser (RHS ())
+pRHS :: String -> IParser (RHS Range)
 pRHS sepsym = do
   -- TODO: parse guards
   inlineWhite
-  pKeySym sepsym <|>> raise Error ("Expected " ++ show sepsym)
-  expr <- pExpr <|>> (raise Error "Expected expression" >> return (ETup () []))
+  pKeySym0 sepsym <|>> raise Error ("Expected " ++ show sepsym)
+  expr <- pExpr <|>> expectedExpression
   return (Plain expr)
 
-pPattern :: Int -> FParser (Pattern ())
+pPattern :: Int -> FParser (Pattern Range)
 pPattern d = inlineWhite >> pPattern0 d
 
-pPattern0 :: Int -> FParser (Pattern ())
+pPattern0 :: Int -> FParser (Pattern Range)
 pPattern0 d = do
+  pos1 <- gets psCur
   p0 <- pPatExprAtom0 (max 10 d)
-  climbRight pPattern (pInfixOp Uppercase) (POp ()) d p0 Nothing
+  climbRight pPattern (pInfixOp Uppercase) POp d pos1 p0 Nothing
 
-pExpr :: FParser (Expr ())
+pExpr :: FParser (Expr Range)
 pExpr = do
   inlineWhite
   -- basics: lit, list, var, con, tup
@@ -298,7 +305,7 @@ pExpr = do
             ,pEIf0
             ,pExprOpExpr0 0]
 
-pPatExprAtom0 :: Int -> FParser (Pattern ())
+pPatExprAtom0 :: Int -> FParser (Pattern Range)
 pPatExprAtom0 d =
   faasum' [pPatWildcard0
           ,pPatVarOrAs0
@@ -306,34 +313,43 @@ pPatExprAtom0 d =
           ,pPatList0
           ,pPatParens0]
   where
-    pPatWildcard0 = pKeySym "_" >> return (PWildcard ())
+    pPatWildcard0 = PWildcard <$> ranged' (pKeySym0 "_")
     pPatVarOrAs0 = do
-      var <- pIdentifier0 InBlock Lowercase WCBacktrack
-      facatch (return (PVar () var)) $ do
+      (varrng@(Range pos1 _), var) <- ranged $ pIdentifier0 InBlock Lowercase WCBacktrack
+      facatch (return (PVar varrng var)) $ do
         inlineWhite
-        pKeySym "@"
+        pKeySym0 "@"
         noFail $ do
-          p <- pPattern 11 <|>> (raise Error "Expected pattern after '@'" >> return (PWildcard ()))
-          return (PAs () var p)
+          pos <- gets psCur
+          facatch (do raise Error "Expected pattern after '@'"
+                      return (PWildcard (Range pos pos))) $ do
+            pat <- pPattern 11
+            pos2 <- gets psCur
+            return (PAs (Range pos1 pos2) var pat)
     pPatCon0 = do
-      con <- pIdentifier0 InBlock Uppercase WCBacktrack
+      (conrng@(Range pos1 _), con) <- ranged $ pIdentifier0 InBlock Uppercase WCBacktrack
       noFail $ if d > 10
-        then return (PCon () con [])
+        then return (PCon conrng con [])
         else do args <- famany (pPattern 11)
-                return (PCon () con args)
+                pos2 <- gets psCur
+                return (PCon (Range pos1 pos2) con args)
     pPatList0 = do
+      pos1 <- gets psCur
       char '['  -- special syntax, no need for pKeySym
       noFail $ do
         ps <- pPattern 0 `sepBy` (inlineWhite >> char ',')
         inlineWhite
         char ']' <|>> raise Error "Expected ']'"
-        return (PList () ps)
+        pos2 <- gets psCur
+        return (PList (Range pos1 pos2) ps)
     pPatParens0 = do
+      pos1 <- gets psCur
       char '('
       inlineWhite
       faasum'
         [do char ')'
-            return (PTup () [])
+            pos2 <- gets psCur
+            return (PTup (Range pos1 pos2) [])
         ,do p <- pPattern0 0
             inlineWhite
             faasum'
@@ -341,10 +357,14 @@ pPatExprAtom0 d =
                   return p
               ,do char ','
                   ps <- pPattern 0 `sepBy1` (inlineWhite >> char ',')
-                  return (PTup () (p : ps))]]
+                  inlineWhite
+                  char ')' <|>> raise Error "Expected ')'"
+                  pos2 <- gets psCur
+                  return (PTup (Range pos1 pos2) (p : ps))]]
 
-pELet0 :: FParser (Expr ())
+pELet0 :: FParser (Expr Range)
 pELet0 = do
+  pos1 <- gets psCur
   pKeyword "let"
   noFail $ do
     inlineWhite
@@ -366,21 +386,26 @@ pELet0 = do
 
     let defs = concat defss
     inlineWhite
-    facatch (do raise Error "Expected 'in' after 'let'"
-                return (ELet () defs (ETup () []))) $ do
+    facatch (do pos2 <- gets psCur
+                raise Error "Expected 'in' after 'let'"
+                return (ELet (Range pos1 pos2) defs EParseError)) $ do
       pKeyword "in"
       noFail $ do
         inlineWhite
-        body <- pExpr <|>> (raise Error "Expected expression" >> return (ETup () []))
-        return (ELet () defs body)
+        body <- pExpr <|>> expectedExpression
+        pos2 <- gets psCur
+        return (ELet (Range pos1 pos2) defs body)
 
-pECase0 :: FParser (Expr ())
+pECase0 :: FParser (Expr Range)
 pECase0 = do
+  pos1 <- gets psCur
   pKeyword "case"
   noFail $ do
-    e <- pExpr <|>> (raise Error "Expected expression" >> return (ETup () []))
+    e <- pExpr <|>> expectedExpression
     inlineWhite
-    facatch (raise Error "Expected 'of' after 'case'" >> return (ECase () e [])) $ do
+    facatch (do pos2 <- gets psCur
+                raise Error "Expected 'of' after 'case'"
+                return (ECase (Range pos1 pos2) e [])) $ do
       pKeyword "of"
       noFail $ do
         inlineWhite
@@ -395,65 +420,77 @@ pECase0 = do
                   rhs <- pRHS "->"
                   return (pat, rhs)
           clauses <- famany pClause
-          return (ECase () e clauses)
+          pos2 <- gets psCur
+          return (ECase (Range pos1 pos2) e clauses)
 
-pEIf0 :: FParser (Expr ())
+pEIf0 :: FParser (Expr Range)
 pEIf0 = do
+  pos1 <- gets psCur
   pKeyword "if"
   noFail $ do
-    e1 <- pExpr <|>> (raise Error "Expected expression" >> return (ECon () (Name "True")))
+    e1 <- pExpr <|>> expectedExpression
     inlineWhite
-    facatch (raise Error "Expected 'then' after 'if'" >> return (EIf () e1 (ETup () []) (ETup () []))) $ do
+    facatch (do pos2 <- gets psCur
+                raise Error "Expected 'then' after 'if'"
+                return (EIf (Range pos1 pos2) e1 EParseError EParseError)) $ do
       pKeyword "then"
       noFail $ do
-        e2 <- pExpr <|>> (raise Error "Expected expression" >> return (ETup () []))
+        e2 <- pExpr <|>> expectedExpression
         inlineWhite
-        facatch (raise Error "Expected else after 'if'" >> return (EIf () e1 (ETup () []) (ETup () []))) $ do
+        facatch (do pos2 <- gets psCur
+                    raise Error "Expected else after 'then'"
+                    return (EIf (Range pos1 pos2) e1 e2 EParseError)) $ do
           pKeyword "else"
           noFail $ do
-            e3 <- pExpr <|>> (raise Error "Expected expression" >> return (ETup () []))
-            return (EIf () e1 e2 e3)
+            e3 <- pExpr <|>> expectedExpression
+            pos2 <- gets psCur
+            return (EIf (Range pos1 pos2) e1 e2 e3)
 
-pExprOpExpr :: Int -> FParser (Expr ())
+pExprOpExpr :: Int -> FParser (Expr Range)
 pExprOpExpr d = inlineWhite >> pExprOpExpr0 d
 
-pExprOpExpr0 :: Int -> FParser (Expr ())
+pExprOpExpr0 :: Int -> FParser (Expr Range)
 pExprOpExpr0 d = do
+  pos1 <- gets psCur
   e0 <- pEApp0
-  climbRight pExprOpExpr (snd <$> pInfixOp Don'tCare) (EOp ()) d e0 Nothing
+  climbRight pExprOpExpr (snd <$> pInfixOp Don'tCare) EOp d pos1 e0 Nothing
 
 climbRight
   :: (Int -> FParser e)  -- ^ Parse an expression at the given precedence level
   -> FParser ParsedOperator  -- ^ Parse an operator
-  -> (e -> Operator -> e -> e)  -- ^ Build an operator application experssion
+  -> (Range -> e -> Operator -> e -> e)  -- ^ Build an operator application experssion
   -> Int  -- ^ Ambient precedence level: minimum precedence of top-level operator in result
+  -> Pos  -- ^ Start of lhs
   -> e  -- ^ lhs: Initial non-operator expression already parsed
   -> Maybe ParsedOperator  -- ^ Top-level operator in lhs (initialise with Nothing)
   -> FParser e
-climbRight pExpr' pOper makeOp d lhs mlhsop =
+climbRight pExpr' pOper makeOp d lhspos lhs mlhsop =
   facatch (return lhs) $ do
-    paop@(PaOp op d2 a2) <- pOper
+    paop@(PaOp op d2 a2 _) <- pOper
     faguard (d2 >= d)  -- respect global minimum precedence
     case mlhsop of  -- check operator compatibility
-      Just (PaOp _ d1 a1) ->
+      Just (PaOp _ d1 a1 _) ->
         faguard (d1 > d2 || (d1 == d2 && a1 == a2 && a1 /= AssocNone))
       Nothing ->
         return ()
     let oprhsd = case a2 of AssocRight -> d2 ; _ -> d2 + 1
     rhs <- pExpr' oprhsd
-    climbRight pExpr' pOper makeOp d (makeOp lhs op rhs) (Just paop)
+    pos2 <- gets psCur
+    climbRight pExpr' pOper makeOp d lhspos (makeOp (Range lhspos pos2) lhs op rhs) (Just paop)
 
-pEApp0 :: FParser (Expr ())
+pEApp0 :: FParser (Expr Range)
 pEApp0 = do
+  pos1 <- gets psCur
   e1 <- pEAtom0
   es <- noFail $ famany (inlineWhite >> pEAtom0)
+  pos2 <- gets psCur
   case es of
     [] -> return e1
-    _ -> return (EApp () e1 es)
+    _ -> return (EApp (Range pos1 pos2) e1 es)
 
-pEAtom0 :: FParser (Expr ())
+pEAtom0 :: FParser (Expr Range)
 pEAtom0 = faasum'
-  [ELit () <$> pLiteral0
+  [uncurry ELit <$> ranged pLiteral0
   ,pEList0
   ,pEVarOrCon0 
   ,pEParens0]
@@ -513,34 +550,39 @@ pStringChar = faasum'
                 return '?'
   ,do satisfy (\c -> c `notElem` "\n\r\\\'")]
 
-pEList0 :: FParser (Expr ())
+pEList0 :: FParser (Expr Range)
 pEList0 = do
+  pos1 <- gets psCur
   char '['  -- special syntax, no need for pKeySym
   noFail $ do
     es <- sepBy pExpr (inlineWhite >> char ',')
     inlineWhite
     char ']' <|>> raise Error "Expected closing ']'"
-    return (EList () es)
+    pos2 <- gets psCur
+    return (EList (Range pos1 pos2) es)
 
-pEVarOrCon0 :: FParser (Expr ())
+pEVarOrCon0 :: FParser (Expr Range)
 pEVarOrCon0 =
-  pIdentifier0 InBlock Don'tCare () >>= \case
-    (Lowercase, name) -> return (EVar () name)
-    (Uppercase, name) -> return (ECon () name)
+  ranged (pIdentifier0 InBlock Don'tCare ()) >>= \case
+    (rng, (Lowercase, name)) -> return (EVar rng name)
+    (rng, (Uppercase, name)) -> return (ECon rng name)
 
-pEParens0 :: FParser (Expr ())
+pEParens0 :: FParser (Expr Range)
 pEParens0 = do
   char '('
   noFail $ do
-    e <- pExpr <|>> (raise Error "Expected expression" >> return (ETup () []))
+    e <- pExpr <|>> expectedExpression
     inlineWhite
     char ')' <|>> raise Error "Expected closing ')'"
     return e
 
+expectedExpression :: IParser (Expr Range)
+expectedExpression = raise Error "Expected expression" >> return EParseError
+
 data Associativity = AssocLeft | AssocRight | AssocNone
   deriving (Show, Eq)
 
-data ParsedOperator = PaOp Operator Int Associativity
+data ParsedOperator = PaOp Operator Int Associativity Range
   deriving (Show)
 
 pInfixOp :: Case care -> FParser (WithCaseOutput care ParsedOperator)
@@ -554,24 +596,24 @@ pInfixOp cs = do
 
 pLowerInfixOp0 :: FParser ParsedOperator
 pLowerInfixOp0 =
-  faasum' [PaOp OEqu 4 AssocNone  <$ pKeySym "=="
-          ,PaOp OAdd 6 AssocLeft  <$ pKeySym "+"
-          ,PaOp OSub 6 AssocLeft  <$ pKeySym "-"
-          ,PaOp OMul 7 AssocLeft  <$ pKeySym "*"
-          ,PaOp ODiv 7 AssocLeft  <$ pKeySym "/"
-          ,PaOp OMod 7 AssocLeft  <$ pKeySym "%"
-          ,PaOp OPow 8 AssocRight <$ pKeySym "^"
+  faasum' [PaOp OEqu 4 AssocNone  <$> ranged' (pKeySym0 "==")
+          ,PaOp OAdd 6 AssocLeft  <$> ranged' (pKeySym0 "+")
+          ,PaOp OSub 6 AssocLeft  <$> ranged' (pKeySym0 "-")
+          ,PaOp OMul 7 AssocLeft  <$> ranged' (pKeySym0 "*")
+          ,PaOp ODiv 7 AssocLeft  <$> ranged' (pKeySym0 "/")
+          ,PaOp OMod 7 AssocLeft  <$> ranged' (pKeySym0 "%")
+          ,PaOp OPow 8 AssocRight <$> ranged' (pKeySym0 "^")
           ]
 
 pUpperInfixOp0 :: FParser ParsedOperator
 pUpperInfixOp0 =
-  faasum' [PaOp OCons 5 AssocRight <$ pKeySym ":"]
+  faasum' [PaOp OCons 5 AssocRight <$> ranged' (pKeySym0 ":")]
 
 pStandaloneTypesig0 :: FParser (Name, Type)
 pStandaloneTypesig0 = do
   name <- pIdentifier0 AtLeft Lowercase WCBacktrack
   inlineWhite
-  pKeySym "::"
+  pKeySym0 "::"
   noFail $ pushContext ("type signature for " ++ pretty name) $ do
     ty <- pType <|>> (raise Error "Expected type" >> return (TTup []))
     return (name, ty)
@@ -581,7 +623,7 @@ pType = do
   ty1 <- pTypeApp
   facatch (return ty1) $ do
     inlineWhite
-    pKeySym "->"
+    pKeySym0 "->"
     noFail $ do
       ty2 <- pType <|>> (raise Error "Expected type" >> return (TTup []))
       return (TFun ty1 ty2)
@@ -635,11 +677,21 @@ pKeyword s = do
   notFollowedBy (() <$ satisfy isNameContChar)
 
 -- | Parse the given symbol-like keyword, ensuring that it is the entire symbol.
-pKeySym :: String -> FParser ()
-pKeySym s = do
+pKeySym0 :: String -> FParser ()
+pKeySym0 s = do
   string s
   notFollowedBy (() <$ satisfy isSymbolChar)
 
+ranged :: Parser fail a -> Parser fail (Range, a)
+ranged p = do
+  pos1 <- gets psCur
+  res <- p
+  pos2 <- gets psCur
+  return (Range pos1 pos2, res)
+
+ranged' :: Parser fail () -> Parser fail (Range)
+ranged' p = fst <$> ranged p
+
 data Case care where
   Uppercase :: Case 'True
   Lowercase :: Case 'True
@@ -777,7 +829,7 @@ startLayoutBlock p = do
 
 data Fatality fatal where
   Error :: Fatality 'False
-  Fatal :: Fatality 'True
+  -- Fatal :: Fatality 'True
 deriving instance Show (Fatality fatal)
 
 type family FatalCtx fatal a where
@@ -786,7 +838,7 @@ type family FatalCtx fatal a where
 
 raise_ :: KnownFallible fail => Fatality fatal -> String -> Parser fail ()
 raise_ Error = raise Error
-raise_ Fatal = raise Fatal
+-- raise_ Fatal = raise Fatal
 
 raise :: (KnownFallible fail, FatalCtx fatal a) => Fatality fatal -> String -> Parser fail a
 raise fat msg = gets psCur >>= \pos -> raiseAt pos fat msg
@@ -798,7 +850,7 @@ raiseAt pos fat msg = do
   let err = Diagnostic fp (Range pos pos) stk (srcLines !! posLine pos) msg
   case fat of
     Error -> dictate (pure err)
-    Fatal -> confess (pure err)
+    -- Fatal -> confess (pure err)
 
 describeLocation :: IParser String
 describeLocation = do
@@ -861,7 +913,7 @@ skipWhiteComment :: IParser ()
 skipWhiteComment = do
   inlineSpaces
   _ <- famany (blockComment >> noFail inlineSpaces)
-  optional_ lineComment
+  optional_ lineComment0
   optional_ (consumeNewline >> noFail skipWhiteComment)
   where
     -- | Consumes some inline whitespace. Stops before newlines.
@@ -883,10 +935,10 @@ blockComment = do
 
 -- | Consumes a line comment marker and the rest of the line, excluding
 -- newline.
-lineComment :: FParser ()
-lineComment = do
+lineComment0 :: FParser ()
+lineComment0 = do
   -- '--!' is an operator, so we need to parse a whole symbol here.
-  pKeySym "--"
+  pKeySym0 "--"
   noFail $ readWhileInline (const True)
 
 -- | Raises an error if we're not currently at EOF.
diff --git a/src/HSVIS/Typecheck.hs b/src/HSVIS/Typecheck.hs
new file mode 100644
index 0000000..b1ffbb9
--- /dev/null
+++ b/src/HSVIS/Typecheck.hs
@@ -0,0 +1,63 @@
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+module HSVIS.Typecheck where
+
+import Control.Monad
+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.Bag
+import HSVIS.AST
+import HSVIS.Diagnostic
+
+
+typecheck :: FilePath -> String -> Program Range -> ([Diagnostic], Program Type)
+typecheck fp source prog =
+  let (progtc, (ds1, cs)) =
+        runWriter
+        . flip evalStateT 1
+        . flip runReaderT (fp, source)
+        . runTCM
+        $ tcProgram prog
+      (ds2, sub) = solveConstrs cs
+  in (toList (ds1 <> ds2), substProg sub progtc)
+
+data Constr =
+  CEq Type Type Range  -- ^ These types must be equal because of the expression here
+
+newtype TCM a = TCM { runTCM :: ReaderT (FilePath, String) (StateT Int (Writer (Bag Diagnostic, Bag Constr))) a }
+  deriving newtype (Functor, Applicative, Monad)
+
+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)
+
+tcProgram :: Program Range -> TCM (Program TCType)
+tcProgram (Program ddefs fdefs) = Program ddefs <$> traverse tcFunDef fdefs
+
+tcFunDef :: FunDef Range -> TCM (FunDef TCType)
+tcFunDef (FunDef name msig eqs) = do
+  when (not $ allEq (fmap (length . funeqPats) eqs)) $
+    raise (sconcatne (fmap funeqRange eqs)) "Function equations have differing numbers of arguments"
+
+  _
+
+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