Add support for defining and droppining element segments.

Prepare code for using them in initialize and copy bulk instructions.

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

Author: pkk33

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -3400,6 +3400,28 @@ public class ProgramBuilder {
             b.emit(WasmDropDataSegment(), withInputs: [dataSegment], types: [.wasmDataSegment()])
         }
 
+        public func wasmDropElementSegment(elementSegment: Variable) {
+            b.emit(WasmDropElementSegment(), withInputs: [elementSegment], types: [.wasmElementSegment()])
+        }
+
+        public func wasmTableInit(elementSegment: Variable, table: Variable, tableOffset: Variable, elementSegmentOffset: Variable, nrOfElementsToUpdate: Variable) {
+            // TODO: b/427115604 - assert that table.elemType IS_SUBTYPE_OF elementSegment.elemType (depending on refactor outcome).
+            let addrType = b.type(of: table).wasmTableType!.isTable64 ? ILType.wasmi64 : ILType.wasmi32
+            b.emit(WasmTableInit(), withInputs: [elementSegment, table, tableOffset, elementSegmentOffset, nrOfElementsToUpdate],
+                types: [.wasmElementSegment(), .object(ofGroup: "WasmTable"), addrType, addrType, addrType])
+        }
+
+        public func wasmTableCopy(dstTable: Variable, srcTable: Variable, dstOffset: Variable, srcOffset: Variable, count: Variable) {
+            // TODO: b/427115604 - assert that srcTable.elemType IS_SUBTYPE_OF dstTable.elemType (depending on refactor outcome).
+            let dstTableType = b.type(of: dstTable).wasmTableType!
+            let srcTableType = b.type(of: srcTable).wasmTableType!
+            assert(dstTableType.isTable64 == srcTableType.isTable64)
+
+            let addrType = dstTableType.isTable64 ? ILType.wasmi64 : ILType.wasmi32
+            b.emit(WasmTableCopy(), withInputs: [dstTable, srcTable, dstOffset, srcOffset, count],
+                types: [.object(ofGroup: "WasmTable"), .object(ofGroup: "WasmTable"), addrType, addrType, addrType])
+        }
+
         public func wasmReassign(variable: Variable, to: Variable) {
             assert(b.type(of: variable) == b.type(of: to))
             b.emit(WasmReassign(variableType: b.type(of: variable)), withInputs: [variable, to])
@@ -3922,6 +3944,12 @@ public class ProgramBuilder {
                 withInputs: definedEntryValues, types: inputTypes).output
         }
 
+        @discardableResult
+        public func addElementSegment(elementsType: ILType,  elements: [Variable]) -> Variable {
+            let inputTypes = Array(repeating: elementsType, count: elements.count)
+            return b.emit(WasmDefineElementSegment(size: UInt32(elements.count)), withInputs: elements, types: inputTypes).output
+        }
+
         // This result can be ignored right now, as we can only define one memory per module
         // Also this should be tracked like a global / table.
         @discardableResult
