1 files changed, 12 insertions, 9 deletions
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: " ++