type signature parsing for loop and if blocks

Author: SPY

src/Language/Wasm/Binary.hs

@@ -331,17 +331,17 @@ instance Serialize (Instruction Natural) where
         putWord8 0x02
         putBlockType blockType
         putExpression body
-    put (Loop result body) = do
+    put (Loop blockType body) = do
         putWord8 0x03
-        putResultType result
+        putBlockType blockType
         putExpression body
-    put If {resultType, true, false = []} = do
+    put If {blockType, true, false = []} = do
         putWord8 0x04
-        putResultType resultType
+        putBlockType blockType
         putExpression true
-    put If {resultType, true, false} = do
+    put If {blockType, true, false} = do
         putWord8 0x04
-        putResultType resultType
+        putBlockType blockType
         mapM_ put true
         putWord8 0x05 -- ELSE
         putExpression false
@@ -528,12 +528,12 @@ instance Serialize (Instruction Natural) where
             0x00 -> return Unreachable
             0x01 -> return Nop
             0x02 -> Block <$> getBlockType <*> getExpression
-            0x03 -> Loop <$> getResultType <*> getExpression
+            0x03 -> Loop <$> getBlockType <*> getExpression
             0x04 -> do
-                resultType <- getResultType
+                blockType <- getBlockType
                 (true, hasElse) <- getTrueBranch
                 false <- if hasElse then getExpression else return []
-                return $ If resultType true false
+                return $ If blockType true false
             0x0C -> Br <$> getULEB128 32
             0x0D -> BrIf <$> getULEB128 32
             0x0E -> BrTable <$> (map unIndex <$> getVec) <*> getULEB128 32

src/Language/Wasm/Builder.hs