@@ -4292,7 +4320,8 @@ public class ProgramBuilder {
         case .wasmDefineGlobal(_),
              .wasmDefineTable(_),
              .wasmDefineMemory(_),
-             .wasmDefineDataSegment(_):
+             .wasmDefineDataSegment(_),
+             .wasmDefineElementSegment(_):
             break
         case .wasmDefineTag(_):
             break

Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift

@@ -235,6 +235,10 @@ public let codeGeneratorWeights = [
     "WasmMemoryInitGenerator":                  5,
     "WasmDefineGlobalGenerator":                2,
     "WasmDefineTableGenerator":                 2,
+    // TODO(427115604): update onece both init and copy instructions are implemented.
+    "WasmDefineElementSegmentGenerator":        1,
+    // TODO(427115604): update onece both init and copy instructions are implemented.
+    "WasmDropElementSegmentGenerator":          1,
     "WasmTableSizeGenerator":                   5,
     "WasmTableGrowGenerator":                   1,
     "WasmGlobalStoreGenerator":                 2,

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -528,6 +528,20 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         module.addTable(elementType: elementType, minSize: minSize, maxSize: maxSize, definedEntries: definedEntries, definedEntryValues: definedEntryValues, isTable64: probability(0.5))
     },
 
+    CodeGenerator("WasmDefineElementSegmentGenerator", inContext: .wasm) { b in
+        let elementsType: ILType = .wasmFunctionDef() | .function()
+        if b.randomVariable(ofType: elementsType) == nil {
+            return
+        }
+
+        var elements: [Variable] = (0...Int.random(in: 0...8)).map {_ in b.randomVariable(ofType: elementsType)!}
+        b.currentWasmModule.addElementSegment(elementsType: elementsType, elements: elements)
+    },
+
+    CodeGenerator("WasmDropElementSegmentGenerator", inContext: .wasmFunction, inputs: .required(.wasmElementSegment())) { b, elementSegment in
+        b.currentWasmFunction.wasmDropElementSegment(elementSegment: elementSegment)
+    },
+
     CodeGenerator("WasmTableSizeGenerator", inContext: .wasmFunction, inputs: .required(.object(ofGroup: "WasmTable"))) { b, table in
         let function = b.currentWasmModule.currentWasmFunction
         function.wasmTableSize(table: table)

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1231,6 +1231,16 @@ extension Instruction: ProtobufConvertible {
                     }
                     $0.isTable64 = op.isTable64
                 }
+            case .wasmDefineElementSegment(let op):
+                $0.wasmDefineElementSegment = Fuzzilli_Protobuf_WasmDefineElementSegment.with {
+                    $0.size = op.size
+                }
+            case .wasmDropElementSegment(_):
+                $0.wasmDropElementSegment = Fuzzilli_Protobuf_WasmDropElementSegment()
+            case .wasmTableCopy(_):
+                $0.wasmTableCopy = Fuzzilli_Protobuf_WasmTableCopy()
+            case .wasmTableInit(_):
+                $0.wasmTableInit = Fuzzilli_Protobuf_WasmTableInit()
             case .wasmDefineMemory(let op):
                 assert(op.wasmMemory.isWasmMemoryType)
                 let mem = op.wasmMemory.wasmMemoryType!
@@ -2322,6 +2332,14 @@ extension Instruction: ProtobufConvertible {
                                      WasmTableType.IndexInTableAndWasmSignature(indexInTable: Int(entry.index), signature: WasmSignatureFromProto(entry.signature))
                                  },
                                  isTable64: p.isTable64)
+        case .wasmDefineElementSegment(let p):
+            op = WasmDefineElementSegment(size: p.size)
+        case .wasmDropElementSegment(_):
+            op = WasmDropElementSegment()
+        case .wasmTableInit(_):
+            op = WasmTableInit()
+        case .wasmTableCopy(_):
+            op = WasmTableCopy()
         case .wasmDefineMemory(let p):
             let maxPages = p.wasmMemory.hasMaxPages ? Int(p.wasmMemory.maxPages) : nil
             op = WasmDefineMemory(limits: Limits(min: Int(p.wasmMemory.minPages), max: maxPages), isShared: p.wasmMemory.isShared, isMemory64: p.wasmMemory.isMemory64)

Sources/Fuzzilli/FuzzIL/JSTyper.swift

@@ -673,6 +673,14 @@ public struct JSTyper: Analyzer {
                     let jsSignature = ProgramBuilder.convertWasmSignatureToJsSignature(entry.signature)
                     dynamicObjectGroupManager.addWasmFunction(withSignature: jsSignature, forDefinition: definingInstruction, forVariable: instr.input(idx))
                 }
