Trees that explode

1 files changed, 105 insertions, 63 deletions

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