[wasm] Add support for if else blocks with result

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

Author: Liedtke

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -3278,6 +3278,15 @@ public class ProgramBuilder {
             b.emit(WasmEndIf())
         }
 
+        @discardableResult
+        public func wasmBuildIfElseWithResult(_ condition: Variable, signature: Signature, args: [Variable], ifBody: (Variable, [Variable]) -> Variable, elseBody: (Variable, [Variable]) -> Variable) -> Variable{
+            let beginBlock = b.emit(WasmBeginIf(with: signature), withInputs: args + [condition])
+            let trueResult = ifBody(beginBlock.innerOutput(0), Array(beginBlock.innerOutputs(1...)))
+            let elseBlock = b.emit(WasmBeginElse(with: signature), withInputs: [trueResult])
+            let falseResult = elseBody(elseBlock.innerOutput(0), Array(elseBlock.innerOutputs(1...)))
+            return b.emit(WasmEndIf(outputType: signature.outputType), withInputs: [falseResult]).output
+        }
+
         // The first output of this block is a label variable, which is just there to explicitly mark control-flow and allow branches.
         public func wasmBuildLoop(with signature: Signature, body: (Variable, [Variable]) -> Void) {
             let instr = b.emit(WasmBeginLoop(with: signature))
@@ -3558,6 +3567,11 @@ public class ProgramBuilder {
         return params => returnType
     }
 
+    public func randomWasmBlockOutputType() -> ILType {
+        // TODO(mliedtke): The selection of types is in sync with ProgramBuilder::randomWasmSignature(). This should allow more types.
+        return chooseUniform(from: [.wasmi32, .wasmi64, .wasmf32, .wasmf64, .nothing])
+    }
+
     @discardableResult
     public func buildWasmModule(_ body: (WasmModule) -> ()) -> WasmModule {
         emit(BeginWasmModule())

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -560,8 +560,7 @@ 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): 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])
+        let outputType = b.randomWasmBlockOutputType()
         if outputType != .nothing {
             function.wasmBuildBlockWithResult(with: parameters => outputType, args: args) { label, args in
                 b.buildRecursive()
@@ -645,10 +644,21 @@ 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))}
-        function.wasmBuildIfElse(conditionVar, signature: parameters => .nothing, args: args) { label, args in
-            b.buildRecursive(block: 1, of: 2, n: 4)
-        } elseBody: { label, args in
-            b.buildRecursive(block: 2, of: 2, n: 4)
+        let outputType = b.randomWasmBlockOutputType()
+        if outputType != .nothing {
+            function.wasmBuildIfElseWithResult(conditionVar, signature: parameters => outputType, args: args) { label, args in
+                b.buildRecursive(block: 1, of: 2, n: 4)
+                return b.randomVariable(ofType: outputType) ?? function.generateRandomWasmVar(ofType: outputType)
+            } elseBody: { label, args in
+                b.buildRecursive(block: 2, of: 2, n: 4)
+                return b.randomVariable(ofType: outputType) ?? function.generateRandomWasmVar(ofType: outputType)
+            }
+        } else {
+            function.wasmBuildIfElse(conditionVar, signature: parameters => outputType, args: args) { label, args in
+                b.buildRecursive(block: 1, of: 2, n: 4)
+            } elseBody: { label, args in
+                b.buildRecursive(block: 2, of: 2, n: 4)
+            }
         }
     },
 

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1258,8 +1258,10 @@ extension Instruction: ProtobufConvertible {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
                     $0.returnType = ILTypeToWasmTypeEnum(op.signature.outputType)
                 }
-            case .wasmEndIf(_):
-                $0.wasmEndIf = Fuzzilli_Protobuf_WasmEndIf()
+            case .wasmEndIf(let op):
+                $0.wasmEndIf = Fuzzilli_Protobuf_WasmEndIf.with {
+                    $0.returnType = ILTypeToWasmTypeEnum(op.outputType)
+                }
             case .wasmNop(_):
                 fatalError("Should never be serialized")
             case .wasmUnreachable(_):
@@ -2053,8 +2055,8 @@ extension Instruction: ProtobufConvertible {
                 Parameter.plain(WasmTypeEnumToILType(param))
             })
             op = WasmBeginElse(with: parameters => WasmTypeEnumToILType(p.returnType))
-        case .wasmEndIf(_):
-            op = WasmEndIf()
+        case .wasmEndIf(let p):
+            op = WasmEndIf(outputType: WasmTypeEnumToILType(p.returnType))
         case .wasmNop(_):
             fatalError("Should never be deserialized!")
         case .wasmUnreachable(_):

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -979,7 +979,7 @@ final class WasmBeginIf: WasmOperation {
         // 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(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext, .isNotInputMutable], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
+        super.init(inputTypes: parameterTypes + [.wasmi32], outputType: .nothing, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext, .isNotInputMutable], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
     }
 }
 