+            case .wasmDefineElementSegment(let op):
+                setType(of: instr.output, to: .wasmElementSegment(segmentLength: Int(op.size)))
+            case .wasmDropElementSegment(_):
+                type(of: instr.input(0)).wasmElementSegmentType!.markAsDropped()
+            case .wasmTableInit(_),
+                 .wasmTableCopy(_):
+                // TODO(427115604): - implement both init and copy instructions.
+                break
             case .wasmDefineMemory(let op):
                 setType(of: instr.output, to: op.wasmMemory)
                 registerWasmMemoryUse(for: instr.output)

Sources/Fuzzilli/FuzzIL/Opcodes.swift

@@ -352,4 +352,8 @@ enum Opcode {
     case wasmAnyConvertExtern(WasmAnyConvertExtern)
     case wasmExternConvertAny(WasmExternConvertAny)
     case wasmMemoryCopy(WasmMemoryCopy)
+    case wasmDefineElementSegment(WasmDefineElementSegment)
+    case wasmDropElementSegment(WasmDropElementSegment)
+    case wasmTableInit(WasmTableInit)
+    case wasmTableCopy(WasmTableCopy)
 }

Sources/Fuzzilli/FuzzIL/TypeSystem.swift

@@ -212,6 +212,12 @@ public struct ILType: Hashable {
         return ILType(definiteType: .wasmDataSegment, ext: typeExtension)
     }
 
+    public static func  wasmElementSegment(segmentLength: Int? = nil) -> ILType {
+        let maybeWasmExtention = segmentLength.map { WasmElementSegmentType(segmentLength: $0) }
+        let typeExtension = TypeExtension(group: "WasmElementSegment", properties: Set(), methods: Set(), signature: nil, wasmExt: maybeWasmExtention)
+        return ILType(definiteType: .wasmElementSegment, ext: typeExtension)
+    }
+
     public static func wasmTable(wasmTableType: WasmTableType) -> ILType {
         return .object(ofGroup: "WasmTable", withProperties: ["length"], withMethods: ["get", "grow", "set"], withWasmType: wasmTableType)
     }
@@ -518,6 +524,14 @@ public struct ILType: Hashable {
         return wasmDataSegmentType != nil
     }
 
