add return type for select

src/Language/Wasm/Binary.hs

@@ -368,7 +368,7 @@ instance Serialize (Instruction Natural) where
     put (CallIndirect tableIdx typeIdx) = putWord8 0x11 >> putULEB128 typeIdx >> putULEB128 tableIdx
     -- Parametric instructions
     put Drop = putWord8 0x1A
-    put Select = putWord8 0x1B
+    put (Select _) = putWord8 0x1B
     -- Variable instructions
     put (GetLocal idx) = putWord8 0x20 >> putULEB128 idx
     put (SetLocal idx) = putWord8 0x21 >> putULEB128 idx
@@ -569,7 +569,7 @@ instance Serialize (Instruction Natural) where
                 return $ CallIndirect tableIdx typeIdx
             -- Parametric instructions
             0x1A -> return $ Drop
-            0x1B -> return $ Select
+            0x1B -> return $ Select Nothing
             -- Variable instructions
             0x20 -> GetLocal <$> getULEB128 32
             0x21 -> SetLocal <$> getULEB128 32

src/Language/Wasm/Builder.hs

@@ -197,7 +197,7 @@ select pred a b = select' (produce pred) (produce a) (produce b)
             a
             res <- b
             pred
-            appendExpr [Select]
+            appendExpr [Select Nothing]
             return res
 
 iBinOp :: (Producer a, Producer b, OutType a ~ OutType b, IsInt (OutType a) ~ True) => IBinOp -> a -> b -> GenFun (OutType a)

src/Language/Wasm/Interpreter.hs