@@ -991,15 +991,19 @@ final class WasmBeginElse: WasmOperation {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
         let labelTypes = signature.outputType != .nothing ? [signature.outputType] : []
-        super.init(outputType: signature.outputType, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .isBlockEnd, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
+        // The WasmBeginElse acts both as a block end for the true case and as a block start for the
+        // false case. As such, its input types are the results from the true block and its inner
+        // output types are the same as for the corresponding WasmBeginIf.
+        super.init(inputTypes: labelTypes, outputType: .nothing, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .isBlockEnd, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
     }
 }
 
 final class WasmEndIf: WasmOperation {
     override var opcode: Opcode { .wasmEndIf(self) }
 
-    init() {
-        super.init(attributes: [.isBlockEnd], requiredContext: [.wasmBlock, .wasmFunction])
+    init(outputType: ILType = .nothing) {
+        let inputTypes = outputType != .nothing ? [outputType] : []
+        super.init(inputTypes: inputTypes, outputType: outputType, attributes: [.isBlockEnd], requiredContext: [.wasmBlock, .wasmFunction])
     }
 }
 
@@ -1035,7 +1039,7 @@ final class WasmBeginTry: WasmOperation {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
         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])
+        super.init(inputTypes: parameterTypes, outputType: .nothing, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmTry])
     }
 }
 