+    public var wasmElementSegmentType: WasmElementSegmentType? {
+        return ext?.wasmExt as? WasmElementSegmentType
+    }
+
+    public var isWasmElementSegmentType: Bool {
+        return wasmElementSegmentType != nil
+    }
+
 
     public var wasmTableType: WasmTableType? {
         return ext?.wasmExt as? WasmTableType
@@ -1082,6 +1096,8 @@ extension ILType: CustomStringConvertible {
             return ".exceptionLabel"
         case .wasmDataSegment:
             return ".wasmDataSegment"
+        case .wasmElementSegment:
+            return ".wasmElementSegment"
         default:
             break
         }
@@ -1157,6 +1173,7 @@ struct BaseType: OptionSet, Hashable {
     static let wasmPackedI16 = BaseType(rawValue: 1 << 23)
 
     static let wasmDataSegment = BaseType(rawValue: 1 << 24)
+    static let wasmElementSegment = BaseType(rawValue: 1 << 25)
 
     static let jsAnything    = BaseType([.undefined, .integer, .float, .string, .boolean, .object, .function, .constructor, .unboundFunction, .bigint, .regexp, .iterable])
 
@@ -1659,6 +1676,30 @@ public class WasmDataSegmentType: WasmTypeExtension {
     }
 }
 
+public class WasmElementSegmentType: WasmTypeExtension {
+    let segmentLength: Int
+    private(set) var isDropped: Bool
+
+    override func isEqual(to other: WasmTypeExtension) -> Bool {
+        guard let other = other as? WasmElementSegmentType else { return false }
+        return self.segmentLength == other.segmentLength && self.isDropped == other.isDropped
+    }
+
+    override public func hash(into hasher: inout Hasher) {
+        hasher.combine(segmentLength)
+        hasher.combine(isDropped)
+    }
+
+    init(segmentLength: Int) {
+        self.segmentLength = segmentLength
+        self.isDropped = false
+    }
+
+    public func markAsDropped() {
+        self.isDropped = true
+    }
+}
+
 public class WasmTableType: WasmTypeExtension {
     public struct IndexInTableAndWasmSignature: Hashable {
         let indexInTable: Int

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -810,6 +810,44 @@ final class WasmDefineTable: WasmOperation {
     }
 }
 
+final class WasmDefineElementSegment: WasmOperation {
+    override var opcode: Opcode { .wasmDefineElementSegment(self) }
+
+    public let size: UInt32
+
+    init(size: UInt32) {
+        self.size = size
+        super.init(numInputs: Int(size), numOutputs: 1, requiredContext: [.wasm])
+    }
+}
+
+class WasmDropElementSegment: WasmOperation {
+    override var opcode: Opcode { .wasmDropElementSegment(self) }
+
+    init() {
+        super.init(
+            numInputs: 1, numOutputs: 0, requiredContext: [.wasmFunction])
+    }
+}
+
+class WasmTableInit: WasmOperation {
+    override var opcode: Opcode { .wasmTableInit(self) }
+
+    init() {
+        super.init(
+            numInputs: 5, numOutputs: 0, requiredContext: [.wasmFunction])
+    }
+}
+
+class WasmTableCopy: WasmOperation {
+    override var opcode: Opcode { .wasmTableCopy(self) }
+
+    init() {
+        super.init(
+            numInputs: 5, numOutputs: 0, requiredContext: [.wasmFunction])
+    }
+}
+
 // TODO: Wasm memory can be initialized in the data segment, theoretically one could initialize them with this instruction as well.
 // Currently they are by default zero initialized and fuzzilli should then just store or load from there.
 // Also note: https://webassembly.github.io/spec/core/syntax/modules.html#memories
@@ -2355,4 +2393,3 @@ final class WasmAtomicCmpxchg: WasmOperation {
         super.init(numInputs: 4, numOutputs: 1, attributes: [.isMutable], requiredContext: [.wasmFunction])
     }
 }
-

Sources/Fuzzilli/Lifting/FuzzILLifter.swift

@@ -822,6 +822,18 @@ public class FuzzILLifter: Lifter {
             let isTable64Str = op.isTable64 ? ", table64" : ""
             w.emit("\(output()) <- WasmDefineTable \(op.elementType)\(isTable64Str), (\(op.limits.min), \(String(describing: op.limits.max))), [\(entries)]")
 
+        case .wasmDefineElementSegment(_):
+            w.emit("\(output()) <- WasmDefineElementSegment [...]")
+
+        case .wasmDropElementSegment:
+            w.emit("WasmDropElementSegment \(input(0))")
+
+        case .wasmTableInit:
+            w.emit("WasmTableInit \(input(0)), \(input(1)), \(input(2)), \(input(3)), \(input(4))")
+
+        case .wasmTableCopy:
+            w.emit("WasmTableCopy \(input(0)), \(input(1)), \(input(2)), \(input(3)), \(input(4))")
+
         case .wasmDefineMemory(let op):
             assert(op.wasmMemory.isWasmMemoryType)
             let mem = op.wasmMemory.wasmMemoryType!

Sources/Fuzzilli/Lifting/JavaScriptLifter.swift

@@ -1624,6 +1624,10 @@ public class JavaScriptLifter: Lifter {
                  .wasmReassign(_),
                  .wasmDefineGlobal(_),
                  .wasmDefineTable(_),
+                 .wasmDefineElementSegment(_),
+                 .wasmDropElementSegment(_),
+                 .wasmTableInit(_),
+                 .wasmTableCopy(_),
                  .wasmDefineMemory(_),
                  .wasmDefineDataSegment(_),
                  .wasmLoadGlobal(_),