Implement "guard" as a builtin

6 files changed, 72 insertions, 60 deletions

diff --git a/src/Parser/Core.hs b/src/Parser/Core.hs
index dd2df12..ab6079a 100644
--- a/src/Parser/Core.hs
+++ b/src/Parser/Core.hs
@@ -55,12 +55,12 @@ lookupVarType off name = do
             gets (fromMaybe t . M.lookup tvar . testTypeUnif)
         Just x -> return x
 
-lookupVarExpr :: Int -> VarName -> TestParser SomeExpr
-lookupVarExpr off name = do
+lookupVarExpr :: Int -> SourceLine -> VarName -> TestParser SomeExpr
+lookupVarExpr off sline name = do
     lookupVarType off name >>= \case
-        ExprTypePrim (Proxy :: Proxy a) -> return $ SomeExpr $ (Variable name :: Expr a)
-        ExprTypeVar tvar -> return $ SomeExpr $ DynVariable tvar name
-        ExprTypeFunction args (_ :: Proxy a) -> return $ SomeExpr $ (FunVariable args name :: Expr (FunctionType a))
+        ExprTypePrim (Proxy :: Proxy a) -> return $ SomeExpr $ (Variable sline name :: Expr a)
+        ExprTypeVar tvar -> return $ SomeExpr $ DynVariable tvar sline name
+        ExprTypeFunction args (_ :: Proxy a) -> return $ SomeExpr $ (FunVariable args sline name :: Expr (FunctionType a))
 
 unify :: Int -> SomeExprType -> SomeExprType -> TestParser SomeExprType
 unify _ (ExprTypeVar aname) (ExprTypeVar bname) | aname == bname = do
@@ -127,10 +127,10 @@ unifyExpr off pa expr = if
     | Just (Refl :: a :~: b) <- eqT
     -> return expr
 
-    | DynVariable tvar name <- expr
+    | DynVariable tvar sline name <- expr
     -> do
         _ <- unify off (ExprTypePrim (Proxy :: Proxy a)) (ExprTypeVar tvar)
-        return $ Variable name
+        return $ Variable sline name
 
     | Just (Refl :: FunctionType a :~: b) <- eqT
     -> do
@@ -198,3 +198,13 @@ listOf :: TestParser a -> TestParser [a]
 listOf item = do
     x <- item
     (x:) <$> choice [ symbol "," >> listOf item, return [] ]
+
+
+getSourceLine :: TestParser SourceLine
+getSourceLine = do
+    pstate <- statePosState <$> getParserState
+    return $ SourceLine $ T.concat
+        [ T.pack $ sourcePosPretty $ pstateSourcePos pstate
+        , T.pack ": "
+        , TL.toStrict $ TL.takeWhile (/='\n') $ pstateInput pstate
+        ]
diff --git a/src/Parser/Expr.hs b/src/Parser/Expr.hs
index 04035c1..8ae0f77 100644
--- a/src/Parser/Expr.hs
+++ b/src/Parser/Expr.hs
@@ -66,8 +66,9 @@ someExpansion = do
     void $ char '$'
     choice
         [do off <- stateOffset <$> getParserState
+            sline <- getSourceLine
             name <- VarName . TL.toStrict <$> takeWhile1P Nothing (\x -> isAlphaNum x || x == '_')
-            lookupVarExpr off name
+            lookupVarExpr off sline name
         , between (char '{') (char '}') someExpr
         ]
 
@@ -348,9 +349,10 @@ literal = label "literal" $ choice
 variable :: TestParser SomeExpr
 variable = label "variable" $ do
     off <- stateOffset <$> getParserState
+    sline <- getSourceLine
     name <- varName
-    lookupVarExpr off name >>= \case
-        SomeExpr e'@(FunVariable (FunctionArguments argTypes) _) -> do
+    lookupVarExpr off sline name >>= \case
+        SomeExpr e'@(FunVariable (FunctionArguments argTypes) _ _) -> do
             let check poff kw expr = do
                     case M.lookup kw argTypes of
                         Just expected -> do
@@ -364,7 +366,7 @@ variable = label "variable" $ do
                                     Nothing                    -> "unexpected parameter"
                             return expr
 
-            args <- functionArguments check someExpr literal (\poff -> lookupVarExpr poff . VarName)
+            args <- functionArguments check someExpr literal (\poff -> lookupVarExpr poff sline . VarName)
             return $ SomeExpr $ ArgsApp args e'
         e -> do
             return e
diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index 94a5583..6434a53 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -8,9 +8,8 @@ import Control.Monad.State
 
 import Data.Kind
 import Data.Maybe
