[wasm] Add support for loops with parameters and result value

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

Author: Liedtke

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -3294,6 +3294,13 @@ public class ProgramBuilder {
             b.emit(WasmEndLoop())
         }
 
+        @discardableResult
+        public func wasmBuildLoop(with signature: Signature, args: [Variable], body: (Variable, [Variable]) -> Variable) -> Variable {
+            let instr = b.emit(WasmBeginLoop(with: signature), withInputs: args)
+            let fallthroughResult = body(instr.innerOutput(0), Array(instr.innerOutputs(1...)))
+            return b.emit(WasmEndLoop(outputType: signature.outputType), withInputs: [fallthroughResult]).output
+        }
+
         public func wasmBuildLegacyTry(with signature: Signature, args: [Variable], body: (Variable, [Variable]) -> Void, catchAllBody: (() -> Void)? = nil) {
             assert(signature.parameters.count == args.count)
             let instr = b.emit(WasmBeginTry(with: signature), withInputs: args)
@@ -3567,9 +3574,10 @@ public class ProgramBuilder {
         return params => returnType
     }
 
-    public func randomWasmBlockOutputType() -> ILType {
+    public func randomWasmBlockOutputType(allowVoid: Bool = true) -> 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])
+        let possibleTypes: [ILType] = [.wasmi32, .wasmi64, .wasmf32, .wasmf64]
+        return chooseUniform(from: allowVoid ? possibleTypes + [.nothing] : possibleTypes)
     }
 
     @discardableResult

Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift

@@ -273,6 +273,7 @@ public let codeGeneratorWeights = [
     "WasmBlockGenerator":                       8,
     "WasmBlockWithSignatureGenerator":          8,
     "WasmLoopGenerator":                        8,
+    "WasmLoopWithSignatureGenerator":           8,
     "WasmLegacyTryGenerator":                   8,
     "WasmLegacyCatchGenerator":                 8,
     "WasmLegacyTryDelegateGenerator":           8,

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -590,6 +590,28 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         }
     },
 
+    RecursiveCodeGenerator("WasmLoopWithSignatureGenerator", inContext: .wasmFunction) { b in
+        let function = b.currentWasmModule.currentWasmFunction
+        // Count upwards here to make it slightly more different from the other loop generator.
+        // Also, instead of using reassign, this generator uses the signature to pass and update the loop counter.
+        let randomArgs = (0..<5).map {_ in b.findVariable {b.type(of: $0).Is(.wasmPrimitive)}}.filter {$0 != nil}.map {$0!}
+        let randomArgTypes = randomArgs.map{b.type(of: $0)}
+        let args = [function.consti32(0)] + randomArgs
+        let parameters = args.map {arg in Parameter.plain(b.type(of: arg))}
+        let outputType = b.randomWasmBlockOutputType(allowVoid: false)
+        // Note that due to the do-while style implementation, the actual iteration count is at least 1.
+        let iterationCount = Int32.random(in: 0...16)
+
+        function.wasmBuildLoop(with: parameters => outputType, args: args) { label, loopArgs in
+            b.buildRecursive()
+            let loopCtr = function.wasmi32BinOp(args[0], function.consti32(1), binOpKind: .Add)
+            let condition = function.wasmi32CompareOp(loopCtr, function.consti32(iterationCount), using: .Lt_s)
+            let backedgeArgs = [loopCtr] + randomArgTypes.map{b.randomVariable(ofType: $0)!}
+            function.wasmBranchIf(condition, to: label, args: backedgeArgs)
+            return b.randomVariable(ofType: outputType) ?? function.generateRandomWasmVar(ofType: outputType)
+        }
+    },
+
     RecursiveCodeGenerator("WasmLegacyTryGenerator", inContext: .wasmFunction) { b in
         let function = b.currentWasmModule.currentWasmFunction
         // Choose a few random wasm values as arguments if available.

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1203,8 +1203,10 @@ extension Instruction: ProtobufConvertible {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
                     $0.returnType = ILTypeToWasmTypeEnum(op.signature.outputType)
                 }
-            case .wasmEndLoop(_):
-                $0.wasmEndLoop = Fuzzilli_Protobuf_WasmEndLoop()
+            case .wasmEndLoop(let op):
+                $0.wasmEndLoop = Fuzzilli_Protobuf_WasmEndLoop.with {
+                    $0.returnType = ILTypeToWasmTypeEnum(op.outputType)
+                }
             case .wasmBeginTry(let op):
                 $0.wasmBeginTry = Fuzzilli_Protobuf_WasmBeginTry.with {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
@@ -2006,8 +2008,8 @@ extension Instruction: ProtobufConvertible {
                 Parameter.plain(WasmTypeEnumToILType(param))
             })
             op = WasmBeginLoop(with: parameters => WasmTypeEnumToILType(p.returnType))
-        case .wasmEndLoop(_):
-            op = WasmEndLoop()
+        case .wasmEndLoop(let p):
+            op = WasmEndLoop(outputType: WasmTypeEnumToILType(p.returnType))
         case .wasmBeginTry(let p):
             let parameters: [Parameter] = p.parameters.map({ param in
                 Parameter.plain(WasmTypeEnumToILType(param))

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -1018,15 +1018,16 @@ final class WasmBeginLoop: WasmOperation {
         // 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])
+        super.init(inputTypes: parameterTypes, outputType: .nothing, innerOutputTypes: [.label(parameterTypes)] + parameterTypes, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction])
     }
 }
 
 final class WasmEndLoop: WasmOperation {
     override var opcode: Opcode { .wasmEndLoop(self) }
 
-    init() {
-        super.init(attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction])
+    init(outputType: ILType = .nothing) {
+        let inputTypes = outputType != .nothing ? [outputType] : []
+        super.init(inputTypes: inputTypes, outputType: outputType, attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction])
     }
 }
 

