From 5356e0b4351b2c403c2868421f25013091848bed Mon Sep 17 00:00:00 2001
From: Ilya Rezvov <ilya.rezvov@revjet.com>
Date: Sat, 21 Apr 2018 11:24:05 -0700
Subject: [PATCH] pass all tests

---
 src/Language/Wasm/Interpreter.hs | 18 +++++++++++-------
 src/Language/Wasm/Script.hs      |  7 ++++++-
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/src/Language/Wasm/Interpreter.hs b/src/Language/Wasm/Interpreter.hs
index 5e5de8d..4038468 100644
--- a/src/Language/Wasm/Interpreter.hs
+++ b/src/Language/Wasm/Interpreter.hs
@@ -465,6 +465,9 @@ allocTables tables = Vector.fromList $ map allocTable tables
                 elements = Vector.fromList $ replicate (fromIntegral from) Nothing
             }
 
+defaultBudget :: Natural
+defaultBudget = 300
+
 pageSize :: Int
 pageSize = 64 * 1024
 
@@ -491,7 +494,7 @@ initialize inst Module {elems, datas, start} store = do
             case start of
                 Just (StartFunction idx) -> do
                     let funInst = funcInstances store ! (funcaddrs inst ! fromIntegral idx)
-                    mainRes <- eval st funInst []
+                    mainRes <- eval defaultBudget st funInst []
                     case mainRes of
                         Just [] -> return $ Right st
                         _ -> return $ Left "Start function terminated with trap"
@@ -586,8 +589,9 @@ data EvalResult =
     | ReturnFn [Value]
     deriving (Show, Eq)
 
-eval :: Store -> FunctionInstance -> [Value] -> IO (Maybe [Value])
-eval store FunctionInstance { funcType, moduleInstance, code = Function { localTypes, body} } args = do
+eval :: Natural -> Store -> FunctionInstance -> [Value] -> IO (Maybe [Value])
+eval 0 _ _ _ = return Nothing
+eval budget store FunctionInstance { funcType, moduleInstance, code = Function { localTypes, body} } args = do
     case sequence $ zipWith checkValType (params funcType) args of
         Just checkedArgs -> do
             let initialContext = EvalCtx {
@@ -676,7 +680,7 @@ eval store FunctionInstance { funcType, moduleInstance, code = Function { localT
             let args = params ft
             case sequence $ zipWith checkValType args $ reverse $ take (length args) $ stack ctx of
                 Just params -> do
-                    res <- eval store funInst params
+                    res <- eval (budget - 1) store funInst params
                     case res of
                         Just res -> return $ Done ctx { stack = reverse res ++ (drop (length args) $ stack ctx) }
                         Nothing -> return Trap
@@ -696,7 +700,7 @@ eval store FunctionInstance { funcType, moduleInstance, code = Function { localT
                         then return Trap
                         else case sequence $ zipWith checkValType args $ reverse $ take (length args) rest of
                             Just params -> do
-                                res <- eval store funInst params
+                                res <- eval (budget - 1) store funInst params
                                 case res of
                                     Just res -> return $ Done ctx { stack = reverse res ++ (drop (length args) rest) }
                                     Nothing -> return Trap
@@ -1301,10 +1305,10 @@ eval store FunctionInstance { funcType, moduleInstance, code = Function { localT
         step ctx@EvalCtx{ stack = (VI64 v:rest) } (FReinterpretI BS64) =
             return $ Done ctx { stack = VF64 (wordToDouble v) : rest }
         step EvalCtx{ stack } instr = error $ "Error during evaluation of instruction: " ++ show instr ++ ". Stack " ++ show stack
-eval _ HostInstance { funcType, hostCode } args = Just <$> hostCode args
+eval _ _ HostInstance { funcType, hostCode } args = Just <$> hostCode args
 
 invoke :: Store -> Address -> [Value] -> IO (Maybe [Value])
-invoke st funcIdx = eval st $ funcInstances st ! funcIdx
+invoke st funcIdx = eval defaultBudget st $ funcInstances st ! funcIdx
 
 invokeExport :: Store -> ModuleInstance -> TL.Text -> [Value] -> IO (Maybe [Value])
 invokeExport st ModuleInstance { exports } name args =
diff --git a/src/Language/Wasm/Script.hs b/src/Language/Wasm/Script.hs
index 813765c..0f49494 100644
--- a/src/Language/Wasm/Script.hs
+++ b/src/Language/Wasm/Script.hs
@@ -212,6 +212,7 @@ runScript onAssertFail script = do
             case Binary.decodeModuleLazy binaryRep of
                 Right _ -> onAssertFail ("Module decoding should fail with failure string " ++ show failureString) assert
                 Left _ -> return ()
+        runAssert st assert@(AssertMalformed (RawModDef _ _) failureString) = return ()
         runAssert st assert@(AssertUnlinkable moduleDef failureString) =
             let (_, m) = buildModule moduleDef in
             case Validate.validate m of
@@ -235,7 +236,11 @@ runScript onAssertFail script = do
                         Left "Start function terminated with trap" -> return ()
                         _ -> onAssertFail ("Module linking should fail with trap during execution of a start function") assert
                 reason -> error $ "Module linking failed dut to invalid module with reason: " ++ show reason
-        runAssert _ _ = return ()
+        runAssert st assert@(AssertExhaustion action failureString) = do
+            result <- runAction st action
+            if isNothing result
+            then return ()
+            else onAssertFail ("Expected exhaustion, but action returned " ++ show (fromJust result)) assert
 
         runCommand :: ScriptState -> Command -> IO ScriptState
         runCommand st (ModuleDef moduleDef) =