-import qualified Data.Set as S
+import Data.Set qualified as S
 import Data.Text qualified as T
-import qualified Data.Text.Lazy as TL
 import Data.Typeable
 
 import Text.Megaparsec hiding (State)
@@ -24,16 +23,6 @@ import Process (Process)
 import Test
 import Util
 
-getSourceLine :: TestParser SourceLine
-getSourceLine = do
-    pstate <- statePosState <$> getParserState
-    return $ SourceLine $ T.concat
-        [ T.pack $ sourcePosPretty $ pstateSourcePos pstate
-        , T.pack ": "
-        , TL.toStrict $ TL.takeWhile (/='\n') $ pstateInput pstate
-        ]
-
-
 letStatement :: TestParser [TestStep]
 letStatement = do
     line <- getSourceLine
@@ -313,11 +302,6 @@ testFlush = command "flush" $ Flush
     <$> paramOrContext "from"
     <*> param ""
 
-testGuard :: TestParser [TestStep]
-testGuard = command "guard" $ Guard
-    <$> cmdLine
-    <*> param ""
-
 testDisconnectNode :: TestParser [TestStep]
 testDisconnectNode = command "disconnect_node" $ DisconnectNode
     <$> paramOrContext ""
@@ -364,7 +348,6 @@ testStep = choice
     , testSend
     , testExpect
     , testFlush
-    , testGuard
     , testDisconnectNode
     , testDisconnectNodes
     , testDisconnectUpstream
diff --git a/src/Run.hs b/src/Run.hs
index 24bba48..a1692cb 100644
--- a/src/Run.hs
+++ b/src/Run.hs
@@ -188,7 +188,7 @@ evalSteps = mapM_ $ \case
 
 
 withVar :: ExprType e => VarName -> e -> TestRun a -> TestRun a
-withVar name value = local (fmap $ \s -> s { tsVars = ( name, SomeVarValue mempty $ const value ) : tsVars s })
+withVar name value = local (fmap $ \s -> s { tsVars = ( name, SomeVarValue mempty $ const $ const value ) : tsVars s })
 
 withInternet :: (Network -> TestRun a) -> TestRun a
 withInternet inner = do
@@ -280,7 +280,10 @@ exprFailed desc (SourceLine sline) pname expr = do
     exprVars <- gatherVars expr
     outLine OutputMatchFail (Just prompt) $ T.concat [desc, T.pack " failed on ", sline]
     forM_ exprVars $ \((name, sel), value) ->
-        outLine OutputMatchFail (Just prompt) $ T.concat ["  ", textVarName name, T.concat (map ("."<>) sel), " = ", textSomeVarValue value]
+        outLine OutputMatchFail (Just prompt) $ T.concat
+            [ "  ", textVarName name, T.concat (map ("."<>) sel)
+            , " = ", textSomeVarValue (SourceLine sline) value
+            ]
     throwError Failed
 
 expect :: SourceLine -> Process -> Expr Regex -> [TypedVarName Text] -> TestRun () -> TestRun ()
@@ -310,7 +313,7 @@ expect (SourceLine sline) p expr tvars inner = do
                      throwError Failed
 
              outProc OutputMatch p line
-             local (fmap $ \s -> s { tsVars = zip vars (map (SomeVarValue mempty . const) capture) ++ tsVars s }) inner
+             local (fmap $ \s -> s { tsVars = zip vars (map (SomeVarValue mempty . const . const) capture) ++ tsVars s }) inner
 
          Nothing -> exprFailed (T.pack "expect") (SourceLine sline) (Just $ procName p) expr
 
diff --git a/src/Test.hs b/src/Test.hs
index 8c5a3ef..b0a91bd 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -160,8 +160,8 @@ data SomeExprType
     | forall a. ExprType a => ExprTypeFunction (FunctionArguments SomeExprType) (Proxy a)
 
 someExprType :: SomeExpr -> SomeExprType
-someExprType (SomeExpr (DynVariable tvar _)) = ExprTypeVar tvar
-someExprType (SomeExpr fun@(FunVariable params _)) = ExprTypeFunction params (proxyOfFunctionType fun)
+someExprType (SomeExpr (DynVariable tvar _ _)) = ExprTypeVar tvar
+someExprType (SomeExpr fun@(FunVariable params _ _)) = ExprTypeFunction params (proxyOfFunctionType fun)
   where
     proxyOfFunctionType :: Expr (FunctionType a) -> Proxy a
     proxyOfFunctionType _ = Proxy
