check expression types

src/Language/Wasm/Validate.hs

@@ -12,14 +12,24 @@ import Language.Wasm.Structure
 import qualified Data.Set as Set
 import Data.List (foldl')
 import qualified Data.Text.Lazy as TL
-import Data.Maybe (fromMaybe)
+import Data.Maybe (fromMaybe, maybeToList)
 import Data.Monoid ((<>))
+import Numeric.Natural (Natural)
+
+import Control.Monad.State.Lazy (StateT, evalStateT, get, put)
+import Control.Monad.Reader (ReaderT, runReaderT, withReaderT, ask)
+import Control.Monad.Except (Except, runExcept, throwError)
 
 data ValidationResult =
     DuplicatedExportNames [String]
     | InvalidTableType
     | MoreThanOneMemory
     | MoreThanOneTable
+    | IndexOutOfRange
+    | ResultTypeDoesntMatch
+    | NoTableInModule
+    | NoMemoryInModule
+    | TypeMismatch
     | Valid
     deriving (Show, Eq)
 
@@ -37,12 +47,8 @@ type Validator = Module -> ValidationResult
 
 data VType =
     Val ValueType
-    | Var
+    | Var Int
-    | LabelRef LabelIndex
-    | LocalRef LocalIndex
-    | GlobalRef GlobalIndex
     | Any
-    | Result
     deriving (Show, Eq)
 
 type End = [VType]
@@ -70,93 +76,321 @@ data Arrow = Arrow End End deriving (Show, Eq)
 (==>) :: (ToEnd a, ToEnd b) => a -> b -> Arrow
 (==>) a b = Arrow (toEnd a) (toEnd b)
 
-getInstrType :: Instruction -> Arrow
+asArrow :: FuncType -> Arrow
-getInstrType Unreachable = Any ==> Any
+asArrow (FuncType params results) = Arrow (map Val params) (map Val results)
-getInstrType Nop = empty ==> empty
+
-getInstrType Block { result } = empty ==> result
+data Ctx = Ctx {
-getInstrType Loop { result } = empty ==> result
+    types :: [FuncType],
-getInstrType If { result } = I32 ==> result
+    funcs :: [FuncType],
-getInstrType (Br lbl) = [Any, LabelRef lbl] ==> Any
+    tables :: [TableType],
-getInstrType (BrIf lbl) = [LabelRef lbl, Val I32] ==> LabelRef lbl
+    mems :: [Limit],
-getInstrType (BrTable _ lbl) = [Any, LabelRef lbl, Val I32] ==> Any
+    globals :: [GlobalType],
-getInstrType Return = [Any, Result] ==> Any
+    locals :: [ValueType],
-getInstrType (Call _) = Any ==> Any
+    labels :: [Maybe ValueType],
-getInstrType (CallIndirect _) = [Any, Val I32] ==> Any
+    returns :: Maybe ValueType
-getInstrType Drop = Var ==> empty
+} deriving (Show, Eq)
-getInstrType Select = [Var, Var, Val I32] ==> Var
+
-getInstrType (GetLocal local) = empty ==> LocalRef local
+type Checker = ReaderT Ctx (StateT Int (Except ValidationResult))
-getInstrType (SetLocal local) = LocalRef local ==> empty
+
-getInstrType (TeeLocal local) = LocalRef local ==> LocalRef local
+freshVar :: Checker VType
-getInstrType (GetGlobal global) = empty ==> GlobalRef global
+freshVar = do
-getInstrType (SetGlobal global) = GlobalRef global ==> empty
+    i <- get
-getInstrType (I32Load _) = I32 ==> I32
+    put (i + 1)
-getInstrType (I64Load _) = I32 ==> I64
+    return $ Var i
-getInstrType (F32Load _) = I32 ==> F32
+
-getInstrType (F64Load _) = I32 ==> F64
+runChecker :: Ctx -> Checker a -> Either ValidationResult a
-getInstrType (I32Load8S _) = I32 ==> I32
+runChecker ctx = runExcept . flip evalStateT 0 . flip runReaderT ctx
-getInstrType (I32Load8U _) = I32 ==> I32
+
-getInstrType (I32Load16S _) = I32 ==> I32
+(!?) :: [a] -> Natural -> Maybe a
-getInstrType (I32Load16U _) = I32 ==> I32
+(!?) (x:_) 0 = Just x
-getInstrType (I64Load8S _) = I32 ==> I64
+(!?) (_:rest) n = rest !? (n - 1)
-getInstrType (I64Load8U _) = I32 ==> I64
+(!?) [] _ = Nothing
-getInstrType (I64Load16S _) = I32 ==> I64
+
-getInstrType (I64Load16U _) = I32 ==> I64
+safeHead :: [a] -> Maybe a
-getInstrType (I64Load32S _) = I32 ==> I64
+safeHead (x: _) = Just x
-getInstrType (I64Load32U _) = I32 ==> I64
+safeHead [] = Nothing
-getInstrType (I32Store _) = [I32, I32] ==> empty
+
-getInstrType (I64Store _) = [I32, I64] ==> empty
+maybeToEither :: ValidationResult -> Maybe a -> Checker a
-getInstrType (F32Store _) = [I32, F32] ==> empty
+maybeToEither _ (Just a) = return a
-getInstrType (F64Store _) = [I32, F64] ==> empty
+maybeToEither l Nothing = throwError l
-getInstrType (I32Store8 _) = [I32, I32] ==> empty
+
-getInstrType (I32Store16 _) = [I32, I32] ==> empty
+asType :: GlobalType -> VType
-getInstrType (I64Store8 _) = [I32, I64] ==> empty
+asType (Const v) = Val v
-getInstrType (I64Store16 _) = [I32, I64] ==> empty
+asType (Mut v) = Val v
-getInstrType (I64Store32 _) = [I32, I64] ==> empty
+
-getInstrType CurrentMemory = empty ==> I32
+getLabel :: LabelIndex -> Checker (Maybe ValueType)
-getInstrType GrowMemory = I32 ==> I32
+getLabel lbl = do
-getInstrType (I32Const _) = empty ==> I32
+    Ctx { labels } <- ask
-getInstrType (I64Const _) = empty ==> I64
+    case labels !? lbl of
-getInstrType (F32Const _) = empty ==> F32
+        Nothing -> throwError IndexOutOfRange
-getInstrType (F64Const _) = empty ==> F64
+        Just v -> return v
-getInstrType (IUnOp BS32 _) = I32 ==> I32
+
-getInstrType (IUnOp BS64 _) = I64 ==> I64
+withLabel :: [ValueType] -> Checker a -> Checker a
-getInstrType (IBinOp BS32 _) = [I32, I32] ==> I32
+withLabel result = withReaderT (\ctx -> ctx { labels = safeHead result : labels ctx })
-getInstrType (IBinOp BS64 _) = [I64, I64] ==> I64
+
-getInstrType I32Eqz = I32 ==> I32
+getInstrType :: Instruction -> Checker Arrow
-getInstrType I64Eqz = I64 ==> I32
+getInstrType Unreachable = return $ Any ==> Any
-getInstrType (IRelOp BS32 _) = [I32, I32] ==> I32
+getInstrType Nop = return $ empty ==> empty
-getInstrType (IRelOp BS64 _) = [I64, I64] ==> I32
+getInstrType Block { result, body } = do
-getInstrType (FUnOp BS32 _) = F32 ==> F32
+    let blockType = empty ==> result
-getInstrType (FUnOp BS64 _) = F64 ==> F64
+    t <- withLabel result $ getExpressionType body
-getInstrType (FBinOp BS32 _) = [F32, F32] ==> F32
+    if t == blockType
-getInstrType (FBinOp BS64 _) = [F64, F64] ==> F64
+    then return $ empty ==> result
-getInstrType (FRelOp BS32 _) = [F32, F32] ==> I32
+    else throwError TypeMismatch
-getInstrType (FRelOp BS64 _) = [F64, F64] ==> I32
+getInstrType Loop { result, body } = do
-getInstrType I32WrapI64 = I64 ==> I32
+    let blockType = empty ==> result
-getInstrType (ITruncFU BS32 BS32) = F32 ==> I32
+    t <- withLabel result $ getExpressionType body
-getInstrType (ITruncFU BS32 BS64) = F64 ==> I32
+    if t == blockType
-getInstrType (ITruncFU BS64 BS32) = F32 ==> I64
+    then return $ empty ==> result
-getInstrType (ITruncFU BS64 BS64) = F64 ==> I64
+    else throwError TypeMismatch
-getInstrType (ITruncFS BS32 BS32) = F32 ==> I32
+getInstrType If { result, true, false } = do
-getInstrType (ITruncFS BS32 BS64) = F64 ==> I32
+    let blockType = empty ==> result
-getInstrType (ITruncFS BS64 BS32) = F32 ==> I64
+    l <- withLabel result $ getExpressionType true
-getInstrType (ITruncFS BS64 BS64) = F64 ==> I64
+    r <- withLabel result $ getExpressionType false
-getInstrType I64ExtendSI32 = I32 ==> I64
+    if l == blockType && r == blockType
-getInstrType I64ExtendUI32 = I32 ==> I64
+    then return $ I32 ==> result
-getInstrType (FConvertIU BS32 BS32) = I32 ==> F32
+    else throwError TypeMismatch
-getInstrType (FConvertIU BS32 BS64) = I64 ==> F32
+getInstrType (Br lbl) = do
-getInstrType (FConvertIU BS64 BS32) = I32 ==> F64
+    r <- map Val . maybeToList <$> getLabel lbl
-getInstrType (FConvertIU BS64 BS64) = I64 ==> F64
+    return $ (Any : r) ==> Any
-getInstrType (FConvertIS BS32 BS32) = I32 ==> F32
+getInstrType (BrIf lbl) = do
-getInstrType (FConvertIS BS32 BS64) = I64 ==> F32
+    r <- map Val . maybeToList <$> getLabel lbl
-getInstrType (FConvertIS BS64 BS32) = I32 ==> F64
+    return $ (r ++ [Val I32]) ==> r
-getInstrType (FConvertIS BS64 BS64) = I64 ==> F64
+getInstrType (BrTable lbls lbl) = do
-getInstrType F32DemoteF64 = F64 ==> F32
+    r <- getLabel lbl
-getInstrType F64PromoteF32 = F32 ==> F64
+    rs <- mapM getLabel lbls
-getInstrType (IReinterpretF BS32) = F32 ==> I32
+    if all (== r) rs
-getInstrType (IReinterpretF BS64) = F64 ==> I64
+    then return $ ([Any] ++ (map Val $ maybeToList r) ++ [Val I32]) ==> Any
-getInstrType (FReinterpretI BS32) = I32 ==> F32
+    else throwError ResultTypeDoesntMatch
-getInstrType (FReinterpretI BS64) = I64 ==> F64
+getInstrType Return = do
+    Ctx { returns } <- ask
+    return $ (Any : (map Val $ maybeToList returns)) ==> Any
+getInstrType (Call fun) = do
+    Ctx { funcs } <- ask
+    maybeToEither IndexOutOfRange $ asArrow <$> funcs !? fun
+getInstrType (CallIndirect sign) = do
+    Ctx { types, tables } <- ask
+    if length tables < 1
+    then throwError NoTableInModule
+    else do
+        Arrow from to <- maybeToEither IndexOutOfRange $ asArrow <$> types !? sign
+        return $ (from ++ [Val I32]) ==> to
+getInstrType Drop = do
+    var <- freshVar
+    return $ var ==> empty
+getInstrType Select = do
+    var <- freshVar
+    return $ [var, var, Val I32] ==> var
+getInstrType (GetLocal local) = do
+    Ctx { locals }  <- ask
+    t <- maybeToEither IndexOutOfRange $ locals !? local
+    return $ empty ==> Val t
+getInstrType (SetLocal local) = do
+    Ctx { locals } <- ask
+    t <- maybeToEither IndexOutOfRange $ locals !? local
+    return $ Val t ==> empty
+getInstrType (TeeLocal local) = do
+    Ctx { locals } <- ask
+    t <- maybeToEither IndexOutOfRange $ locals !? local
+    return $ Val t ==> Val t
+getInstrType (GetGlobal global) = do
+    Ctx { globals } <- ask
+    t <- maybeToEither IndexOutOfRange $ asType <$> globals !? global
+    return $ empty ==> t
+getInstrType (SetGlobal global) = do
+    Ctx { globals } <- ask
+    t <- maybeToEither IndexOutOfRange $ asType <$> globals !? global
+    return $ t ==> empty
+-- TODO: check memory alignment
+getInstrType (I32Load _) = do
+    Ctx { mems } <- ask
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I32
+getInstrType (I64Load _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I64
+getInstrType (F32Load _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> F32
+getInstrType (F64Load _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> F64
+getInstrType (I32Load8S _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I32
+getInstrType (I32Load8U _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I32
+getInstrType (I32Load16S _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I32
+getInstrType (I32Load16U _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I32
+getInstrType (I64Load8S _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I64
+getInstrType (I64Load8U _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I64
+getInstrType (I64Load16S _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I64
+getInstrType (I64Load16U _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I64
+getInstrType (I64Load32S _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I64
+getInstrType (I64Load32U _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I64
+getInstrType (I32Store _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, I32] ==> empty
+getInstrType (I64Store _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, I64] ==> empty
+getInstrType (F32Store _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, F32] ==> empty
+getInstrType (F64Store _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, F64] ==> empty
+getInstrType (I32Store8 _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, I32] ==> empty
+getInstrType (I32Store16 _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, I32] ==> empty
+getInstrType (I64Store8 _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, I64] ==> empty
+getInstrType (I64Store16 _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, I64] ==> empty
+getInstrType (I64Store32 _) = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ [I32, I64] ==> empty
+getInstrType CurrentMemory = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ empty ==> I32
+getInstrType GrowMemory = do
+    Ctx { mems } <- ask 
+    if length mems < 1 then throwError NoMemoryInModule else return $ I32 ==> I32
+getInstrType (I32Const _) = return $ empty ==> I32
+getInstrType (I64Const _) = return $ empty ==> I64
+getInstrType (F32Const _) = return $ empty ==> F32
+getInstrType (F64Const _) = return $ empty ==> F64
+getInstrType (IUnOp BS32 _) = return $ I32 ==> I32
+getInstrType (IUnOp BS64 _) = return $ I64 ==> I64
+getInstrType (IBinOp BS32 _) = return $ [I32, I32] ==> I32
+getInstrType (IBinOp BS64 _) = return $ [I64, I64] ==> I64
+getInstrType I32Eqz = return $ I32 ==> I32
+getInstrType I64Eqz = return $ I64 ==> I32
+getInstrType (IRelOp BS32 _) = return $ [I32, I32] ==> I32
+getInstrType (IRelOp BS64 _) = return $ [I64, I64] ==> I32
+getInstrType (FUnOp BS32 _) = return $ F32 ==> F32
+getInstrType (FUnOp BS64 _) = return $ F64 ==> F64
+getInstrType (FBinOp BS32 _) = return $ [F32, F32] ==> F32
+getInstrType (FBinOp BS64 _) = return $ [F64, F64] ==> F64
+getInstrType (FRelOp BS32 _) = return $ [F32, F32] ==> I32
+getInstrType (FRelOp BS64 _) = return $ [F64, F64] ==> I32
+getInstrType I32WrapI64 = return $ I64 ==> I32
+getInstrType (ITruncFU BS32 BS32) = return $ F32 ==> I32
+getInstrType (ITruncFU BS32 BS64) = return $ F64 ==> I32
+getInstrType (ITruncFU BS64 BS32) = return $ F32 ==> I64
+getInstrType (ITruncFU BS64 BS64) = return $ F64 ==> I64
+getInstrType (ITruncFS BS32 BS32) = return $ F32 ==> I32
+getInstrType (ITruncFS BS32 BS64) = return $ F64 ==> I32
+getInstrType (ITruncFS BS64 BS32) = return $ F32 ==> I64
+getInstrType (ITruncFS BS64 BS64) = return $ F64 ==> I64
+getInstrType I64ExtendSI32 = return $ I32 ==> I64
+getInstrType I64ExtendUI32 = return $ I32 ==> I64
+getInstrType (FConvertIU BS32 BS32) = return $ I32 ==> F32
+getInstrType (FConvertIU BS32 BS64) = return $ I64 ==> F32
+getInstrType (FConvertIU BS64 BS32) = return $ I32 ==> F64
+getInstrType (FConvertIU BS64 BS64) = return $ I64 ==> F64
+getInstrType (FConvertIS BS32 BS32) = return $ I32 ==> F32
+getInstrType (FConvertIS BS32 BS64) = return $ I64 ==> F32
+getInstrType (FConvertIS BS64 BS32) = return $ I32 ==> F64
+getInstrType (FConvertIS BS64 BS64) = return $ I64 ==> F64
+getInstrType F32DemoteF64 = return $ F64 ==> F32
+getInstrType F64PromoteF32 = return $ F32 ==> F64
+getInstrType (IReinterpretF BS32) = return $ F32 ==> I32
+getInstrType (IReinterpretF BS64) = return $ F64 ==> I64
+getInstrType (FReinterpretI BS32) = return $ I32 ==> F32
+getInstrType (FReinterpretI BS64) = return $ I64 ==> F64
+
+
+replace :: (Eq a) => a -> a -> [a] -> [a]
+replace _ _ [] = []
+replace x y (v:r) = (if x == v then y else v) : replace x y r
+
+unify :: Arrow -> Arrow -> Checker Arrow
+unify (f `Arrow` []) (f' `Arrow` t') =
+    return $ (f ++ f') `Arrow` t'
+unify (f `Arrow` t) ([] `Arrow` t') =
+    return $ f `Arrow` (t' ++ t)
+unify (f `Arrow` (Val v':t)) ((Val v:f') `Arrow` t') =
+    if v == v'
+    then unify (f `Arrow` t) (f' `Arrow` t')
+    else throwError TypeMismatch
+unify (f `Arrow` (Var r:t)) ((Val v:f') `Arrow` t') =
+    let subst = replace (Var r) (Val v) in
+    unify (subst f `Arrow` subst t) (f' `Arrow` t')
+unify (f `Arrow` (Val v:t)) ((Var r:f') `Arrow` t') =
+    let subst = replace (Var r) (Val v) in
+    unify (f `Arrow` t) (subst f' `Arrow` subst t')
+unify (f `Arrow` (Var r:t)) ((Var r':f') `Arrow` t') =
+    let subst = replace (Var r') (Var r) in
+    unify (f `Arrow` t) (subst f' `Arrow` subst t')
+unify (f `Arrow` (Any:t)) (f' `Arrow` t') =
+    return $ f `Arrow` t'
+unify (f `Arrow` t) ((Any:f') `Arrow` t') =
+    return $ f `Arrow` t'
+
+unify' (f `Arrow` t) (f' `Arrow` t') = unify (reverse f `Arrow` reverse t) (reverse f' `Arrow` reverse t')
+
+getExpressionType :: [Instruction] -> Checker Arrow
+getExpressionType instrs =
+    case reverse instrs of
+        [] -> return $ Any ==> Any
+        (i:rest) -> do
+            arr <- getInstrType i
+            go arr rest
+    where
+        go :: Arrow -> [Instruction] -> Checker Arrow
+        go arr [] = return arr
+        go arr (i:rest) = do
+            a <- getInstrType i
+            arr' <- unify' a arr
+            go arr' rest
+
+ctxFromModule :: [ValueType] -> [Maybe ValueType] -> Maybe ValueType -> Module -> Ctx
+ctxFromModule locals labels returns Module {types, functions, tables, mems, globals} =
+    Ctx {
+        types,
+        funcs = map ((types !!) . fromIntegral . funcType) functions,
+        tables = map (\(Table t) -> t) tables,
+        mems = map (\(Memory l) -> l) mems,
+        globals = map (\(Global g _) -> g) globals,
+        locals,
+        labels,
+        returns
+    }
+
+isFunctionValid :: Function -> Validator
+isFunctionValid Function {funcType, locals, body} mod@Module {types} =
+    let ft@(FuncType params results) = types !! fromIntegral funcType in
+    let r = safeHead results in
+    let ctx = ctxFromModule (params ++ locals) [r] r mod in
+    case runChecker ctx $ getExpressionType body of
+        Left err -> err
+        Right arr -> if arr == asArrow ft then Valid else TypeMismatch
+
+functionsShouldBeValid :: Validator
+functionsShouldBeValid mod@Module {functions} =
+    foldMap (flip isFunctionValid mod) functions
 
 tablesShouldBeValid :: Validator
 tablesShouldBeValid Module { imports, tables } =
@@ -202,5 +436,6 @@ validate mod = foldMap ($ mod) validators
         validators = [
                 tablesShouldBeValid,
                 shouldBeAtMostOneMemory,
-                exportNamesShouldBeDifferent
+                exportNamesShouldBeDifferent,
+                functionsShouldBeValid
             ]