[wasm] Add parameters and labels to if and else blocks

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

Author: Liedtke

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -3261,6 +3261,14 @@ public class ProgramBuilder {
             b.emit(WasmEndIf())
         }
 
+        public func wasmBuildIfElse(_ condition: Variable, signature: Signature, args: [Variable], ifBody: (Variable, [Variable]) -> Void, elseBody: (Variable, [Variable]) -> Void) {
+            let beginBlock = b.emit(WasmBeginIf(with: signature), withInputs: args + [condition])
+            ifBody(beginBlock.innerOutput(0), Array(beginBlock.innerOutputs(1...)))
+            let elseBlock = b.emit(WasmBeginElse(with: signature))
+            elseBody(elseBlock.innerOutput(0), Array(elseBlock.innerOutputs(1...)))
+            b.emit(WasmEndIf())
+        }
+
         // 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))

Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift

@@ -268,6 +268,7 @@ public let codeGeneratorWeights = [
     // Control Flow Generators
     "WasmFunctionGenerator":                    30,
     "WasmIfElseGenerator":                      15,
+    "WasmIfElseWithSignatureGenerator":         10,
     "WasmReturnGenerator":                      15,
     "WasmBlockGenerator":                       8,
     "WasmBlockWithSignatureGenerator":          8,

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -633,6 +633,18 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         }
     },
 