@@ -173,24 +173,24 @@ textSomeExprType (ExprTypeVar (TypeVar name)) = name
 textSomeExprType (ExprTypeFunction _ r) = "function:" <> textExprType r
 
 
-data SomeVarValue = forall a. ExprType a => SomeVarValue (FunctionArguments SomeExprType) (FunctionArguments SomeExpr -> a)
+data SomeVarValue = forall a. ExprType a => SomeVarValue (FunctionArguments SomeExprType) (SourceLine -> FunctionArguments SomeExpr -> a)
 
-fromSomeVarValue :: forall a m. (ExprType a, MonadFail m) => VarName -> SomeVarValue -> m a
-fromSomeVarValue name (SomeVarValue args (value :: args -> b)) = do
+fromSomeVarValue :: forall a m. (ExprType a, MonadFail m) => SourceLine -> VarName -> SomeVarValue -> m a
+fromSomeVarValue sline name (SomeVarValue args (value :: SourceLine -> args -> b)) = do
     maybe (fail err) return $ do
         guard $ anull args
-        cast $ value mempty
+        cast $ value sline mempty
   where
     err = T.unpack $ T.concat [ T.pack "expected ", textExprType @a Proxy, T.pack ", but variable '", textVarName name, T.pack "' has type ",
             if anull args then textExprType @b Proxy else "function type" ]
 
-textSomeVarValue :: SomeVarValue -> Text
-textSomeVarValue (SomeVarValue args value)
-    | anull args = textExprValue $ value mempty
+textSomeVarValue :: SourceLine -> SomeVarValue -> Text
+textSomeVarValue sline (SomeVarValue args value)
+    | anull args = textExprValue $ value sline mempty
     | otherwise  =  "<function>"
 
 someVarValueType :: SomeVarValue -> SomeExprType
-someVarValueType (SomeVarValue args (_ :: args -> a))
+someVarValueType (SomeVarValue args (_ :: SourceLine -> args -> a))
     | anull args = ExprTypePrim (Proxy @a)
     | otherwise  = ExprTypeFunction args (Proxy @a)
 
@@ -203,9 +203,9 @@ data ExprEnumerator a = ExprEnumerator (a -> a -> [a]) (a -> a -> a -> [a])
 
 
 data Expr a where
-    Variable :: ExprType a => VarName -> Expr a
-    DynVariable :: TypeVar -> VarName -> Expr DynamicType
-    FunVariable :: ExprType a => FunctionArguments SomeExprType -> VarName -> Expr (FunctionType a)
+    Variable :: ExprType a => SourceLine -> VarName -> Expr a
+    DynVariable :: TypeVar -> SourceLine -> VarName -> Expr DynamicType
+    FunVariable :: ExprType a => FunctionArguments SomeExprType -> SourceLine -> VarName -> Expr (FunctionType a)
     ArgsApp :: FunctionArguments SomeExpr -> Expr (FunctionType a) -> Expr (FunctionType a)
     FunctionEval :: Expr (FunctionType a) -> Expr a
     Pure :: a -> Expr a
@@ -226,9 +226,9 @@ instance Applicative Expr where
     (<*>) = App AnnNone
 
 eval :: MonadEval m => Expr a -> m a
-eval (Variable name) = fromSomeVarValue name =<< lookupVar name
-eval (DynVariable _ _) = fail "ambiguous type"
-eval (FunVariable _ name) = funFromSomeVarValue name =<< lookupVar name
+eval (Variable sline name) = fromSomeVarValue sline name =<< lookupVar name
+eval (DynVariable _ _ _) = fail "ambiguous type"
+eval (FunVariable _ sline name) = funFromSomeVarValue sline name =<< lookupVar name
 eval (ArgsApp args efun) = do
     FunctionType fun <- eval efun
     return $ FunctionType $ \args' -> fun (args <> args')
@@ -247,15 +247,15 @@ eval (RootNetwork) = rootNetwork
 eval (Undefined err) = fail err
 
 evalSome :: MonadEval m => SomeExpr -> m SomeVarValue
-evalSome (SomeExpr expr) = SomeVarValue mempty . const <$> eval expr
+evalSome (SomeExpr expr) = SomeVarValue mempty . const . const <$> eval expr
 
 gatherVars :: forall a m. MonadEval m => Expr a -> m [((VarName, [Text]), SomeVarValue)]
 gatherVars = fmap (uniqOn fst . sortOn fst) . helper
   where
     helper :: forall b. Expr b -> m [((VarName, [Text]), SomeVarValue)]
