validate table.init and use mutable vector as a table storage

src/Language/Wasm/Interpreter.hs

@@ -34,6 +34,8 @@ import Data.Maybe (fromMaybe, isNothing)
 
 import Data.Vector (Vector, (!), (!?), (//))
 import qualified Data.Vector as Vector
+import Data.Vector.Mutable (IOVector)
+import qualified Data.Vector.Mutable as MVector
 import qualified Data.Primitive.ByteArray as ByteArray
 import qualified Data.Primitive.Types as Primitive
 import qualified Control.Monad.Primitive as Primitive
@@ -168,9 +170,11 @@ data Label = Label ResultType deriving (Show, Eq)
 
 type Address = Int
 
+type TableStore = IOVector (Maybe Address)
+
 data TableInstance = TableInstance {
     lim :: Limit,
-    elements :: Vector (Maybe Address)
+    items :: TableStore
 }
 
 type MemoryStore = ByteArray.MutableByteArray (Primitive.PrimState IO)
@@ -298,8 +302,8 @@ makeHostModule st items = do
         makeHostTables :: (Store, ModuleInstance) -> IO (Store, ModuleInstance)
         makeHostTables (st, inst) = do
             let tableLen = Vector.length $ tableInstances st
-            let (names, tables) = unzip [(name, Table (TableType lim FuncRef)) | (name, (HostTable lim)) <- items]
+            let (names, tables) = unzip [(name, Table (TableType lim FuncRef)) | (name, HostTable lim) <- items]
-            let instances = allocTables tables
+            instances <- allocTables tables
             let exps = Vector.fromList $ zipWith (\name i -> ExportInstance name (ExternTable i)) names [tableLen..]
             let inst' = inst {
                     tableaddrs = Vector.fromList [tableLen..tableLen + length instances - 1],
@@ -465,15 +469,13 @@ allocAndInitGlobals inst store globs = Vector.fromList <$> mapM allocGlob globs
             val <- runIniter initer
             GIMut vt <$> newIORef val
 
-allocTables :: [Table] -> Vector TableInstance
+allocTables :: [Table] -> IO (Vector TableInstance)
-allocTables = Vector.fromList . map allocTable
+allocTables = fmap Vector.fromList . mapM allocTable
     where
-        allocTable :: Table -> TableInstance
+        allocTable :: Table -> IO TableInstance
         allocTable (Table (TableType lim@(Limit from to) _)) =
-            TableInstance {
+            let elements = MVector.replicate (fromIntegral from) Nothing in
-                lim,
+            TableInstance lim <$> elements
-                elements = Vector.fromList $ replicate (fromIntegral from) Nothing
-            }
 
 defaultBudget :: Natural
 defaultBudget = 300
@@ -541,15 +543,15 @@ initialize inst Module {elems, datas, start} = do
             let idx = tableaddrs inst ! fromIntegral tableIndex
             let last = from + length funcs
             let TableInstance lim elems = tableInstances st ! idx
-            let len = Vector.length elems
+            let len = MVector.length elems
-            Monad.when (last > len) $ throwError "elements segment does not fit"
+            Monad.when (last > len) $ throwError "out of bounds table access"
             return (idx, from, funcs)
 
         initElem :: (Address, Int, [Maybe Address]) -> Initialize ()
-        initElem (idx, from, funcs) = State.modify $ \st ->
+        initElem (idx, from, funcs) = do
-            let TableInstance lim elems = tableInstances st ! idx in
+            Store {tableInstances} <- State.get
-            let table = TableInstance lim (elems // zip [from..] funcs) in
+            let elems = items $ tableInstances ! idx
-            st { tableInstances = tableInstances st Vector.// [(idx, table)] }
+            Monad.forM_ (zip [from..] funcs) $ uncurry $ MVector.unsafeWrite elems
 
         checkData :: DataSegment -> Initialize (Int, MemoryStore, LBS.ByteString)
         checkData DataSegment {memIndex, offset, chunk} = do
@@ -574,7 +576,7 @@ instantiate st imps mod = flip State.runStateT st $ runExceptT $ do
     inst <- calcInstance st imps m
     let functions = funcInstances st <> allocFunctions inst (Struct.functions m)
     globals <- liftIO $ (globalInstances st <>) <$> allocAndInitGlobals inst st (Struct.globals m)
-    let tables = tableInstances st <> allocTables (Struct.tables m)
+    tables <- (tableInstances st <>) <$> liftIO (allocTables (Struct.tables m))
     mems <- liftIO $ (memInstances st <>) <$> allocMems (Struct.mems m)
     elems <- liftIO $ (elemInstances st <>) <$> allocElems inst st (Struct.elems m)
     let datas = dataInstances st <> allocDatas inst st (Struct.datas m)
@@ -750,23 +752,28 @@ eval budget store FunctionInstance { funcType, moduleInstance, code = Function {
                 Nothing -> return Trap
         step ctx@EvalCtx{ stack = (VI32 v): rest } (CallIndirect typeIdx) = do
             let funcType = funcTypes moduleInstance ! fromIntegral typeIdx
-            let TableInstance { elements } = tableInstances store ! (tableaddrs moduleInstance ! 0)
+            let TableInstance { items } = tableInstances store ! (tableaddrs moduleInstance ! 0)
-            let checks = do
+            let pos = fromIntegral v
-                    addr <- Monad.join $ elements !? fromIntegral v
+            if pos >= MVector.length items
-                    let funcInst = funcInstances store ! addr
+            then return Trap
-                    let targetType = Language.Wasm.Interpreter.funcType funcInst
+            else do
-                    Monad.guard $ targetType == funcType
+                maybeAddr <- liftIO $ MVector.read items pos
-                    let args = params targetType
+                let checks = do
-                    Monad.guard $ length args <= length rest
+                        addr <- maybeAddr
-                    params <- sequence $ zipWith checkValType args $ reverse $ take (length args) rest
+                        let funcInst = funcInstances store ! addr
-                    return (funcInst, params)
+                        let targetType = Language.Wasm.Interpreter.funcType funcInst
-            case checks of
+                        Monad.guard $ targetType == funcType
-                Just (funcInst, params) -> do
+                        let args = params targetType
-                    res <- eval (budget - 1) store funcInst params
+                        Monad.guard $ length args <= length rest
-                    case res of
+                        params <- sequence $ zipWith checkValType args $ reverse $ take (length args) rest
-                        Just res -> return $ Done ctx { stack = reverse res ++ (drop (length params) rest) }
+                        return (funcInst, params)
-                        Nothing -> return Trap
+                case checks of
-                Nothing -> return Trap
+                    Just (funcInst, params) -> do
+                        res <- eval (budget - 1) store funcInst params
+                        case res of
+                            Just res -> return $ Done ctx { stack = reverse res ++ (drop (length params) rest) }
+                            Nothing -> return Trap
+                    Nothing -> return Trap
         step ctx@EvalCtx{ stack = st } (RefNull FuncRef) =
             return $ Done ctx { stack = RF Nothing : st }
         step ctx@EvalCtx{ stack = st } (RefNull ExternRef) =
@@ -880,6 +887,11 @@ eval budget store FunctionInstance { funcType, moduleInstance, code = Function {
                     else return $ -1
                 )
             return $ Done ctx { stack = VI32 (asWord32 $ fromIntegral result) : rest }
+        -- step ctx@EvalCtx{ stack = (VI32 n:rest) } (TableInit tableIdx elemIdx) = do
+        --     let tableAddr = tableaddrs moduleInstance ! fromIntegral tableIdx
+        --     let TableInstance { items } = tableInstances store ! tableAddr
+
+        --     return $ Done ctx { stack = rest }
         step ctx (I32Const v) = return $ Done ctx { stack = VI32 v : stack ctx }
         step ctx (I64Const v) = return $ Done ctx { stack = VI64 v : stack ctx }
         step ctx (F32Const v) = return $ Done ctx { stack = VF32 v : stack ctx }

src/Language/Wasm/Script.hs

@@ -260,9 +260,10 @@ runScript onAssertFail script = do
                     st <- fst <$> State.get
                     (res, store') <- liftIO $ Interpreter.instantiate (store st) (buildImports st) m
                     case res of
+                        Left failureString -> return ()
                         Left "Start function terminated with trap" ->
                             State.modify $ \(st, pos) -> (st { store = store' }, pos)
-                        _ -> printFailedAssert ("Module linking should fail with trap during execution of a start function") assert
+                        r -> printFailedAssert "Module linking should fail with trap during execution of a start function" assert
                 Left reason -> error $ "Module linking failed due to invalid module with reason: " ++ show reason
         runAssert assert@(AssertExhaustion action failureString) = do
             result <- runActionInAssert action

src/Language/Wasm/Validate.hs

@@ -38,10 +38,12 @@ data ValidationError =
     | MemoryIndexOutOfRange Natural
     | LocalIndexOutOfRange Natural
     | GlobalIndexOutOfRange Natural
+    | ElemIndexOutOfRange Natural
     | LabelIndexOutOfRange
     | TypeIndexOutOfRange
     | ResultTypeDoesntMatch
     | TypeMismatch { actual :: Arrow, expected :: Arrow }
+    | RefTypeMismatch ElemType ElemType
     | InvalidResultArity
     | InvalidConstantExpr
     | InvalidStartFunctionType
@@ -127,6 +129,7 @@ data Ctx = Ctx {
     types :: [FuncType],
     funcs :: [FuncType],
     tables :: [TableType],
+    elems :: [ElemType],
     mems :: [Limit],
     globals :: [GlobalType],
     locals :: [ValueType],
@@ -362,8 +365,16 @@ getInstrType CurrentMemory = do
     Ctx { mems } <- ask 
     if length mems < 1 then throwError (MemoryIndexOutOfRange 0) else return $ empty ==> I32
 getInstrType GrowMemory = do
-    Ctx { mems } <- ask 
+    Ctx { mems } <- ask
     if length mems < 1 then throwError (MemoryIndexOutOfRange 0) else return $ I32 ==> I32
+getInstrType (TableInit tableIdx elemIdx) = do
+    Ctx { tables, elems } <- ask
+    when (length tables <= fromIntegral tableIdx) $ throwError (TableIndexOutOfRange tableIdx)
+    when (length elems <= fromIntegral elemIdx) $ throwError (ElemIndexOutOfRange elemIdx)
+    let TableType _ tableType = tables !! fromIntegral tableIdx
+    let elemType = elems !! fromIntegral elemIdx
+    when (elemType /= tableType) $ throwError (RefTypeMismatch tableType elemType)
+    return $ [I32, I32, I32] ==> empty
 getInstrType (I32Const _) = return $ empty ==> I32
 getInstrType (I64Const _) = return $ empty ==> I64
 getInstrType (F32Const _) = return $ empty ==> F32
@@ -476,7 +487,7 @@ getFuncTypes Module {types, functions, imports} =
         getFuncType _ = Nothing
 
 ctxFromModule :: [ValueType] -> [[ValueType]] -> [ValueType] -> Module -> Ctx
-ctxFromModule locals labels returns m@Module {types, tables, mems, globals, imports} =
+ctxFromModule locals labels returns m@Module {types, tables, mems, globals, imports, elems} =
     let tableImports = catMaybes $ map getTableType imports in
     let memsImports = catMaybes $ map getMemType imports in
     let globalImports = catMaybes $ map getGlobalType imports in
@@ -484,6 +495,7 @@ ctxFromModule locals labels returns m@Module {types, tables, mems, globals, impo
         types,
         funcs = getFuncTypes m,
         tables = tableImports ++ map (\(Table t) -> t) tables,
+        elems = map elemType elems,
         mems = memsImports ++ map (\(Memory l) -> l) mems,
         globals = globalImports ++ map (\(Global g _) -> g) globals,
         locals,