+    RecursiveCodeGenerator("WasmIfElseWithSignatureGenerator", inContext: .wasmFunction, inputs: .required(.wasmi32)) { b, conditionVar 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!}
+        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)
+        }
+    },
+
     CodeGenerator("WasmSelectGenerator", inContext: .wasmFunction, inputs: .required(.wasmi32)) { b, condition in
         let function = b.currentWasmModule.currentWasmFunction
         let supportedTypes : ILType = .wasmi32 | .wasmi64 | .wasmf32 | .wasmf64 | .wasmExternRef

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1242,10 +1242,16 @@ extension Instruction: ProtobufConvertible {
                 $0.wasmBranch = Fuzzilli_Protobuf_WasmBranch()
             case .wasmBranchIf(_):
                 $0.wasmBranchIf = Fuzzilli_Protobuf_WasmBranchIf()
-            case .wasmBeginIf(_):
-                $0.wasmBeginIf = Fuzzilli_Protobuf_WasmBeginIf()
-            case .wasmBeginElse(_):
-                $0.wasmBeginElse = Fuzzilli_Protobuf_WasmBeginElse()
+            case .wasmBeginIf(let op):
+                $0.wasmBeginIf = Fuzzilli_Protobuf_WasmBeginIf.with {
+                    $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
+                    $0.returnType = ILTypeToWasmTypeEnum(op.signature.outputType)
+                }
+            case .wasmBeginElse(let op):
+                $0.wasmBeginElse = Fuzzilli_Protobuf_WasmBeginElse.with {
+                    $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
+                    $0.returnType = ILTypeToWasmTypeEnum(op.signature.outputType)
+                }
             case .wasmEndIf(_):
                 $0.wasmEndIf = Fuzzilli_Protobuf_WasmEndIf()
             case .wasmNop(_):
@@ -2031,10 +2037,16 @@ extension Instruction: ProtobufConvertible {
             op = WasmBranch()
         case .wasmBranchIf(_):
             op = WasmBranchIf()
-        case .wasmBeginIf(_):
-            op = WasmBeginIf()
-        case .wasmBeginElse(_):
-            op = WasmBeginElse()
+        case .wasmBeginIf(let p):
+            let parameters: [Parameter] = p.parameters.map({ param in
+                Parameter.plain(WasmTypeEnumToILType(param))
+            })
+            op = WasmBeginIf(with: parameters => WasmTypeEnumToILType(p.returnType))
+        case .wasmBeginElse(let p):
+            let parameters: [Parameter] = p.parameters.map({ param in
+                Parameter.plain(WasmTypeEnumToILType(param))
+            })
+            op = WasmBeginElse(with: parameters => WasmTypeEnumToILType(p.returnType))
         case .wasmEndIf(_):
             op = WasmEndIf()
         case .wasmNop(_):

Sources/Fuzzilli/FuzzIL/JSTyper.swift

@@ -111,6 +111,8 @@ public struct JSTyper: Analyzer {
                 activeWasmModuleDefinition!.methodSignatures.append((instr.output, signature))
             case .wasmBeginBlock(_),
                  .wasmBeginLoop(_),
+                 .wasmBeginIf(_),
+                 .wasmBeginElse(_),
                  .wasmBeginTry(_),
                  .wasmBeginCatch(_),
                  .wasmBeginTryDelegate(_):

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -967,17 +967,28 @@ final class WasmEndBlock: WasmOperation {
 
 final class WasmBeginIf: WasmOperation {
     override var opcode: Opcode { .wasmBeginIf(self) }
+    let signature: Signature
 
-    init() {
-        super.init(inputTypes: [.wasmi32], attributes: [.isBlockStart, .propagatesSurroundingContext, .isNotInputMutable], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
+    init(with signature: Signature = [] => .nothing) {
+        self.signature = signature
+        let parameterTypes = signature.parameters.convertPlainToILTypes()
+        // 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])
     }
 }
 
 final class WasmBeginElse: WasmOperation {
     override var opcode: Opcode { .wasmBeginElse(self) }
+    let signature: Signature
 
-    init() {
-        super.init(attributes: [.isBlockStart, .isBlockEnd, .propagatesSurroundingContext], requiredContext: [.wasmFunction], contextOpened: [.wasmBlock])
+    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])
     }
 }
 

Sources/Fuzzilli/Lifting/FuzzILLifter.swift

@@ -1047,13 +1047,18 @@ public class FuzzILLifter: Lifter {
         case .wasmBranchIf(_):
             w.emit("wasmBranchIf \(instr.input(1)), \(instr.input(0))")
 
-        case .wasmBeginIf(_):
-            w.emit("wasmBeginIf \(instr.input(0))")
+        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.increaseIndentionLevel()
 
-        case .wasmBeginElse(_):
+        case .wasmBeginElse(let op):
             w.decreaseIndentionLevel()
-            w.emit("wasmBeginElse")
+            w.emit("wasmBeginElse L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
             w.increaseIndentionLevel()
 
         case .wasmEndIf(_):

Sources/Fuzzilli/Lifting/WasmLifter.swift

@@ -866,11 +866,21 @@ public class WasmLifter {
 
         switch instr.op.opcode {
         case .wasmBeginBlock(let op):
-            // TODO(mliedtke): Repeat this for loops and if-else blocks.
+            // TODO(mliedtke): Repeat this for loops.
             registerSignature(op.signature)
             self.currentFunction!.labelBranchDepthMapping[instr.innerOutput(0)] = self.currentFunction!.variableAnalyzer.wasmBranchDepth
             // Needs typer analysis
             return true
+        case .wasmBeginIf(let op):
+            registerSignature(op.signature)
+            self.currentFunction!.labelBranchDepthMapping[instr.innerOutput(0)] = self.currentFunction!.variableAnalyzer.wasmBranchDepth
+            // Needs typer analysis
+            return true
+        case .wasmBeginElse(_):
+            // Note: We need to subtract one because the begin else block closes the if block before opening the else block!
+            self.currentFunction!.labelBranchDepthMapping[instr.innerOutput(0)] = self.currentFunction!.variableAnalyzer.wasmBranchDepth - 1
+            // Needs typer analysis
+            return true
         case .wasmBeginTry(let op):
             registerSignature(op.signature)
             self.currentFunction!.labelBranchDepthMapping[instr.innerOutput(0)] = self.currentFunction!.variableAnalyzer.wasmBranchDepth
@@ -954,20 +964,12 @@ public class WasmLifter {
 
             // TODO(mliedtke): Make this an attribute.
             // Instruction has to be a glue instruction now, maybe add an attribute to the instruction that it may have non-wasm inputs, i.e. inputs that do not have a local slot.
-            if instr.op is WasmLoadGlobal       ||
-               instr.op is WasmStoreGlobal      ||
-               instr.op is WasmJsCall           ||
-               instr.op is WasmMemoryStore      ||
-               instr.op is WasmMemoryLoad       ||
-               instr.op is WasmI64x2LoadSplat   ||
-               instr.op is WasmTableGet         ||
-               instr.op is WasmTableSet         ||
-               instr.op is WasmBeginCatch       ||
-               instr.op is WasmThrow            ||
-               instr.op is WasmRethrow          ||
-               instr.op is WasmBeginBlock       ||
-               instr.op is WasmBeginTry         ||
-               instr.op is WasmBeginTryDelegate {
+            if instr.op is WasmLoadGlobal || instr.op is WasmStoreGlobal || instr.op is WasmJsCall
+                || instr.op is WasmMemoryStore || instr.op is WasmMemoryLoad || instr.op is WasmTableGet
+                || instr.op is WasmTableSet || instr.op is WasmBeginCatch || instr.op is WasmThrow
+                || instr.op is WasmRethrow || instr.op is WasmBeginBlock || instr.op is WasmBeginTry
+                || instr.op is WasmI64x2LoadSplat || instr.op is WasmBeginTryDelegate
+                || instr.op is WasmBeginIf || instr.op is WasmBeginElse {
                 continue
             }
             fatalError("unreachable")
@@ -995,7 +997,8 @@ public class WasmLifter {
         }
 
         // TODO(mliedtke): Reuse this for handling parameters in loops, if-else, ...
-        if instr.op is WasmBeginCatch || instr.op is WasmBeginBlock || instr.op is WasmBeginTry || instr.op is WasmBeginTryDelegate {
+        if instr.op is WasmBeginCatch || instr.op is WasmBeginBlock || instr.op is WasmBeginTry
+            || instr.op is WasmBeginTryDelegate || instr.op is WasmBeginIf || instr.op is WasmBeginElse {
             // As the parameters are pushed "in order" to the stack, they need to be popped in reverse order.
             for innerOutput in instr.innerOutputs(1...).reversed() {
                 currentFunction!.spillLocal(forVariable: innerOutput)
@@ -1399,8 +1402,8 @@ public class WasmLifter {
         case .wasmBranchIf(_):
             let branchDepth = self.currentFunction!.variableAnalyzer.wasmBranchDepth - self.currentFunction!.labelBranchDepthMapping[wasmInstruction.input(0)]! - 1
             return Data([0x0D]) + Leb128.unsignedEncode(branchDepth)
-        case .wasmBeginIf(_):
-            return Data([0x04] + [0x40])
+        case .wasmBeginIf(let op):
+            return Data([0x04] + Leb128.unsignedEncode(signatureIndexMap[op.signature]!))
         case .wasmBeginElse(_):
             // 0x05 is the else block instruction.
             return Data([0x05])

Sources/Fuzzilli/Protobuf/operations.pb.swift

@@ -4140,6 +4140,10 @@ public struct Fuzzilli_Protobuf_WasmBeginIf: Sendable {
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
   // methods supported on all messages.
 
+  public var parameters: [Fuzzilli_Protobuf_WasmILType] = []
+
+  public var returnType: Fuzzilli_Protobuf_WasmILType = .consti32
+
   public var unknownFields = SwiftProtobuf.UnknownStorage()
 
   public init() {}
@@ -4150,6 +4154,10 @@ public struct Fuzzilli_Protobuf_WasmBeginElse: Sendable {
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
   // methods supported on all messages.
 
+  public var parameters: [Fuzzilli_Protobuf_WasmILType] = []
+
+  public var returnType: Fuzzilli_Protobuf_WasmILType = .consti32
+
   public var unknownFields = SwiftProtobuf.UnknownStorage()
 
   public init() {}
@@ -12105,37 +12113,75 @@ extension Fuzzilli_Protobuf_WasmReassign: SwiftProtobuf.Message, SwiftProtobuf._
 
 extension Fuzzilli_Protobuf_WasmBeginIf: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   public static let protoMessageName: String = _protobuf_package + ".WasmBeginIf"
-  public static let _protobuf_nameMap = SwiftProtobuf._NameMap()
+  public static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
+    1: .same(proto: "parameters"),
+    2: .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.decodeRepeatedEnumField(value: &self.parameters) }()
+      case 2: try { try decoder.decodeSingularEnumField(value: &self.returnType) }()
+      default: break
+      }
+    }
   }
 
   public func traverse<V: SwiftProtobuf.Visitor>(visitor: inout V) throws {
+    if !self.parameters.isEmpty {
+      try visitor.visitPackedEnumField(value: self.parameters, fieldNumber: 1)
+    }
+    if self.returnType != .consti32 {
+      try visitor.visitSingularEnumField(value: self.returnType, fieldNumber: 2)
+    }
     try unknownFields.traverse(visitor: &visitor)
   }
 
   public static func ==(lhs: Fuzzilli_Protobuf_WasmBeginIf, rhs: Fuzzilli_Protobuf_WasmBeginIf) -> Bool {
+    if lhs.parameters != rhs.parameters {return false}
+    if lhs.returnType != rhs.returnType {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }
 }
 
 extension Fuzzilli_Protobuf_WasmBeginElse: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   public static let protoMessageName: String = _protobuf_package + ".WasmBeginElse"
-  public static let _protobuf_nameMap = SwiftProtobuf._NameMap()
+  public static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
+    1: .same(proto: "parameters"),
+    2: .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.decodeRepeatedEnumField(value: &self.parameters) }()
+      case 2: try { try decoder.decodeSingularEnumField(value: &self.returnType) }()
+      default: break
+      }
+    }
   }
 
   public func traverse<V: SwiftProtobuf.Visitor>(visitor: inout V) throws {
+    if !self.parameters.isEmpty {
+      try visitor.visitPackedEnumField(value: self.parameters, fieldNumber: 1)
+    }
+    if self.returnType != .consti32 {
+      try visitor.visitSingularEnumField(value: self.returnType, fieldNumber: 2)
+    }
     try unknownFields.traverse(visitor: &visitor)
   }
 
   public static func ==(lhs: Fuzzilli_Protobuf_WasmBeginElse, rhs: Fuzzilli_Protobuf_WasmBeginElse) -> Bool {
+    if lhs.parameters != rhs.parameters {return false}
+    if lhs.returnType != rhs.returnType {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }

Sources/Fuzzilli/Protobuf/operations.proto

@@ -1196,9 +1196,13 @@ message WasmReassign {
 }
 
 message WasmBeginIf {
+    repeated WasmILType parameters = 1;
+    WasmILType returnType = 2;
 }
 
 message WasmBeginElse {
+    repeated WasmILType parameters = 1;
+    WasmILType returnType = 2;
 }
 
 message WasmEndIf {