-    helper (Variable var) = (:[]) . ((var, []),) <$> lookupVar var
-    helper (DynVariable _ var) = (:[]) . ((var, []),) <$> lookupVar var
-    helper (FunVariable _ var) = (:[]) . ((var, []),) <$> lookupVar var
+    helper (Variable _ var) = (:[]) . ((var, []),) <$> lookupVar var
+    helper (DynVariable _ _ var) = (:[]) . ((var, []),) <$> lookupVar var
+    helper (FunVariable _ _ var) = (:[]) . ((var, []),) <$> lookupVar var
     helper (ArgsApp (FunctionArguments args) fun) = do
         v <- helper fun
         vs <- mapM (\(SomeExpr e) -> helper e) $ M.elems args
@@ -264,7 +264,7 @@ gatherVars = fmap (uniqOn fst . sortOn fst) . helper
     helper (Pure _) = return []
     helper e@(App (AnnRecord sel) _ x)
         | Just (var, sels) <- gatherSelectors x
-        = do val <- SomeVarValue mempty . const <$> eval e
+        = do val <- SomeVarValue mempty . const . const <$> eval e
              return [((var, sels ++ [sel]), val)]
         | otherwise = helper x
     helper (App _ f x) = (++) <$> helper f <*> helper x
@@ -275,7 +275,7 @@ gatherVars = fmap (uniqOn fst . sortOn fst) . helper
 
     gatherSelectors :: forall b. Expr b -> Maybe (VarName, [Text])
     gatherSelectors = \case
-        Variable var -> Just (var, [])
+        Variable _ var -> Just (var, [])
         App (AnnRecord sel) _ x -> do
             (var, sels) <- gatherSelectors x
             return (var, sels ++ [sel])
@@ -292,17 +292,17 @@ anull :: FunctionArguments a -> Bool
 anull (FunctionArguments args) = M.null args
 
 exprArgs :: Expr (FunctionType a) -> FunctionArguments SomeExprType
-exprArgs (FunVariable args _) = args
+exprArgs (FunVariable args _ _) = args
 exprArgs (ArgsApp (FunctionArguments applied) expr) =
     let FunctionArguments args = exprArgs expr
      in FunctionArguments (args `M.difference` applied)
 exprArgs _ = error "exprArgs on unexpected type"
 
-funFromSomeVarValue :: forall a m. (ExprType a, MonadFail m) => VarName -> SomeVarValue -> m (FunctionType a)
-funFromSomeVarValue name (SomeVarValue args (value :: args -> b)) = do
+funFromSomeVarValue :: forall a m. (ExprType a, MonadFail m) => SourceLine -> VarName -> SomeVarValue -> m (FunctionType a)
+funFromSomeVarValue sline name (SomeVarValue args (value :: SourceLine -> args -> b)) = do
     maybe (fail err) return $ do
         guard $ not $ anull args
-        FunctionType <$> cast value
+        FunctionType <$> cast (value sline)
   where
     err = T.unpack $ T.concat [ T.pack "expected function returning ", textExprType @a Proxy, T.pack ", but variable '", textVarName name, T.pack "' has ",
             (if anull args then "type" else "function type returting ") <> textExprType @b Proxy ]
diff --git a/src/Test/Builtins.hs b/src/Test/Builtins.hs
index 2ab38aa..b768bb9 100644
--- a/src/Test/Builtins.hs
+++ b/src/Test/Builtins.hs
@@ -2,12 +2,26 @@ module Test.Builtins (
     builtins,
 ) where
 
+import Data.Map qualified as M
+import Data.Typeable
+
 import Test
 
 builtins :: [ ( VarName, SomeVarValue ) ]
 builtins =
-    [ ( VarName "wait", builtinWait )
+    [ ( VarName "guard", builtinGuard )
+    , ( VarName "wait", builtinWait )
     ]
 
+getArg :: Typeable a => FunctionArguments SomeExpr -> Maybe ArgumentKeyword -> a
+getArg (FunctionArguments args) kw =
+    case M.lookup kw args of
+        Just (SomeExpr expr) | Just expr' <- cast expr -> expr'
+        _ -> error "parameter mismatch"
+
+builtinGuard :: SomeVarValue
+builtinGuard = SomeVarValue (FunctionArguments $ M.singleton Nothing (ExprTypePrim (Proxy @Bool))) $
+    \sline args -> TestBlock [ Guard sline (getArg args Nothing) ]
+
 builtinWait :: SomeVarValue
-builtinWait = SomeVarValue mempty $ const $ TestBlock [ Wait ]
+builtinWait = SomeVarValue mempty $ const . const $ TestBlock [ Wait ]