Sources/Fuzzilli/Lifting/FuzzILLifter.swift

@@ -982,12 +982,18 @@ public class FuzzILLifter: Lifter {
             }
 
         case .wasmBeginLoop(let op):
-            w.emit("WasmBeginLoop L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")
+            let inputs = instr.inputs.map(lift).joined(separator: ", ")
+            w.emit("WasmBeginLoop (\(op.signature)) [\(inputs)] -> L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))]")
             w.increaseIndentionLevel()
 
-        case .wasmEndLoop(_):
+        case .wasmEndLoop(let op):
             w.decreaseIndentionLevel()
-            w.emit("WasmEndLoop")
+            let inputs = instr.inputs.map(lift).joined(separator: ", ")
+            if op.numOutputs > 0 {
+                w.emit("\(output()) <- WasmEndLoop \(inputs)")
+            } else {
+                w.emit("WasmEndLoop \(inputs)")
+            }
 
         case .wasmBeginTry(let op):
             w.emit("WasmBeginTry L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))] (\(op.signature))")

Sources/Fuzzilli/Lifting/WasmLifter.swift

@@ -891,7 +891,8 @@ public class WasmLifter {
             self.currentFunction!.labelBranchDepthMapping[instr.innerOutput(0)] = self.currentFunction!.variableAnalyzer.wasmBranchDepth
             // Needs typer analysis
             return true
-        case .wasmBeginLoop(_):
+        case .wasmBeginLoop(let op):
+            registerSignature(op.signature)
             self.currentFunction!.labelBranchDepthMapping[instr.innerOutput(0)] = self.currentFunction!.variableAnalyzer.wasmBranchDepth
             // Needs typer analysis
             return true
@@ -996,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 || instr.op is WasmBeginIf || instr.op is WasmBeginElse {
+            || instr.op is WasmBeginTryDelegate || instr.op is WasmBeginIf || instr.op is WasmBeginElse
+            || instr.op is WasmBeginLoop {
             // 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)
@@ -1363,9 +1365,8 @@ public class WasmLifter {
             // A Block can "produce" (push) an item on the value stack, just like a function. Similarly, a block can also have parameters.
             // Ref: https://webassembly.github.io/spec/core/binary/instructions.html#binary-blocktype
             return Data([0x02] + Leb128.unsignedEncode(signatureIndexMap[op.signature]!))
-        case .wasmBeginLoop(_):
-            // 0x03 is the loop instruction and 0x40 is the empty block type, just like in .wasmBeginBlock
-            return Data([0x03] + [0x40])
+        case .wasmBeginLoop(let op):
+            return Data([0x03] + Leb128.unsignedEncode(signatureIndexMap[op.signature]!))
         case .wasmBeginTry(let op):
             return Data([0x06] + Leb128.unsignedEncode(signatureIndexMap[op.signature]!))
         case .wasmBeginTryDelegate(let op):

Sources/Fuzzilli/Protobuf/operations.pb.swift

@@ -3982,6 +3982,8 @@ public struct Fuzzilli_Protobuf_WasmEndLoop: 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() {}
@@ -11759,18 +11761,31 @@ extension Fuzzilli_Protobuf_WasmBeginLoop: SwiftProtobuf.Message, SwiftProtobuf.
 
 extension Fuzzilli_Protobuf_WasmEndLoop: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   public static let protoMessageName: String = _protobuf_package + ".WasmEndLoop"
-  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_WasmEndLoop, rhs: Fuzzilli_Protobuf_WasmEndLoop) -> Bool {
+    if lhs.returnType != rhs.returnType {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }

Sources/Fuzzilli/Protobuf/operations.proto

@@ -1144,6 +1144,7 @@ message WasmBeginLoop {
 }
 
 message WasmEndLoop {
+  WasmILType returnType = 1;
 }
 
 message WasmBeginTry {

Tests/FuzzilliTests/WasmTests.swift

@@ -753,6 +753,33 @@ class WasmFoundationTests: XCTestCase {
         testForOutput(program: jsProg, runner: runner, outputString: "1368\n")
     }
 
+    func testLoopWithParametersAndResult() 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, .wasmi32] => .wasmi32) { function, args in
+                let loopResult = function.wasmBuildLoop(with: [.wasmi32, .wasmi32] => .wasmi32, args: args) { loopLabel, loopArgs in
+                    let incFirst = function.wasmi32BinOp(loopArgs[0], function.consti32(1), binOpKind: .Add)
+                    let incSecond = function.wasmi32BinOp(loopArgs[1], function.consti32(2), binOpKind: .Add)
+                    let condition = function.wasmi32CompareOp(incFirst, incSecond, using: .Gt_s)
+                    function.wasmBranchIf(condition, to: loopLabel, args: [incFirst, incSecond])
+                    return incFirst
+                }
+                function.wasmReturn(loopResult)
+            }
+        }
+        let exports = module.loadExports()
+        let wasmOut = b.callMethod(module.getExportedMethod(at: 0), on: exports, withArgs: [b.loadInt(10), b.loadInt(0)])
+        b.callFunction(outputFunc, withArgs: [b.callMethod("toString", on: wasmOut)])
+
+        let prog = b.finalize()
+        let jsProg = fuzzer.lifter.lift(prog)
+        testForOutput(program: jsProg, runner: runner, outputString: "20\n")
+    }
+
     func testIfs() throws {
         let runner = try GetJavaScriptExecutorOrSkipTest()
         let liveTestConfig = Configuration(logLevel: .error, enableInspection: true)