@@ -41,7 +41,7 @@ module Language.Wasm.Builder (
     memorySize, growMemory,
     nop, Language.Wasm.Builder.drop, select,
     call, callIndirect, finish, br, brIf, brTable,
-    if', loop, {-block,-} when, for, while,
+    {-if', loop, block, when, for, while,-}
     trap, unreachable,
     appendExpr, after,
     Producer, OutType, produce, Consumer, (.=)
@@ -681,7 +681,7 @@ finish val = do
     appendExpr [Return]
 
 newtype Label i = Label Natural deriving (Show, Eq)
-
+{-
 when :: (Producer pred, OutType pred ~ Proxy I32)
     => pred
     -> GenFun ()
@@ -704,28 +704,28 @@ while pred body = do
             body
             loopLabel <- label
             if' () pred (br loopLabel) (return ())
-    if' () pred (loop () loopBody) (return ())
+    if' () pred (loop () loopBody) (return ())-}
 
 label :: GenFun (Label t)
 label = Label <$> ask
 
-if' :: (Producer pred, OutType pred ~ Proxy I32, Returnable res)
-    => res
-    -> pred
-    -> GenFun res
-    -> GenFun res
-    -> GenFun res
-if' res pred true false = do
-    produce pred
-    deep <- (+1) <$> ask
-    appendExpr [If (asResultValue res) (genExpr deep $ true) (genExpr deep $ false)]
-    return returnableValue
+-- if' :: (Producer pred, OutType pred ~ Proxy I32, Returnable res)
+--     => res
+--     -> pred
+--     -> GenFun res
+--     -> GenFun res
+--     -> GenFun res
+-- if' res pred true false = do
+--     produce pred
+--     deep <- (+1) <$> ask
+--     appendExpr [If (asResultValue res) (genExpr deep $ true) (genExpr deep $ false)]
+--     return returnableValue
 
-loop :: (Returnable res) => res -> GenFun res -> GenFun res
-loop res body = do
-    deep <- (+1) <$> ask
-    appendExpr [Loop (asResultValue res) (genExpr deep $ body)]
-    return returnableValue
+-- loop :: (Returnable res) => res -> GenFun res -> GenFun res
+-- loop res body = do
+--     deep <- (+1) <$> ask
+--     appendExpr [Loop (asResultValue res) (genExpr deep $ body)]
+--     return returnableValue
 
 -- block :: (Returnable res) => res -> GenFun res -> GenFun res
 -- block res body = do
@@ -988,39 +988,39 @@ asWord64 i
     | i >= 0 = fromIntegral i
     | otherwise = 0xFFFFFFFFFFFFFFFF - (fromIntegral (abs i)) + 1
 
-rts :: Module
-rts = genMod $ do
-    gc <- importFunction "rts" "gc" () [I32]
-    memory 10 Nothing
-
-    stackStart <- global Const i32 0
-    stackEnd <- global Const i32 0
-    stackBase <- global Mut i32 0
-    stackTop <- global Mut i32 0
-
-    retReg <- global Mut i32 0
-    tmpReg <- global Mut i32 0
-
-    heapStart <- global Mut i32 0
-    heapNext <- global Mut i32 0
-    heapEnd <- global Mut i32 0
-
-    aligned <- fun i32 $ do
-        size <- param i32
-        (size `add` i32c 3) `and` i32c 0xFFFFFFFC
-    alloc <- funRec i32 $ \self -> do
-        size <- param i32
-        alignedSize <- local i32
-        addr <- local i32
-        alignedSize .= call aligned [arg size]
-        if' i32 ((heapNext `add` alignedSize) `lt_u` heapEnd)
-            (do
-                addr .= heapNext
-                heapNext .= heapNext `add` alignedSize
-                ret addr
-            )
-            (do
-                call gc []
-                call self [arg size]
-            )
-    return ()
+-- rts :: Module
+-- rts = genMod $ do
+--     gc <- importFunction "rts" "gc" () [I32]
+--     memory 10 Nothing
+
+--     stackStart <- global Const i32 0
+--     stackEnd <- global Const i32 0
+--     stackBase <- global Mut i32 0
+--     stackTop <- global Mut i32 0
+
+--     retReg <- global Mut i32 0
+--     tmpReg <- global Mut i32 0
+
+--     heapStart <- global Mut i32 0
+--     heapNext <- global Mut i32 0
+--     heapEnd <- global Mut i32 0
+
+--     aligned <- fun i32 $ do
+--         size <- param i32
+--         (size `add` i32c 3) `and` i32c 0xFFFFFFFC
+--     alloc <- funRec i32 $ \self -> do
+--         size <- param i32
+--         alignedSize <- local i32
+--         addr <- local i32
+--         alignedSize .= call aligned [arg size]
+--         if' i32 ((heapNext `add` alignedSize) `lt_u` heapEnd)
+--             (do
+--                 addr .= heapNext
+--                 heapNext .= heapNext `add` alignedSize
+--                 ret addr
+--             )
+--             (do
+--                 call gc []
+--                 call self [arg size]
+--             )
+--     return ()

src/Language/Wasm/Interpreter.hs

@@ -646,14 +646,22 @@ eval budget store FunctionInstance { funcType, moduleInstance, code = Function {
                 Break n r ctx' -> return $ Break (n - 1) r ctx'
                 Done ctx'@EvalCtx{ labels = (_:rest) } -> return $ Done ctx' { labels = rest }
                 command -> return command
-        step ctx loop@(Loop resType expr) = do
+        step ctx loop@(Loop blockType expr) = do
+            let resType = case blockType of
+                    Inline Nothing -> []
+                    Inline (Just valType) -> [valType]
+                    TypeIndex typeIdx -> results $ funcTypes moduleInstance ! fromIntegral typeIdx
             res <- go ctx { labels = Label resType : labels ctx } expr
             case res of
                 Break 0 r EvalCtx{ locals = ls } -> step ctx { locals = ls, stack = r ++ stack ctx } loop
                 Break n r ctx' -> return $ Break (n - 1) r ctx'
                 Done ctx'@EvalCtx{ labels = (_:rest) } -> return $ Done ctx' { labels = rest }
                 command -> return command
-        step ctx@EvalCtx{ stack = (VI32 v): rest } (If resType true false) = do
+        step ctx@EvalCtx{ stack = (VI32 v): rest } (If blockType true false) = do
+            let resType = case blockType of
+                    Inline Nothing -> []
+                    Inline (Just valType) -> [valType]
+                    TypeIndex typeIdx -> results $ funcTypes moduleInstance ! fromIntegral typeIdx
             let expr = if v /= 0 then true else false
             res <- go ctx { labels = Label resType : labels ctx, stack = rest } expr
             case res of

src/Language/Wasm/Parser.y

@@ -660,68 +660,32 @@ raw_instr :: { [Instruction] }
         then Right $ [BlockInstr $2 tu (instr ++ instr')]
         else Left "Block labels have to match"
     }
-    | 'loop' opt(ident) raw_loop {% (: []) `fmap` $3 $2 }
-    | 'if' opt(ident) raw_if_result {% $3 $2 }
-
-raw_loop :: { Maybe Ident -> Either String Instruction }
-    : 'end' opt(ident) {
-        \ident ->
-            if ident == $2 || isNothing $2
-            then Right $ LoopInstr ident [] []
-            else Left "Loop labels have to match"
-    }
-    | raw_instr list(instruction) 'end' opt(ident) {
-        \ident ->
-            if ident == $4 || isNothing $4
-            then Right $ LoopInstr ident [] ($1 ++ concat $2)
-            else Left "Loop labels have to match"
-    }
-    | '(' raw_loop1 { $2 }
-
-raw_loop1 :: { Maybe Ident -> Either String Instruction }
-    : 'result' list(valtype) ')' list(instruction) 'end' opt(ident) {
-        \ident ->
-            if ident == $6 || isNothing $6
-            then Right $ LoopInstr ident $2 (concat $4)
-            else Left "Loop labels have to match"
-    }
-    | folded_instr1 list(instruction) 'end' opt(ident) {
-        \ident ->
-            if ident == $4 || isNothing $4
-            then Right $ LoopInstr ident [] ($1 ++ concat $2)
-            else Left "Loop labels have to match"
-    }
-
-raw_if_result :: { Maybe Ident -> Either String [Instruction] }
-    : raw_else {
-        \ident ->
-            if ident == (snd $1) || isNothing (snd $1)
-            then Right [IfInstr ident [] [] $ fst $1]
-            else Left "If labels have to match"
+    | 'loop' opt(ident) typeuse_cont(pair(folded_instr1_list, block_end), block_end) {%
+        let (tu, rest) = $3 in
+        let (instr, (instr', identAfter)) = either (\a -> ([], a)) id rest in
+        if $2 == identAfter || isNothing identAfter
+        then Right $ [LoopInstr $2 tu (instr ++ instr')]
+        else Left "Block labels have to match"
     }
-    | raw_instr list(instruction) raw_else {
-        \ident ->
-            if ident == (snd $3) || isNothing (snd $3)
-            then Right [IfInstr ident [] ($1 ++ concat $2) $ fst $3]
-            else Left "If labels have to match"
+    | 'if' opt(ident) typeuse_cont(pair(folded_instr1_list, raw_if_end), raw_if_end) {%
+        let (tu, rest) = $3 in
+        let (trueBranch, falseBranch, identAfter) = either id (\(t, (t', f, i)) -> (t ++ t', f, i)) rest in
+        if $2 == identAfter || isNothing identAfter
+        then Right $ [IfInstr $2 tu trueBranch falseBranch]
+        else Left "If labels have to match"
     }
-    | '(' raw_if_result1 { $2 }
 
-raw_if_result1 :: { Maybe Ident -> Either String [Instruction] }
-    : 'result' list(valtype) ')' list(instruction) raw_else {
-        \ident ->
-            if ident == (snd $5) || isNothing (snd $5)
-            then Right [IfInstr ident $2 (concat $4) $ fst $5]
-            else Left "If labels have to match"
+raw_if_end
+    : raw_if_else {
+        let (falseBranch, ident) = $1 in
+        ([], falseBranch, ident)
     }
-    | folded_instr1 list(instruction) raw_else {
-        \ident ->
-            if ident == (snd $3) || isNothing (snd $3)
-            then Right [IfInstr ident [] ($1 ++ concat $2) $ fst $3]
-            else Left "If labels have to match"
+    | raw_instr list(instruction) raw_if_else {
+        let (falseBranch, ident) = $3 in
+        ($1 ++ concat $2, falseBranch, ident)
     }
 
-raw_else :: { ([Instruction], Maybe Ident) }
+raw_if_else :: { ([Instruction], Maybe Ident) }
     : 'end' opt(ident) { ([], $2) }
     | 'else' opt(ident) list(instruction) 'end' opt(ident) {%
         if matchIdents $2 $5
@@ -748,28 +712,14 @@ folded_instr1 :: { [Instruction] }
         let (instr, instr') = either (\a -> ([], a)) id rest in
         [BlockInstr $2 typeUse (instr ++ instr')]
     }
-    | 'loop' opt(ident) folded_loop { [$3 $2] }
-    | 'if' opt(ident) '(' folded_if_result { $4 $2 }
-
-folded_loop :: { Maybe Ident -> Instruction }
-    : ')' { \ident -> LoopInstr ident [] [] }
-    | '(' folded_loop1 { $2 }
-    | raw_instr list(instruction) ')' { \ident -> LoopInstr ident [] ($1 ++ concat $2) }
-
-folded_loop1 :: { Maybe Ident -> Instruction }
-    : 'result' list(valtype) ')' list(instruction) ')' { \ident -> LoopInstr ident $2 (concat $4) }
-    | folded_instr1 list(instruction) ')' { \ident -> LoopInstr ident [] ($1 ++ concat $2) }
-
-folded_if_result :: { Maybe Ident -> [Instruction] }
-    : 'result' list(valtype) ')' '(' folded_then_else {
-        \ident ->
-            let (pred, (trueBranch, falseBranch)) = $5 in
-            pred ++ [IfInstr ident $2 trueBranch falseBranch]
+    | 'loop' opt(ident) typeuse_cont(pair(folded_instr1_list, instr_list_closed), instr_list_closed) {
+        let (typeUse, rest) = $3 in
+        let (instr, instr') = either (\a -> ([], a)) id rest in
+        [LoopInstr $2 typeUse (instr ++ instr')]
     }
-    | folded_then_else {
-        \ident ->
-            let (pred, (trueBranch, falseBranch)) = $1 in
-            pred ++ [IfInstr ident [] trueBranch falseBranch]
+    | 'if' opt(ident) '(' typeuse1_cont(folded_then_else, never) {
+        let (typeUse, Right (pred, (trueBranch, falseBranch))) = $4 in
+        pred ++ [IfInstr $2 typeUse trueBranch falseBranch]
     }
 
 folded_then_else :: { ([Instruction], ([Instruction], [Instruction])) }
@@ -1225,12 +1175,12 @@ data Instruction =
     }
     | LoopInstr {
         label :: Maybe Ident,
-        resultType :: [ValueType],
+        blockType :: TypeUse,
         body :: [Instruction]
     }
     | IfInstr {
         label :: Maybe Ident,
-        resultType :: [ValueType],
+        blockType :: TypeUse,
         trueBranch :: [Instruction],
         falseBranch :: [Instruction]
     }
@@ -1503,10 +1453,10 @@ desugarize fields = do
             matchTypeUse defs typeUse
         extractTypeDefFromInstruction defs (BlockInstr { body, blockType }) =
             extractTypeDefFromInstructions (matchTypeUse defs blockType) body
-        extractTypeDefFromInstruction defs (LoopInstr { body }) =
-            extractTypeDefFromInstructions defs body
-        extractTypeDefFromInstruction defs (IfInstr { trueBranch, falseBranch }) =
-            extractTypeDefFromInstructions defs $ trueBranch ++ falseBranch
+        extractTypeDefFromInstruction defs (LoopInstr { body, blockType }) =
+            extractTypeDefFromInstructions (matchTypeUse defs blockType) body
+        extractTypeDefFromInstruction defs (IfInstr { blockType, trueBranch, falseBranch }) =
+            extractTypeDefFromInstructions (matchTypeUse defs blockType) $ trueBranch ++ falseBranch
         extractTypeDefFromInstruction defs _ = defs
 
         funcTypesEq :: FuncType -> FuncType -> Bool
@@ -1659,14 +1609,26 @@ desugarize fields = do
                     Just idx -> return $ S.TypeIndex idx
                     Nothing -> Left "unknown type"
             S.Block bt <$> mapM (synInstrToStruct ctx') body
-        synInstrToStruct ctx LoopInstr {label, resultType, body} =
-            let ctx' = ctx { ctxLabels = label : ctxLabels ctx } in
-            S.Loop resultType <$> mapM (synInstrToStruct ctx') body
-        synInstrToStruct ctx IfInstr {label, resultType, trueBranch, falseBranch} = do
+        synInstrToStruct ctx@FunCtx { ctxMod = Module { types } } LoopInstr {label, blockType, body} = do
+            let ctx' = ctx { ctxLabels = label : ctxLabels ctx }
+            bt <- case blockType of
+                AnonimousTypeUse (FuncType [] []) -> return $ S.Inline Nothing
+                AnonimousTypeUse (FuncType [] [vt]) -> return $ S.Inline (Just vt)
+                typed -> case getTypeIndex types typed of
+                    Just idx -> return $ S.TypeIndex idx
+                    Nothing -> Left "unknown type"
+            S.Loop bt <$> mapM (synInstrToStruct ctx') body
+        synInstrToStruct ctx@FunCtx { ctxMod = Module { types } } IfInstr {label, blockType, trueBranch, falseBranch} = do
             let ctx' = ctx { ctxLabels = label : ctxLabels ctx }
+            bt <- case blockType of
+                AnonimousTypeUse (FuncType [] []) -> return $ S.Inline Nothing
+                AnonimousTypeUse (FuncType [] [vt]) -> return $ S.Inline (Just vt)
+                typed -> case getTypeIndex types typed of
+                    Just idx -> return $ S.TypeIndex idx
+                    Nothing -> Left "unknown type"
             trueBranch' <- mapM (synInstrToStruct ctx') trueBranch
             falseBranch' <- mapM (synInstrToStruct ctx') falseBranch
-            return $ S.If resultType trueBranch' falseBranch'
+            return $ S.If bt trueBranch' falseBranch'
 
         synFunctionToStruct :: Module -> Function -> Either String S.Function
         synFunctionToStruct mod Function { funcType, locals, body } = do