@@ -1109,7 +1113,7 @@ final class WasmBeginTryDelegate: WasmOperation {
         self.signature = signature
         let parameterTypes = signature.parameters.convertPlainToILTypes()
         let labelTypes = signature.outputType != .nothing ? [signature.outputType] : []
-        super.init(inputTypes: parameterTypes, outputType: signature.outputType, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [])
+        super.init(inputTypes: parameterTypes, outputType: .nothing, innerOutputTypes: [.label(labelTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [])
     }
 }
 

Sources/Fuzzilli/Lifting/FuzzILLifter.swift

@@ -797,7 +797,7 @@ public class FuzzILLifter: Lifter {
 
         case .beginWasmFunction(let op):
             // TODO(cffsmith): do this properly?
-            w.emit("BeginWasmFunction [\(liftCallArguments(instr.innerOutputs))] (\(op.signature))")
+            w.emit("BeginWasmFunction (\(op.signature)) -> [\(liftCallArguments(instr.innerOutputs))]")
             w.increaseIndentionLevel()
 
         case .endWasmFunction:
@@ -982,24 +982,15 @@ public class FuzzILLifter: Lifter {
             }
 
         case .wasmBeginLoop(let op):
-            if instr.numOutputs > 0 {
-                w.emit("\(output()) <- WasmBeginLoop L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
-            } else {
-                w.emit("WasmBeginLoop L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
-            }
+            w.emit("WasmBeginLoop L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
             w.increaseIndentionLevel()
 
         case .wasmEndLoop(_):
             w.decreaseIndentionLevel()
             w.emit("WasmEndLoop")
 
         case .wasmBeginTry(let op):
-            if instr.numOutputs > 0 {
-                // TODO(cffsmith): Maybe lift labels as e.g. L7 or something like that?
-                w.emit("\(output()) <- WasmBeginTry L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
-            } else {
-                w.emit("WasmBeginTry L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
-            }
+            w.emit("WasmBeginTry L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
             w.increaseIndentionLevel()
 
         case .wasmBeginCatchAll(_):
@@ -1027,12 +1018,7 @@ public class FuzzILLifter: Lifter {
             w.emit("WasmRethrow \(instr.input(0))")
 
         case .wasmBeginTryDelegate(let op):
-            if instr.numOutputs > 0 {
-                // TODO(cffsmith): Maybe lift labels as e.g. L7 or something like that?
-                w.emit("\(output()) <- WasmBeginTryDelegate L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
-            } else {
-                w.emit("WasmBeginTryDelegate L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
-            }
+            w.emit("WasmBeginTryDelegate L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
             w.increaseIndentionLevel()
 
         case .wasmEndTryDelegate(_):
@@ -1050,21 +1036,25 @@ public class FuzzILLifter: Lifter {
 
         case .wasmBeginIf(let op):
             let inputs = instr.inputs.map(lift).joined(separator: ", ")
-            if instr.numOutputs > 0 {
-                w.emit("\(output()) <- wasmBeginIf L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature)) [\(inputs)]")
-            } else {
-                w.emit("wasmBeginIf L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature)) [\(inputs)]")
-            }
+            w.emit("wasmBeginIf (\(op.signature)) [\(inputs)] -> L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))]")
             w.increaseIndentionLevel()
 
-        case .wasmBeginElse(let op):
+        case .wasmBeginElse(_):
             w.decreaseIndentionLevel()
-            w.emit("wasmBeginElse L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
+            let inputs = instr.inputs.map(lift).joined(separator: ", ")
+            // Note that the signature is printed by the WasmBeginIf, so we skip it here for better
+            // readability.
+            w.emit("wasmBeginElse [\(inputs)] -> L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))]")
             w.increaseIndentionLevel()
 
-        case .wasmEndIf(_):
+        case .wasmEndIf(let op):
             w.decreaseIndentionLevel()
-            w.emit("wasmEndIf")
+            let inputs = instr.inputs.map(lift).joined(separator: ", ")
+            if op.numOutputs > 0 {
+                w.emit("\(output()) <- WasmEndIf \(inputs)")
+            } else {
+                w.emit("WasmEndIf \(inputs)")
+            }
 
         case .print:
             w.emit("Print \(input(0))")

Sources/Fuzzilli/Protobuf/operations.pb.swift

@@ -4174,6 +4174,8 @@ public struct Fuzzilli_Protobuf_WasmEndIf: Sendable {
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
   // methods supported on all messages.
 
+  public var returnType: Fuzzilli_Protobuf_WasmILType = .consti32
+
   public var unknownFields = SwiftProtobuf.UnknownStorage()
 
   public init() {}
@@ -12234,18 +12236,31 @@ extension Fuzzilli_Protobuf_WasmBeginElse: SwiftProtobuf.Message, SwiftProtobuf.
 
 extension Fuzzilli_Protobuf_WasmEndIf: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   public static let protoMessageName: String = _protobuf_package + ".WasmEndIf"
-  public static let _protobuf_nameMap = SwiftProtobuf._NameMap()
+  public static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
+    1: .same(proto: "returnType"),
+  ]
 
   public mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
-    // Load everything into unknown fields
-    while try decoder.nextFieldNumber() != nil {}
+    while let fieldNumber = try decoder.nextFieldNumber() {
+      // The use of inline closures is to circumvent an issue where the compiler
+      // allocates stack space for every case branch when no optimizations are
+      // enabled. https://github.com/apple/swift-protobuf/issues/1034
+      switch fieldNumber {
+      case 1: try { try decoder.decodeSingularEnumField(value: &self.returnType) }()
+      default: break
+      }
+    }
   }
 
   public func traverse<V: SwiftProtobuf.Visitor>(visitor: inout V) throws {
+    if self.returnType != .consti32 {
+      try visitor.visitSingularEnumField(value: self.returnType, fieldNumber: 1)
+    }
     try unknownFields.traverse(visitor: &visitor)
   }
 
   public static func ==(lhs: Fuzzilli_Protobuf_WasmEndIf, rhs: Fuzzilli_Protobuf_WasmEndIf) -> Bool {
+    if lhs.returnType != rhs.returnType {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }

Sources/Fuzzilli/Protobuf/operations.proto

@@ -1209,6 +1209,7 @@ message WasmBeginElse {
 }
 
 message WasmEndIf {
+    WasmILType returnType = 1;
 }
 
 message WasmNop {

Tests/FuzzilliTests/WasmTests.swift

@@ -860,6 +860,33 @@ class WasmFoundationTests: XCTestCase {
         testForOutput(program: jsProg, runner: runner, outputString: "100\n300\n200\n300\n")
     }
 
+    func testIfElseWithResult() throws {
+        let runner = try GetJavaScriptExecutorOrSkipTest()
+        let liveTestConfig = Configuration(logLevel: .error, enableInspection: true)
+        let fuzzer = makeMockFuzzer(config: liveTestConfig, environment: JavaScriptEnvironment())
+        let b = fuzzer.makeBuilder()
+        let outputFunc = b.createNamedVariable(forBuiltin: "output")
+        let module = b.buildWasmModule { wasmModule in
+            wasmModule.addWasmFunction(with: [.wasmi32] => .wasmi64) { function, args in
+                let blockResult = function.wasmBuildIfElseWithResult(args[0], signature: [] => .wasmi64, args: []) {label, args in
+                    return function.consti64(123)
+                } elseBody: {label, args in
+                    return function.consti64(321)
+                }
+                function.wasmReturn(blockResult)
+            }
+        }
+        let exports = module.loadExports()
+        let outTrue = b.callMethod(module.getExportedMethod(at: 0), on: exports, withArgs: [b.loadInt(1)])
+        b.callFunction(outputFunc, withArgs: [b.callMethod("toString", on: outTrue)])
+        let outFalse = b.callMethod(module.getExportedMethod(at: 0), on: exports, withArgs: [b.loadInt(0)])
+        b.callFunction(outputFunc, withArgs: [b.callMethod("toString", on: outFalse)])
+
+        let prog = b.finalize()
+        let jsProg = fuzzer.lifter.lift(prog)
+        testForOutput(program: jsProg, runner: runner, outputString: "123\n321\n")
+    }
+
     func testTryVoid() throws {
         let runner = try GetJavaScriptExecutorOrSkipTest()
         let liveTestConfig = Configuration(logLevel: .error, enableInspection: true)