[wasm] Add support for blocks with result

This changes wasm blocks (only the "plain" block, not ifs, trys etc.)
to have an optional result type. A wasm block with a result produces
a Variable in its WasmEndBlock operation. The value is either the input
argument to the WasmEndBlock (in case of regularly reaching the end of
the block) or the arguments passed into a WasmBranch or WasmBranchIf
operation that branches to the WasmBeginBlock's label.

As these branch instructions need to provide the argument(s) matching
the block signature, the label type is extended to also contain type
information about the arguments to be passed along a branch.

Change-Id: I4a1af782330c7a71320a2d4b1949dba081d1d4de
Reviewed-on: https://chrome-internal-review.googlesource.com/c/v8/fuzzilli/+/7967831
Reviewed-by: Carl Smith <[email protected]>
Commit-Queue: Matthias Liedtke <[email protected]>

Author: Liedtke

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -3236,21 +3236,30 @@ public class ProgramBuilder {
         // TODO(cffsmith): I think the best way to handle these types of blocks is to treat them like inline functions that have a signature. E.g. they behave like a definition and call of a wasmfunction. The output should be the output of the signature.
         public func wasmBuildBlock(with signature: Signature, args: [Variable], body: (Variable, [Variable]) -> ()) {
             assert(signature.parameters.count == args.count)
+            assert(signature.outputType == .nothing)
             let instr = b.emit(WasmBeginBlock(with: signature), withInputs: args)
-            b.setType(ofVariable: instr.innerOutput(0), to: .label)
             body(instr.innerOutput(0), Array(instr.innerOutputs(1...)))
-            b.emit(WasmEndBlock())
+            b.emit(WasmEndBlock(outputType: signature.outputType))
+        }
+
+        @discardableResult
+        public func wasmBuildBlockWithResult(with signature: Signature, args: [Variable], body: (Variable, [Variable]) -> Variable) -> Variable {
+            assert(signature.parameters.count == args.count)
+            assert(signature.outputType != .nothing)
+            let instr = b.emit(WasmBeginBlock(with: signature), withInputs: args)
+            let result = body(instr.innerOutput(0), Array(instr.innerOutputs(1...)))
+            return b.emit(WasmEndBlock(outputType: signature.outputType), withInputs: [result]).output
         }
 
         // This can branch to label variables only, has a variable input for dataflow purposes.
-        public func wasmBranch(to label: Variable) {
-            assert(b.type(of: label) == .label)
-            b.emit(WasmBranch(), withInputs: [label])
+        public func wasmBranch(to label: Variable, args: [Variable] = []) {
+            assert(b.type(of: label).Is(.anyLabel))
+            b.emit(WasmBranch(labelTypes: b.type(of: label).wasmLabelType!.parameters), withInputs: [label] + args)
         }
 
-        public func wasmBranchIf(_ condition: Variable, to label: Variable) {
-            assert(b.type(of: label) == .label)
-            b.emit(WasmBranchIf(), withInputs: [label, condition])
+        public func wasmBranchIf(_ condition: Variable, to label: Variable, args: [Variable] = []) {
+            assert(b.type(of: label).Is(.label(args.map({b.type(of: $0)}))), "label type \(b.type(of: label)) doesn't match argument types \(args.map({b.type(of: $0)}))")
+            b.emit(WasmBranchIf(labelTypes: b.type(of: label).wasmLabelType!.parameters), withInputs: [label] + args + [condition])
         }
 
         public func wasmBuildIfElse(_ condition: Variable, ifBody: () -> Void, elseBody: () -> Void) {

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -560,10 +560,17 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         // Choose a few random wasm values as arguments if available.
         let args = (0..<5).map {_ in b.findVariable {b.type(of: $0).Is(.wasmPrimitive)}}.filter {$0 != nil}.map {$0!}
         let parameters = args.map {arg in Parameter.plain(b.type(of: arg))}
-        // TODO(mliedtke): Also add support for return values (and probably find a suitable mechanism that we can also use for functions).
-        // A good solution would probably be to make the endBlock take the return type as an input.
-        function.wasmBuildBlock(with: parameters => .nothing, args: args) { label, args in
-            b.buildRecursive()
+        // TODO(mliedtke): The selection of types is in sync with ProgramBuilder::randomWasmSignature(). This should allow more types.
+        let outputType: ILType = chooseUniform(from: [.wasmi32, .wasmi64, .wasmf32, .wasmf64, .nothing])
+        if outputType != .nothing {
+            function.wasmBuildBlockWithResult(with: parameters => outputType, args: args) { label, args in
+                b.buildRecursive()
+                return b.randomVariable(ofType: outputType) ?? function.generateRandomWasmVar(ofType: outputType)
+            }
+        } else {
+            function.wasmBuildBlock(with: parameters => outputType, args: args) { label, args in
+                b.buildRecursive()
+            }
         }
     },
 
@@ -603,7 +610,7 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         }
     },
 
-    RecursiveCodeGenerator("WasmLegacyTryDelegateGenerator", inContext: .wasmFunction, inputs: .required(.label)) { b, label in
+    RecursiveCodeGenerator("WasmLegacyTryDelegateGenerator", inContext: .wasmFunction, inputs: .required(.anyLabel)) { b, label in
         let function = b.currentWasmModule.currentWasmFunction
         // Choose a few random wasm values as arguments if available.
         let args = (0..<5).map {_ in b.findVariable {b.type(of: $0).Is(.wasmPrimitive)}}.filter {$0 != nil}.map {$0!}
@@ -687,14 +694,20 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         b.currentWasmModule.addTag(parameterTypes: b.randomTagParameters())
     },
 
-    CodeGenerator("WasmBranchGenerator", inContext: .wasmFunction, inputs: .required(.label)) { b, label in
+    CodeGenerator("WasmBranchGenerator", inContext: .wasmFunction, inputs: .required(.anyLabel)) { b, label in
         let function = b.currentWasmModule.currentWasmFunction
-        function.wasmBranch(to: label)
+        let args = b.type(of: label).wasmLabelType!.parameters.map {
+            b.randomVariable(ofType: $0) ?? function.generateRandomWasmVar(ofType: $0)
+        }
+        function.wasmBranch(to: label, args: args)
     },
 
-    CodeGenerator("WasmBranchIfGenerator", inContext: .wasmFunction, inputs: .required(.label, .wasmi32)) { b, label, conditionVar in
+    CodeGenerator("WasmBranchIfGenerator", inContext: .wasmFunction, inputs: .required(.anyLabel, .wasmi32)) { b, label, conditionVar in
         let function = b.currentWasmModule.currentWasmFunction
-        function.wasmBranchIf(conditionVar, to: label)
+        let args = b.type(of: label).wasmLabelType!.parameters.map {
+            b.randomVariable(ofType: $0) ?? function.generateRandomWasmVar(ofType: $0)
+        }
+        function.wasmBranchIf(conditionVar, to: label, args: args)
     },
 
     CodeGenerator("ConstSimd128Generator", inContext: .wasmFunction) { b in

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1194,8 +1194,10 @@ extension Instruction: ProtobufConvertible {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
                     $0.returnType = ILTypeToWasmTypeEnum(op.signature.outputType)
                 }
-            case .wasmEndBlock(_):
-                $0.wasmEndBlock = Fuzzilli_Protobuf_WasmEndBlock()
+            case .wasmEndBlock(let op):
+                $0.wasmEndBlock = Fuzzilli_Protobuf_WasmEndBlock.with {
+                    $0.returnType = ILTypeToWasmTypeEnum(op.outputType)
+                }
             case .wasmBeginLoop(let op):
                 $0.wasmBeginLoop = Fuzzilli_Protobuf_WasmBeginLoop.with {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
@@ -1238,10 +1240,14 @@ extension Instruction: ProtobufConvertible {
                 $0.wasmDefineTag = Fuzzilli_Protobuf_WasmDefineTag.with {
                     $0.parameters = convertParametersToWasmTypeEnums(op.parameters)
                 }
-            case .wasmBranch(_):
-                $0.wasmBranch = Fuzzilli_Protobuf_WasmBranch()
-            case .wasmBranchIf(_):
-                $0.wasmBranchIf = Fuzzilli_Protobuf_WasmBranchIf()
+            case .wasmBranch(let op):
+                $0.wasmBranch = Fuzzilli_Protobuf_WasmBranch.with {
+                    $0.parameters = op.labelTypes.map(ILTypeToWasmTypeEnum)
+                }
+            case .wasmBranchIf(let op):
+                $0.wasmBranchIf = Fuzzilli_Protobuf_WasmBranchIf.with {
+                    $0.parameters = op.labelTypes.map(ILTypeToWasmTypeEnum)
+                }
             case .wasmBeginIf(let op):
                 $0.wasmBeginIf = Fuzzilli_Protobuf_WasmBeginIf.with {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
@@ -1991,8 +1997,8 @@ extension Instruction: ProtobufConvertible {
                 Parameter.plain(WasmTypeEnumToILType(param))
             })
             op = WasmBeginBlock(with: parameters => WasmTypeEnumToILType(p.returnType))
-        case .wasmEndBlock(_):
-            op = WasmEndBlock()
+        case .wasmEndBlock(let p):
+            op = WasmEndBlock(outputType: WasmTypeEnumToILType(p.returnType))
         case .wasmBeginLoop(let p):
             let parameters: [Parameter] = p.parameters.map({ param in
                 Parameter.plain(WasmTypeEnumToILType(param))
@@ -2033,10 +2039,10 @@ extension Instruction: ProtobufConvertible {
             op = WasmDefineTag(parameters: p.parameters.map({ param in
                 Parameter.plain(WasmTypeEnumToILType(param))
             }))
-        case .wasmBranch(_):
-            op = WasmBranch()
-        case .wasmBranchIf(_):
-            op = WasmBranchIf()
+        case .wasmBranch(let p):
+            op = WasmBranch(labelTypes: p.parameters.map(WasmTypeEnumToILType))
+        case .wasmBranchIf(let p):
+            op = WasmBranchIf(labelTypes: p.parameters.map(WasmTypeEnumToILType))
         case .wasmBeginIf(let p):
             let parameters: [Parameter] = p.parameters.map({ param in
                 Parameter.plain(WasmTypeEnumToILType(param))

Sources/Fuzzilli/FuzzIL/TypeSystem.swift

@@ -175,7 +175,11 @@ public struct ILType: Hashable {
     // Internal types
 
     // This type is used to indicate block labels in wasm.
-    public static let label: ILType = ILType(definiteType: .label)
+    public static func label(_ parameterTypes: [ILType] = []) -> ILType {
+        return ILType(definiteType: .label, ext: TypeExtension(group: "WasmLabel", properties: [], methods: [], signature: nil, wasmExt: WasmLabelType(parameterTypes)))
+    }
+
+    public static let anyLabel: ILType = ILType(definiteType: .label, ext: TypeExtension(group: "WasmLabel", properties: [], methods: [], signature: nil, wasmExt: nil))
 
     /// A label that allows rethrowing the caught exception of a catch block.
     public static let exceptionLabel: ILType = ILType(definiteType: .exceptionLabel)
@@ -409,6 +413,14 @@ public struct ILType: Hashable {
         return wasmTagType != nil && ext?.group == "WasmTag"
     }
 
+    public var wasmLabelType: WasmLabelType? {
+        return wasmType as? WasmLabelType
+    }
+
+    public var isWasmLabelType: Bool {
+        return wasmTagType != nil
+    }
+
     public var properties: Set<String> {
         return ext?.properties ?? Set()
     }
@@ -977,11 +989,12 @@ public class WasmTagType: WasmTypeExtension {
 
     override func isEqual(to other: WasmTypeExtension) -> Bool {
         guard let other = other as? WasmTagType else { return false }
-        return self.parameters == other.parameters
+        return self.parameters == other.parameters && self.isJSTag == other.isJSTag
     }
 
     override public func hash(into hasher: inout Hasher) {
         hasher.combine(parameters)
+        hasher.combine(isJSTag)
     }
 
     init(_ parameters: ParameterList, isJSTag: Bool = false) {
@@ -990,6 +1003,27 @@ public class WasmTagType: WasmTypeExtension {
     }
 }
 
+public class WasmLabelType: WasmTypeExtension {
+    // The parameter types for the label, meaning the types of the values that need to be provided
+    // when branching to this label. This is the list of result types for all wasm blocks excluding
+    // the loop for which the parameter types are the parameter types of the block. (This is caused
+    // by the branch instruction branching to the loop header and not the loop end.)
+    public let parameters: [ILType]
+
+    override func isEqual(to other: WasmTypeExtension) -> Bool {
+        guard let other = other as? WasmLabelType else { return false }
+        return self.parameters == other.parameters
+    }
+
+    override public func hash(into hasher: inout Hasher) {
+        hasher.combine(parameters)
+    }
+
+    init(_ parameters: [ILType]) {
+        self.parameters = parameters
+    }
+}
+
 public struct Limits: Hashable {
     var min: Int
     var max: Int?

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -953,15 +953,16 @@ final class WasmBeginBlock: WasmOperation {
     init(with signature: Signature) {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
-        super.init(inputTypes: parameterTypes, outputType: signature.outputType, innerOutputTypes: [.label] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
+        let labelTypes = signature.outputType != .nothing ? [signature.outputType] : []
+        super.init(inputTypes: parameterTypes, outputType: .nothing, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
     }
 }
 
 final class WasmEndBlock: WasmOperation {
     override var opcode: Opcode { .wasmEndBlock(self) }
 
-    init() {
-        super.init(attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction, .wasmBlock])
+    init(outputType: ILType) {
+        super.init(inputTypes: outputType != .nothing ? [outputType] : [], outputType: outputType, attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction, .wasmBlock])
     }
 }
 
@@ -972,12 +973,13 @@ final class WasmBeginIf: WasmOperation {
     init(with signature: Signature = [] => .nothing) {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
+        let labelTypes = signature.outputType != .nothing ? [signature.outputType] : []
         // TODO(mliedtke): Why does this set .isNotInputMutable? Try to remove it and see if the WasmLifter failure rate is affected.
 
         // Note that the condition is the last input! This is due to how lifting works for the wasm
         // value stack and that the condition is the first value to be removed from the stack, so
         // it needs to be the last one pushed to it.
-        super.init(inputTypes: parameterTypes + [.wasmi32], outputType: signature.outputType, innerOutputTypes: [.label] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext, .isNotInputMutable], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
+        super.init(inputTypes: parameterTypes + [.wasmi32], outputType: signature.outputType, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext, .isNotInputMutable], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
     }
 }
 
@@ -988,7 +990,8 @@ final class WasmBeginElse: WasmOperation {
     init(with signature: Signature = [] => .nothing) {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
-        super.init(outputType: signature.outputType, innerOutputTypes: [.label] + parameterTypes, attributes: [.isBlockStart, .isBlockEnd, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
+        let labelTypes = signature.outputType != .nothing ? [signature.outputType] : []
+        super.init(outputType: signature.outputType, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .isBlockEnd, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
     }
 }
 
@@ -1008,7 +1011,10 @@ final class WasmBeginLoop: WasmOperation {
     init(with signature: Signature) {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
-        super.init(outputType: signature.outputType, innerOutputTypes: [.label] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction])
+        // Note that different to all other blocks the loop's label parameters are the input types
+        // of the block, not the result types (because a branch to a loop label jumps to the
+        // beginning of the loop block instead of the end.)
+        super.init(outputType: .nothing, innerOutputTypes: [.label(parameterTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction])
     }
 }
 
@@ -1028,7 +1034,8 @@ final class WasmBeginTry: WasmOperation {
     init(with signature: Signature) {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
-        super.init(inputTypes: parameterTypes, outputType: signature.outputType, innerOutputTypes: [.label] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmTry])
+        let labelTypes = signature.outputType != .nothing ? [signature.outputType] : []
+        super.init(inputTypes: parameterTypes, outputType: signature.outputType, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmTry])
     }
 }
 
@@ -1101,7 +1108,8 @@ final class WasmBeginTryDelegate: WasmOperation {
     init(with signature: Signature) {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
-        super.init(inputTypes: parameterTypes, outputType: signature.outputType, innerOutputTypes: [.label] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [])
+        let labelTypes = signature.outputType != .nothing ? [signature.outputType] : []
+        super.init(inputTypes: parameterTypes, outputType: signature.outputType, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [])
     }
 }
 
@@ -1111,7 +1119,8 @@ final class WasmEndTryDelegate: WasmOperation {
     override var opcode: Opcode { .wasmEndTryDelegate(self) }
 
     init() {
-        super.init(inputTypes: [.label], attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction])
+        // Note that the actual block signature doesn't matter as the try-delegate "rethrows" the exception at that block level.
+        super.init(inputTypes: [.anyLabel], attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction])
     }
 }
 
@@ -1136,20 +1145,23 @@ final class WasmRethrow: WasmOperation {
 
 final class WasmBranch: WasmOperation {
     override var opcode: Opcode { .wasmBranch(self) }
+    let labelTypes: [ILType]
 
-    init() {
-        super.init(inputTypes: [.label], requiredContext: [.wasmFunction])
+    init(labelTypes: [ILType]) {
+        self.labelTypes = labelTypes
+        super.init(inputTypes: [.label(self.labelTypes)] + labelTypes, requiredContext: [.wasmFunction])
 
     }
 }
 
 final class WasmBranchIf: WasmOperation {
     override var opcode: Opcode { .wasmBranchIf(self) }
+    let labelTypes: [ILType]
 
-    init() {
-        super.init(inputTypes: [.label, .wasmi32], requiredContext: [.wasmFunction])
+    init(labelTypes: [ILType]) {
+        self.labelTypes = labelTypes
+        super.init(inputTypes: [.label(self.labelTypes)] + labelTypes + [.wasmi32], requiredContext: [.wasmFunction])
     }
-
 }
 
 // TODO: make this comprehensive, currently only works for locals, or assumes every thing it reassigns to is a local.