[wasm] Change loops to use wasm-gc signatures

This allows using parameter types which are indexed types
(things like `(ref null 1)`).

Implementation:
- Each WasmLoop instruction now takes its signature as the first input.
- The static signature types are removed from the begin and endLoop.
- The loop code generator emits an "ad hoc" signature in order to emit
  signatures for which we already have corresponding inputs available.

Bug: 445356784
Change-Id: Ic58ab7d6a092a39de77c974142dd7f976786e8e1
Reviewed-on: https://chrome-internal-review.googlesource.com/c/v8/fuzzilli/+/8792956
Reviewed-by: Pawel Krawczyk <pawkra@google.com>
Commit-Queue: Matthias Liedtke <mliedtke@google.com>

Author: Liedtke

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -1345,6 +1345,10 @@ public class ProgramBuilder {
         jsTyper.setType(of: variable, to: variableType)
     }
 
+    public func getWasmTypeDef(for type: ILType) -> Variable {
+        jsTyper.getWasmTypeDef(for: type)
+    }
+
     /// This helper function converts parameter types into argument types, for example by "unrolling" rest parameters and handling optional parameters.
     private static func prepareArgumentTypes(forParameters params: ParameterList) -> [ILType] {
         var argumentTypes = [ILType]()
@@ -3998,18 +4002,25 @@ public class ProgramBuilder {
             return Array(b.emit(WasmEndIf(outputTypes: signature.outputTypes), withInputs: falseResults, types: signature.outputTypes).outputs)
         }
 
-        // 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: WasmSignature, body: (Variable, [Variable]) -> Void) {
-            let instr = b.emit(WasmBeginLoop(with: signature))
-            body(instr.innerOutput(0), Array(instr.innerOutputs(1...)))
-            b.emit(WasmEndLoop())
+        @discardableResult
+        public func wasmBuildLoop(with signature: WasmSignature, args: [Variable], body: (Variable, [Variable]) -> [Variable]) -> [Variable] {
+            let signatureDef = b.wasmDefineAdHocSignatureType(signature: signature)
+            return wasmBuildLoopImpl(signature: signature, signatureDef: signatureDef, args: args, body: body)
         }
 
         @discardableResult
-        public func wasmBuildLoop(with signature: WasmSignature, args: [Variable], body: (Variable, [Variable]) -> [Variable]) -> [Variable] {
-            let instr = b.emit(WasmBeginLoop(with: signature), withInputs: args, types: signature.parameterTypes)
+        public func wasmBuildLoop(with signatureDef: Variable, args: [Variable], body: (Variable, [Variable]) -> [Variable]) -> [Variable] {
+            let signature = b.type(of: signatureDef).wasmFunctionSignatureDefSignature
+            return wasmBuildLoopImpl(signature: signature, signatureDef: signatureDef, args: args, body: body)
+        }
+
+        private func wasmBuildLoopImpl(signature: WasmSignature, signatureDef: Variable, args: [Variable], body: (Variable, [Variable]) -> [Variable]) -> [Variable] {
+            let instr = b.emit(WasmBeginLoop(with: signature), withInputs: [signatureDef] + args,
+                types: [.wasmTypeDef()] + signature.parameterTypes)
             let fallthroughResults = body(instr.innerOutput(0), Array(instr.innerOutputs(1...)))
-            return Array(b.emit(WasmEndLoop(outputTypes: signature.outputTypes), withInputs: fallthroughResults, types: signature.outputTypes).outputs)
+            return Array(b.emit(WasmEndLoop(outputCount: signature.outputTypes.count),
+                withInputs: [signatureDef] + fallthroughResults,
+                types: [.wasmTypeDef()] + signature.outputTypes).outputs)
         }
 
         @discardableResult
@@ -4637,8 +4648,20 @@ public class ProgramBuilder {
 
     /// Like wasmDefineSignatureType but instead of within a type group this defines a signature
     /// type directly inside a wasm function.
+    /// This takes a signature with resolved index types for ease-of-use (meaning it accepts full
+    /// index reference types directly inside the signature).
     @discardableResult
-    func wasmDefineAdHocSignatureType(signature: WasmSignature, indexTypes: [Variable]) -> Variable {
+    func wasmDefineAdHocSignatureType(signature: WasmSignature) -> Variable {
+        let indexTypes = (signature.parameterTypes + signature.outputTypes)
+                .filter {$0.Is(.anyIndexRef)}
+                .map(getWasmTypeDef)
+        let cleanIndexTypes = {(type: ILType) -> ILType in
+            type.Is(.anyIndexRef)
+                ? .wasmRef(.Index(), nullability: type.wasmReferenceType!.nullability)
+                : type
+        }
+        let signature = signature.parameterTypes.map(cleanIndexTypes)
+            => signature.outputTypes.map(cleanIndexTypes)
         return emit(WasmDefineAdHocSignatureType(signature: signature), withInputs: indexTypes).output
     }
 
@@ -4854,16 +4877,13 @@ public class ProgramBuilder {
             activeWasmModule!.blockSignatures.push(op.signature)
         case .wasmBeginBlock(let op):
             activeWasmModule!.blockSignatures.push(op.signature)
-        case .wasmBeginLoop(let op):
-            activeWasmModule!.blockSignatures.push(op.signature)
         case .wasmBeginTry(let op):
             activeWasmModule!.blockSignatures.push(op.signature)
         case .wasmBeginTryDelegate(let op):
             activeWasmModule!.blockSignatures.push(op.signature)
         case .wasmBeginTryTable(let op):
             activeWasmModule!.blockSignatures.push(op.signature)
         case .wasmEndIf(_),
-             .wasmEndLoop(_),
              .wasmEndTry(_),
              .wasmEndTryDelegate(_),
              .wasmEndTryTable(_),

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -1261,7 +1261,9 @@ public let WasmCodeGenerators: [CodeGenerator] = [
                 let function = b.currentWasmModule.currentWasmFunction
                 let loopCtr = function.consti32(10)
                 b.runtimeData.push("loopCounter", loopCtr)
-                b.emit(WasmBeginLoop(with: [] => []))
+                let signature = b.wasmDefineAdHocSignatureType(signature: [] => [])
+                b.runtimeData.push("loopSignature", signature)
+                b.emit(WasmBeginLoop(with: [] => []), withInputs: [signature])
                 // Increase loop counter.
                 let result = function.wasmi32BinOp(
                 loopCtr, function.consti32(1), binOpKind: .Sub)
@@ -1275,38 +1277,59 @@ public let WasmCodeGenerators: [CodeGenerator] = [
                 let loopCtr = b.runtimeData.pop("loopCounter")
                 // Backedge of loop, we continue if it is not equal to zero.
                 let isNotZero = function.wasmi32CompareOp(loopCtr, function.consti32(0), using: .Ne)
-                b.emit(WasmEndLoop(outputTypes: []))
+                b.emit(WasmEndLoop(outputCount: 0), withInputs: [b.runtimeData.pop("loopSignature")])
             },
         ]),
 
-    CodeGenerator("WasmLoopWithSignatureGenerator", inContext: .single(.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 = b.randomWasmBlockArguments(upTo: 5)
-        let randomArgTypes = randomArgs.map { b.type(of: $0) }
-        let args = [function.consti32(0)] + randomArgs
-        let parameters = args.map(b.type)
-        let outputTypes = b.randomWasmBlockOutputTypes(upTo: 5)
-        // Note that due to the do-while style implementation, the actual iteration count is at least 1.
-        let iterationCount = Int32.random(in: 0...16)
+    CodeGenerator("WasmLoopWithSignatureGenerator", [
+        GeneratorStub("WasmBeginLoopWithSignatureGenerator", inContext: .single(.wasmFunction),
+                provides: [.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 = b.randomWasmBlockArguments(upTo: 5)
+            let randomArgTypes = randomArgs.map { b.type(of: $0) }
+            let args = [function.consti32(0)] + randomArgs
+            let parameters = args.map(b.type)
+            // TODO(mliedtke): Also allow index types in the output types.
+            let outputTypes = b.randomWasmBlockOutputTypes(upTo: 5)
+            // Calculate index types for the dynamically created signature.
+            let indexTypes = (parameters + outputTypes)
+                .filter {$0.Is(.anyIndexRef)}
+                .map(b.getWasmTypeDef)
+
+            // TODO(mliedtke): We need to cleanup the index types here!
+            let signature = b.wasmDefineAdHocSignatureType(signature: parameters => outputTypes)
+            let loopBegin = b.emit(WasmBeginLoop(parameterCount: parameters.count),
+                withInputs: [signature] + args)
+            let loopCounter = loopBegin.innerOutput(1)
+            assert(b.type(of: loopCounter).Is(.wasmi32))
+            b.runtimeData.push("loopCounter", loopCounter)
+            b.runtimeData.push("loopSignature", signature)
+            b.runtimeData.push("loopLabel", loopBegin.innerOutput(0))
+        },
+        GeneratorStub("WasmEndLoopWithSignatureGenerator", inContext: .single(.wasmFunction),
+                provides: [.wasmFunction]) { b in
+            let signature = b.runtimeData.pop("loopSignature")
+            let function = b.currentWasmModule.currentWasmFunction
 
-        function.wasmBuildLoop(with: parameters => outputTypes, args: args) {
-            label, loopArgs in
-            b.buildRecursive(n: defaultCodeGenerationAmount)
+            // Note that due to the do-while style implementation, the actual iteration count is at
+            // least 1.
+            let iterationCount = Int32.random(in: 0...16)
             let loopCtr = function.wasmi32BinOp(
-                args[0], function.consti32(1), binOpKind: .Add)
+                b.runtimeData.pop("loopCounter"), 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)! }
+            let wasmSignature = b.type(of: signature).wasmFunctionSignatureDefSignature
+            let backedgeArgs = [loopCtr] + wasmSignature.parameterTypes.dropFirst()
+                .map {b.randomVariable(ofType: $0)!}
             function.wasmBranchIf(
-                condition, to: label, args: backedgeArgs,
+                condition, to: b.runtimeData.pop("loopLabel"), args: backedgeArgs,
                 hint: b.randomWasmBranchHint())
-            return outputTypes.map(function.findOrGenerateWasmVar)
-        }
-    },
+            let outputs = wasmSignature.outputTypes.map(function.findOrGenerateWasmVar)
+            b.emit(WasmEndLoop(outputCount: outputs.count), withInputs: [signature] + outputs)
+        },
+    ]),
 
     // TODO Turn this into a multi-part Generator
     CodeGenerator("WasmLegacyTryCatchGenerator", inContext: .single(.wasmFunction)) {

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1395,12 +1395,11 @@ extension Instruction: ProtobufConvertible {
                 }
             case .wasmBeginLoop(let op):
                 $0.wasmBeginLoop = Fuzzilli_Protobuf_WasmBeginLoop.with {
-                    $0.parameterTypes = op.signature.parameterTypes.map(ILTypeToWasmTypeEnum)
-                    $0.outputTypes = op.signature.outputTypes.map(ILTypeToWasmTypeEnum)
+                    $0.parameterCount = Int32(op.parameterCount)
                 }
             case .wasmEndLoop(let op):
                 $0.wasmEndLoop = Fuzzilli_Protobuf_WasmEndLoop.with {
-                    $0.outputTypes = op.outputTypes.map(ILTypeToWasmTypeEnum)
+                    $0.outputCount = Int32(op.numOutputs)
                 }
             case .wasmBeginTryTable(let op):
                 $0.wasmBeginTryTable = Fuzzilli_Protobuf_WasmBeginTryTable.with {
@@ -2456,11 +2455,9 @@ extension Instruction: ProtobufConvertible {
         case .wasmEndBlock(let p):
             op = WasmEndBlock(outputTypes: p.outputTypes.map(WasmTypeEnumToILType))
         case .wasmBeginLoop(let p):
-            let parameters = p.parameterTypes.map(WasmTypeEnumToILType)
-            let outputs = p.outputTypes.map(WasmTypeEnumToILType)
-            op = WasmBeginLoop(with: parameters => outputs)
+            op = WasmBeginLoop(parameterCount: Int(p.parameterCount))
         case .wasmEndLoop(let p):
-            op = WasmEndLoop(outputTypes: p.outputTypes.map(WasmTypeEnumToILType))
+            op = WasmEndLoop(outputCount: Int(p.outputCount))
         case .wasmBeginTryTable(let p):
             let parameters = p.parameterTypes.map(WasmTypeEnumToILType)
             let outputs = p.outputTypes.map(WasmTypeEnumToILType)

Sources/Fuzzilli/FuzzIL/JSTyper.swift

@@ -38,6 +38,8 @@ public struct JSTyper: Analyzer {
     private var typeGroupDependencies: [Int:Set<Int>] = [:]
     private var selfReferences: [Variable: [(inout JSTyper, Variable?) -> ()]] = [:]
     private var isWithinTypeGroup = false
+    // Tracks the type definition variable for each Wasm index type.
+    private var wasmTypeDefMap = [WasmTypeDescription:Variable]()
 
     // Tracks the active function definitions and contains the instruction that started the function.
     private var activeFunctionDefinitions = Stack<Instruction>()
@@ -401,6 +403,7 @@ public struct JSTyper: Analyzer {
         state.reset()
         signatures.removeAll()
         typeGroups.removeAll()
+        wasmTypeDefMap.removeAll()
         defUseAnalyzer = DefUseAnalyzer()
         isWithinTypeGroup = false
         dynamicObjectGroupManager = ObjectGroupManager()
@@ -432,6 +435,20 @@ public struct JSTyper: Analyzer {
         }
     }
 
+    mutating func registerWasmTypeDef(_ def: Variable) {
+        let desc = type(of: def).wasmTypeDefinition!.description!
+        assert(wasmTypeDefMap[desc] == nil, "duplicate type description")
+        wasmTypeDefMap[desc] = def
+    }
+
+    func getWasmTypeDef(for type: ILType) -> Variable {
+        assert(type.isWasmReferenceType)
+        guard case .Index(let desc) = type.wasmReferenceType!.kind else {
+            fatalError("\(type) is not an index type")
+        }
+        return wasmTypeDefMap[desc.get()!]!
+    }
+
     mutating func addSignatureType(def: Variable, signature: WasmSignature, inputs: ArraySlice<Variable>) {
         assert(isWithinTypeGroup)
         var inputs = inputs.makeIterator()
@@ -445,7 +462,7 @@ public struct JSTyper: Analyzer {
             if paramType.requiredInputCount() == 0 {
                 return paramType
             }
-            assert(paramType.Is(.wasmRef(.Index(), nullability: true)))
+            assert(paramType.Is(.anyIndexRef))
             let typeDef = inputs.next()!
             let elementDesc = type(of: typeDef).wasmTypeDefinition!.description!
             if elementDesc == .selfReference {
@@ -809,16 +826,18 @@ public struct JSTyper: Analyzer {
                 wasmTypeBeginBlock(instr, op.signature)
             case .wasmEndIf(let op):
                 wasmTypeEndBlock(instr, op.outputTypes)
-            case .wasmBeginLoop(let op):
+            case .wasmBeginLoop(_):
                 // 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.)
-                setType(of: instr.innerOutputs.first!, to: .label(op.signature.parameterTypes))
-                for (innerOutput, paramType) in zip(instr.innerOutputs.dropFirst(), op.signature.parameterTypes) {
+                let signature = type(of: instr.input(0)).wasmFunctionSignatureDefSignature
+                setType(of: instr.innerOutputs.first!, to: .label(signature.parameterTypes))
+                for (innerOutput, paramType) in zip(instr.innerOutputs.dropFirst(), signature.parameterTypes) {
                     setType(of: innerOutput, to: paramType)
                 }
-            case .wasmEndLoop(let op):
-                wasmTypeEndBlock(instr, op.outputTypes)
+            case .wasmEndLoop(_):
+                let signature = type(of: instr.input(0)).wasmFunctionSignatureDefSignature
+                wasmTypeEndBlock(instr, signature.outputTypes)
             case .wasmBeginTryTable(let op):
                 wasmTypeBeginBlock(instr, op.signature)
                 instr.inputs.forEach { input in
@@ -903,6 +922,7 @@ public struct JSTyper: Analyzer {
                 startTypeGroup()
                 addSignatureType(def: instr.output, signature: op.signature, inputs: instr.inputs)
                 finishTypeGroup()
+                registerWasmTypeDef(instr.output)
             default:
                 if instr.numInnerOutputs + instr.numOutputs != 0 {
                     fatalError("Missing typing of outputs for \(instr.op.opcode)")
@@ -1886,6 +1906,9 @@ public struct JSTyper: Analyzer {
                 set(output, state.type(of: input))
             }
             finishTypeGroup()
+            for output in instr.outputs {
+                registerWasmTypeDef(output)
+            }
 
         case .wasmDefineSignatureType(let op):
             addSignatureType(def: instr.output, signature: op.signature, inputs: instr.inputs)

Sources/Fuzzilli/FuzzIL/TypeSystem.swift

@@ -291,6 +291,7 @@ public struct ILType: Hashable {
     public static let wasmNumericalPrimitive = .wasmi32 | .wasmi64 | .wasmf32 | .wasmf64
 
     public static let anyNonNullableIndexRef = wasmRef(.Index(), nullability: false)
+    public static let anyIndexRef = wasmRef(.Index(), nullability: true)
 
     //
     // Type testing
@@ -596,6 +597,14 @@ public struct ILType: Hashable {
         return (wasmType as! WasmFunctionDefinition).signature
     }
 
+    public var wasmFunctionSignatureDefSignature: WasmSignature {
+        let desc = (wasmType as? WasmTypeDefinition)?.description
+        guard let desc = desc as? WasmSignatureTypeDescription else {
+            fatalError("\(self) is not a Wasm signature type defintion")
+        }
+        return desc.signature
+    }
+
     public var isWasmDefaultable: Bool {
         let isPacked = Is(.wasmPackedI8) || Is(.wasmPackedI16)
         return isPacked

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -1315,21 +1315,26 @@ final class WasmEndIf: WasmOperation {
 
 final class WasmBeginLoop: WasmOperation {
     override var opcode: Opcode { .wasmBeginLoop(self) }
-    let signature: WasmSignature
 
-    init(with signature: WasmSignature) {
-        self.signature = signature
-        super.init(numInputs: signature.parameterTypes.count, numInnerOutputs: 1 + signature.parameterTypes.count, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction])
+    init(parameterCount: Int) {
+        // Inputs: the signature + the inputs to the loop.
+        // inner outputs: The loop label + the arguments of the loop.
+        super.init(numInputs: 1 + parameterCount, numInnerOutputs: 1 + parameterCount, attributes: [.isBlockStart, .propagatesSurroundingContext], requiredContext: [.wasmFunction])
+    }
+
+    convenience init(with signature: WasmSignature) {
+        self.init(parameterCount: signature.parameterTypes.count)
     }
+
+    var parameterCount: Int {numInputs - 1}
 }
 
 final class WasmEndLoop: WasmOperation {
     override var opcode: Opcode { .wasmEndLoop(self) }
-    let outputTypes: [ILType]
 
-    init(outputTypes: [ILType] = []) {
-        self.outputTypes = outputTypes
-        super.init(numInputs: outputTypes.count, numOutputs: outputTypes.count, attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction])
+    init(outputCount: Int) {
+        // Inputs: the signature + the outputs of the loop.
+        super.init(numInputs: 1 + outputCount, numOutputs: outputCount, attributes: [.isBlockEnd, .resumesSurroundingContext], requiredContext: [.wasmFunction])
     }
 }
 

Sources/Fuzzilli/Lifting/FuzzILLifter.swift

@@ -1097,9 +1097,9 @@ public class FuzzILLifter: Lifter {
                 w.emit("WasmEndBlock \(inputs)")
             }
 
-        case .wasmBeginLoop(let op):
+        case .wasmBeginLoop(_):
             let inputs = instr.inputs.map(lift).joined(separator: ", ")
-            w.emit("WasmBeginLoop (\(op.signature)) [\(inputs)] -> L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))]")
+            w.emit("WasmBeginLoop [\(inputs)] -> L:\(instr.innerOutput(0)) [\(liftCallArguments(instr.innerOutputs(1...)))]")
             w.increaseIndentionLevel()
 
         case .wasmEndLoop(let op):