@@ -808,7 +808,7 @@ eval budget store FunctionInstance { funcType, moduleInstance, code = Function {
         step ctx@EvalCtx{ stack = st } (RefFunc index) =
             return $ Done ctx { stack = RF (Just index) : st }
         step ctx@EvalCtx{ stack = (_:rest) } Drop = return $ Done ctx { stack = rest }
-        step ctx@EvalCtx{ stack = (VI32 test:val2:val1:rest) } Select =
+        step ctx@EvalCtx{ stack = (VI32 test:val2:val1:rest) } (Select _) =
             if test == 0
             then return $ Done ctx { stack = val2 : rest }
             else return $ Done ctx { stack = val1 : rest }

src/Language/Wasm/Parser.y

@@ -437,7 +437,6 @@ plaininstr :: { PlainInstr }
     | 'return'                       { Return }
     | 'call' index                   { Call $2 }
     | 'drop'                         { Drop }
-    | 'select'                       { Select }
     -- reference instructions
     | 'ref.null' heaptype            { RefNull $2 }
     | 'ref.is_null'                  { RefIsNull }
@@ -695,6 +694,20 @@ memarg4 :: { MemArg }
 memarg8 :: { MemArg }
     : opt(offset) opt(align) {% parseMemArg 8 $1 $2 }
 
+select_type_or_instructions(terminator)
+    : terminator { ($1, Nothing, []) }
+    | '(' select_type_or_instructions1(terminator) { $2 }
+
+select_type_or_instructions1(terminator) 
+    : 'result' list(valtype) ')' mixed_instruction_list(terminator) {
+        let (end, instr) = $4 in
+        (end, Just $2, instr)
+    }
+    | folded_instr_list(terminator) {
+        let (end, instr) = $1 in
+        (end, Nothing, instr)
+    }
+
 instruction_list(terminator)
     : terminator { ($1, []) }
     | plaininstr mixed_instruction_list(terminator) { ([PlainInstr $1] ++) `fmap` $2 }
@@ -703,6 +716,9 @@ instruction_list(terminator)
         let (tu, instr, end) = $3 in
         onlyAnonimParams tu >> (return (end, [PlainInstr $ CallIndirect tableIdx tu] ++ instr))
     }
+    | 'select' select_type_or_instructions(terminator) {
+        let (end, t, instr) = $2 in (end, [PlainInstr $ Select t] ++ instr)
+    }
     | 'block' opt(ident) typeuse('end') opt(ident) mixed_instruction_list(terminator) {% do
         let (tu, instr, _) = $3 
         matchIdents $2 $4
@@ -744,6 +760,9 @@ folded_instr1 :: { [Instruction] }
         let (tu, instr, _) = $3 in
         onlyAnonimParams tu >> (return $ instr ++ [PlainInstr $ CallIndirect tableIdx tu])
     }
+    | 'select' select_type_or_instructions(')') {
+        let (_, t, instr) = $2 in instr ++ [PlainInstr $ Select t]
+    }
     | 'block' opt(ident) typeuse(')') {%
         let (typeUse, instr, _) = $3 in
         onlyAnonimParams typeUse >> (return [BlockInstr $2 typeUse instr])
@@ -1206,7 +1225,7 @@ data PlainInstr =
     | RefExtern Natural
     -- Parametric instructions
     | Drop
-    | Select
+    | Select (Maybe [ValueType])
     -- Variable instructions
     | GetLocal LocalIndex
     | SetLocal LocalIndex
@@ -1671,7 +1690,7 @@ desugarize fields = do
                         Nothing -> Left "unknown type"
                 Nothing -> Left "unknown table"
         synInstrToStruct _ (PlainInstr Drop) = return $ S.Drop
-        synInstrToStruct _ (PlainInstr Select) = return $ S.Select
+        synInstrToStruct _ (PlainInstr (Select vt)) = return $ S.Select vt
         synInstrToStruct _ (PlainInstr (RefNull elType)) = return $ S.RefNull elType
         synInstrToStruct _ (PlainInstr RefIsNull) = return $ S.RefIsNull
         synInstrToStruct FunCtx { ctxMod } (PlainInstr (RefFunc funIdx)) =

src/Language/Wasm/Structure.hs

@@ -145,7 +145,7 @@ data Instruction index =
     | RefExtern Natural
     -- Parametric instructions
     | Drop
-    | Select
+    | Select (Maybe [ValueType])
     -- Variable instructions
     | GetLocal index
     | SetLocal index

src/Language/Wasm/Validate.hs

@@ -71,6 +71,7 @@ type Validator = Module -> ValidationResult
 data VType =
     Val ValueType
     | Var
+    | NonRefVar
     | Any
     deriving (Show, Eq)
 
@@ -105,6 +106,11 @@ asArrow (FuncType params results) = Arrow (map Val params) (map Val $ reverse re
 isArrowMatch :: Arrow -> Arrow -> Bool
 isArrowMatch (f `Arrow` t) ( f' `Arrow` t') = isEndMatch f f' && isEndMatch t t'
     where
+        isRef :: VType -> Bool
+        isRef (Val Func) = True
+        isRef (Val Extern) = True
+        isRef _            = False
+
         isEndMatch :: End -> End -> Bool
         isEndMatch (Any:l) (Any:r) =
             let (leftTail, rightTail) = unzip $ zip (takeWhile (/= Any) $ reverse l) (takeWhile (/= Any) $ reverse r) in
@@ -121,6 +127,12 @@ isArrowMatch (f `Arrow` t) ( f' `Arrow` t') = isEndMatch f f' && isEndMatch t t'
         isEndMatch (x:l) (Var:r) =
             let subst = replace Var x in
             isEndMatch (subst l) (subst r)
+        isEndMatch (NonRefVar:l) (x:r) =
+            let subst = replace NonRefVar x in
+            isEndMatch (subst l) (subst r)
+        isEndMatch (x:l) (NonRefVar:r) =
+            let subst = replace NonRefVar x in
+            isEndMatch (subst l) (subst r)
         isEndMatch (Val v:l) (Val v':r) = v == v' && isEndMatch l r
         isEndMatch [] [] = True
         isEndMatch _ _ = False
@@ -263,9 +275,13 @@ getInstrType (CallIndirect tableIdx sign) = do
 getInstrType Drop = do
     var <- freshVar
     return $ var ==> empty
-getInstrType Select = do
+getInstrType (Select Nothing) = do
-    var <- freshVar
+    var <- return NonRefVar
     return $ [var, var, Val I32] ==> var
+getInstrType (Select (Just vt)) =
+    case vt of
+        [t] -> return $ [t, t, I32] ==> t
+        _ -> throwError InvalidResultArity
 getInstrType (RefNull elType) = do
     let t = case elType of { FuncRef -> Func; ExternRef -> Extern }
     return $ empty ==> Val t
@@ -486,6 +502,10 @@ getExpressionTypeWithInput inp = fmap (inp `Arrow`) . foldM go inp
             (f `Arrow` t) <- getInstrType instr
             matchStack stack (reverse f) t
         
+        isRef (Func)   = True
+        isRef (Extern) = True
+        isRef _        = False
+        
         matchStack :: [VType] -> [VType] -> [VType] -> Checker [VType]
         matchStack stack@(Any:_) _arg res = return $ res ++ stack
         matchStack (Val v:stack) (Val v':args) res =
@@ -499,6 +519,16 @@ getExpressionTypeWithInput inp = fmap (inp `Arrow`) . foldM go inp
         matchStack (Var:stack) (Val v:args) res =
             let subst = replace Var (Val v) in
             matchStack stack (subst args) (subst res)
+        matchStack (Val v:stack) (NonRefVar:args) res =
+            let subst = replace NonRefVar (Val v) in
+            if isRef v
+            then throwError $ TypeMismatch (empty ==> empty) (empty ==> empty)
+            else matchStack stack (subst args) (subst res)
+        matchStack (NonRefVar:stack) (Val v:args) res =
+            let subst = replace NonRefVar (Val v) in
+            if isRef v
+            then throwError $ TypeMismatch (empty ==> empty) (empty ==> empty)
+            else matchStack stack (subst args) (subst res)
         matchStack stack [] res = return $ res ++ stack
         matchStack [] args res = throwError $ TypeMismatch ((reverse args) `Arrow` res) ([] `Arrow` [])
         matchStack _ _ _ = error "inconsistent checker state"

tests/Test.hs

@@ -19,7 +19,7 @@ main = do
   files <-
     filter (not . List.isPrefixOf "simd") . filter (List.isSuffixOf ".wast")
       <$> Directory.listDirectory "tests/spec"
-  -- let files = ["table_fill.wast"]
+  -- let files = ["select.wast"]
   scriptTestCases <- (`mapM` files) $ \file -> do
     test <- LBS.readFile ("tests/spec/" ++ file)
     return $ testCase file $ do