From dac931b055372b4cf70675c81896bc82074d58ae Mon Sep 17 00:00:00 2001
From: Ilya Rezvov <ilya.rezvov@revjet.com>
Date: Mon, 9 Apr 2018 15:08:48 -0700
Subject: [PATCH] implement assert_return assertion

---
 src/Language/Wasm/Script.hs | 14 +++++++++++---
 tests/Test.hs               |  4 ++--
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/Language/Wasm/Script.hs b/src/Language/Wasm/Script.hs
index e1caa7e..f2e7639 100644
--- a/src/Language/Wasm/Script.hs
+++ b/src/Language/Wasm/Script.hs
@@ -15,8 +15,7 @@ import Language.Wasm.Parser (
         ModuleDef(..),
         Command(..),
         Action(..),
-        Assertion(..),
-        Meta(..)
+        Assertion(..)
     )
 
 import qualified Language.Wasm.Interpreter as Interpreter
@@ -26,7 +25,7 @@ import qualified Language.Wasm.Parser as Parser
 import qualified Language.Wasm.Lexer as Lexer
 import qualified Language.Wasm.Binary as Binary
 
-type OnAssertFail = Assertion -> IO ()
+type OnAssertFail = String -> Assertion -> IO ()
 
 data ScriptState = ScriptState {
     store :: Interpreter.Store,
@@ -118,6 +117,14 @@ runScript onAssertFail script = do
             case getModule st ident of
                 Just m -> Interpreter.getGlobalValueByName (store st) m name >>= return . (: [])
                 Nothing -> error $ "Cannot invoke function on module with identifier '" ++ show ident  ++ "'. No such module"
+        
+        runAssert :: ScriptState -> Assertion -> IO ()
+        runAssert st assert@(AssertReturn action expected) = do
+            result <- runAction st action
+            if result == map asArg expected
+            then return ()
+            else onAssertFail ("Expected " ++ show (map asArg expected) ++ ", but action returned " ++ show result) assert
+        runAssert _ _ = return ()
 
         runCommand :: ScriptState -> Command -> IO ScriptState
         runCommand st (ModuleDef (RawModDef ident m)) = addModule ident m st
@@ -129,4 +136,5 @@ runScript onAssertFail script = do
             addModule ident m st
         runCommand st (Register name i) = return $ addToRegistery name i st
         runCommand st (Action action) = runAction st action >> return st
+        runCommand st (Assertion assertion) = runAssert st assertion >> return st
         runCommand st _ = return st
diff --git a/tests/Test.hs b/tests/Test.hs
index de3129b..2c8073b 100644
--- a/tests/Test.hs
+++ b/tests/Test.hs
@@ -34,10 +34,10 @@ compile file = do
 main :: IO ()
 main = do
   files <- Directory.listDirectory "tests/samples"
-  -- let files = ["labels.wast"]
+  let files = ["linking.wast"]
   scriptTestCases <- (`mapM` files) $ \file -> do
     content <- LBS.readFile $ "tests/samples/" ++ file
     let Right script = Parser.parseScript <$> Lexer.scanner content
     return $ testCase file $ do
-      Script.runScript (assertFailure . ("Failed assert: " ++) . show) script
+      Script.runScript (\msg assert -> assertFailure ("Failed assert: " ++ msg ++ ". Assert " ++ show assert)) script
   defaultMain $ testGroup "Wasm Core Test Suit" scriptTestCases