Partial application in expression statements

2 files changed, 53 insertions, 20 deletions

diff --git a/src/Parser/Expr.hs b/src/Parser/Expr.hs
index 322bb9b..4ed0215 100644
--- a/src/Parser/Expr.hs
+++ b/src/Parser/Expr.hs
@@ -10,6 +10,7 @@ module Parser.Expr (
     literal,
     variable,
 
+    checkFunctionArguments,
     functionArguments,
 ) where
 
@@ -344,20 +345,8 @@ variable = label "variable" $ do
     sline <- getSourceLine
     name <- varName
     lookupVarExpr off sline name >>= \case
-        SomeExpr e'@(FunVariable (FunctionArguments argTypes) _ _) -> do
-            let check poff kw expr = do
-                    case M.lookup kw argTypes of
-                        Just (SomeArgumentType (_ :: ArgumentType expected)) -> do
-                            withRecovery registerParseError $ do
-                                void $ unify poff (ExprTypePrim (Proxy @expected)) (someExprType expr)
-                            return expr
-                        Nothing -> do
-                            registerParseError $ FancyError poff $ S.singleton $ ErrorFail $ T.unpack $
-                                case kw of
-                                    Just (ArgumentKeyword tkw) -> "unexpected parameter with keyword `" <> tkw <> "'"
-                                    Nothing                    -> "unexpected parameter"
-                            return expr
-
+        SomeExpr e'@(FunVariable argTypes _ _) -> do
+            let check = checkFunctionArguments argTypes
             args <- functionArguments check someExpr literal (\poff -> lookupVarExpr poff sline . VarName)
             return $ SomeExpr $ ArgsApp args e'
         e -> do
@@ -378,6 +367,22 @@ variable = label "variable" $ do
     applyRecordSelector m e (RecordSelector f) = SomeExpr $ App (AnnRecord m) (pure f) e
 
 
+checkFunctionArguments :: FunctionArguments SomeArgumentType
+                       -> Int -> Maybe ArgumentKeyword -> SomeExpr -> TestParser SomeExpr
+checkFunctionArguments (FunctionArguments argTypes) poff kw expr = do
+    case M.lookup kw argTypes of
+        Just (SomeArgumentType (_ :: ArgumentType expected)) -> do
+            withRecovery registerParseError $ do
+                void $ unify poff (ExprTypePrim (Proxy @expected)) (someExprType expr)
+            return expr
+        Nothing -> do
+            registerParseError $ FancyError poff $ S.singleton $ ErrorFail $ T.unpack $
+                case kw of
+                    Just (ArgumentKeyword tkw) -> "unexpected parameter with keyword `" <> tkw <> "'"
+                    Nothing                    -> "unexpected parameter"
+            return expr
+
+
 functionArguments :: (Int -> Maybe ArgumentKeyword -> a -> TestParser b) -> TestParser a -> TestParser a -> (Int -> Text -> TestParser a) -> TestParser (FunctionArguments b)
 functionArguments check param lit promote = do
     args <- parseArgs True
diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index 2d45a21..912366b 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -65,9 +65,34 @@ forStatement = do
         return [For line tname (unpack <$> e) body]
 
 exprStatement :: TestParser [ TestStep ]
-exprStatement  = do
-    expr <- typedExpr
-    return [ ExprStatement expr ]
+exprStatement = do
+    ref <- L.indentLevel
+    off <- stateOffset <$> getParserState
+    SomeExpr expr <- someExpr
+    choice
+        [ do
+            continuePartial off ref expr
+        , do
+            stmt <- unifyExpr off Proxy expr
+            return [ ExprStatement stmt ]
+        ]
+  where
+    continuePartial :: ExprType a => Int -> Pos -> Expr a -> TestParser [ TestStep ]
+    continuePartial off ref expr = do
+        symbol ":"
+        void eol
+        (fun :: Expr (FunctionType TestBlock)) <- unifyExpr off Proxy expr
+        scn
+        indent <- L.indentGuard scn GT ref
+        blockOf indent $ do
+            coff <- stateOffset <$> getParserState
+            sline <- getSourceLine
+            args <- functionArguments (checkFunctionArguments (exprArgs fun)) someExpr literal (\poff -> lookupVarExpr poff sline . VarName)
+            let fun' = ArgsApp args fun
+            choice
+                [ continuePartial coff indent fun'
+                , (: []) . ExprStatement <$> unifyExpr coff Proxy fun'
+                ]
 
 class (Typeable a, Typeable (ParamRep a)) => ParamType a where
     type ParamRep a :: Type
@@ -327,8 +352,11 @@ testPacketLoss = command "packet_loss" $ PacketLoss
     <*> innerBlock
 
 
-testBlock :: Pos -> TestParser [TestStep]
-testBlock indent = concat <$> go
+testBlock :: Pos -> TestParser [ TestStep ]
+testBlock indent = blockOf indent testStep
+
+blockOf :: Pos -> TestParser [ a ] -> TestParser [ a ]
+blockOf indent step = concat <$> go
   where
     go = do
         scn
@@ -336,7 +364,7 @@ testBlock indent = concat <$> go
         optional eof >>= \case
             Just _ -> return []
             _ | pos <  indent -> return []
-              | pos == indent -> (:) <$> testStep <*> go
+              | pos == indent -> (:) <$> step <*> go
               | otherwise     -> L.incorrectIndent EQ indent pos
 
 testStep :: TestParser [TestStep]