Fix leak of iseq buffer and default locals

swiftwasm · Jul 1, 2024 · 5339d89 · 5339d89
1 parent 967e7e7
commit 5339d89
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 38 deletions.
diff --git a/Sources/WasmKit/Execution/Instructions/Expression.swift b/Sources/WasmKit/Execution/Instructions/Expression.swift
@@ -6,21 +6,12 @@ struct InstructionSequence: Equatable {
     /// This height does not count the locals.
     let maxStackHeight: Int
 
-    init(instructions: [Instruction], maxStackHeight: Int) {
-        assert(_isPOD(Instruction.self))
-        let buffer = UnsafeMutableBufferPointer<Instruction>.allocate(capacity: instructions.count + 1)
-        for (idx, instruction) in instructions.enumerated() {
-            buffer[idx] = instruction
-        }
-        buffer[instructions.count] = .endOfFunction
-        self.instructions = UnsafeBufferPointer(buffer)
+    init(instructions: UnsafeBufferPointer<Instruction>, maxStackHeight: Int) {
+        self.instructions = instructions
+        assert(self.instructions.last == .endOfFunction)
         self.maxStackHeight = maxStackHeight
     }
 
-    func deallocate() {
-        instructions.deallocate()
-    }
-
     var baseAddress: UnsafePointer<Instruction> {
         self.instructions.baseAddress!
     }

diff --git a/Sources/WasmKit/Execution/Runtime/Runtime.swift b/Sources/WasmKit/Execution/Runtime/Runtime.swift
@@ -87,18 +87,21 @@ extension Runtime {
                         .numericConst(.i32(UInt32(element.initializer.count))),
                         .tableInit(tableIndex, elementIndex),
                         .tableElementDrop(elementIndex),
+                        .endOfFunction,
                     ])
-                    let initIseq = InstructionSequence(instructions: instructions, maxStackHeight: 2)
-                    defer { initIseq.deallocate() }
-                    try evaluateConstExpr(initIseq, instance: instance)
+                    try instructions.withUnsafeBufferPointer {
+                        let initIseq = InstructionSequence(instructions: $0, maxStackHeight: 2)
+                        try evaluateConstExpr(initIseq, instance: instance)
+                    }
 
                 case .declarative:
-                    let initIseq = InstructionSequence(
-                        instructions: [.tableElementDrop(elementIndex)],
-                        maxStackHeight: 0
-                    )
-                    defer { initIseq.deallocate() }
-                    try evaluateConstExpr(initIseq, instance: instance)
+                    let instructions: [Instruction] = [.tableElementDrop(elementIndex), .endOfFunction]
+                    try instructions.withUnsafeBufferPointer {
+                        let initIseq = InstructionSequence(
+                            instructions: $0, maxStackHeight: 0
+                        )
+                        try evaluateConstExpr(initIseq, instance: instance)
+                    }
 
                 case .passive:
                     continue
@@ -125,14 +128,17 @@ extension Runtime {
                 default:
                     throw InstantiationError.unsupported("init expr in data section \(data.offset)")
                 }
-                let iseq = InstructionSequence(instructions: instructions + [
+                instructions.append(contentsOf: [
                     .numericConst(.i32(0)),
                     .numericConst(.i32(UInt32(data.initializer.count))),
                     .memoryInit(UInt32(dataIndex)),
                     .memoryDataDrop(UInt32(dataIndex)),
-                ], maxStackHeight: 2)
-                defer { iseq.deallocate() }
-                try evaluateConstExpr(iseq, instance: instance)
+                    .endOfFunction,
+                ])
+                try instructions.withUnsafeBufferPointer {
+                    let iseq = InstructionSequence(instructions: $0, maxStackHeight: 2)
+                    try evaluateConstExpr(iseq, instance: instance)
+                }
             }
         } catch Trap.outOfBoundsMemoryAccess {
             throw InstantiationError.outOfBoundsMemoryAccess
@@ -193,10 +199,12 @@ extension Runtime {
                 default:
                     throw InstantiationError.unsupported("init expr in global section \(global.initializer)")
                 }
-                let iseq = InstructionSequence(instructions: instructions, maxStackHeight: 1)
-                defer { iseq.deallocate() }
-                return try evaluateConstExpr(iseq, instance: globalModuleInstance, arity: 1) { _, stack in
-                    return stack.popValue()
+                instructions.append(.endOfFunction)
+                return try instructions.withUnsafeBufferPointer {
+                    let iseq = InstructionSequence(instructions: $0, maxStackHeight: 1)
+                    return try evaluateConstExpr(iseq, instance: globalModuleInstance, arity: 1) { _, stack in
+                        return stack.popValue()
+                    }
                 }
             }
 

diff --git a/Sources/WasmKit/ModuleParser.swift b/Sources/WasmKit/ModuleParser.swift
@@ -129,7 +129,7 @@ func parseModule<Stream: ByteStream>(stream: Stream, features: WasmFeatureSet =
         let funcTypeIndex = typeIndices[index]
         let funcType = module.types[Int(funcTypeIndex)]
         return GuestFunction(
-            type: typeIndices[index], locals: code.locals,
+            type: typeIndices[index], locals: code.locals, allocator: module.allocator,
             body: {
                 let enableAssert: Bool
                 #if ASSERT

diff --git a/Sources/WasmKit/Translator.swift b/Sources/WasmKit/Translator.swift
@@ -11,6 +11,22 @@ class ISeqAllocator {
         return buffer
     }
 
+    func allocateDefaultLocals(_ locals: [ValueType]) -> UnsafeBufferPointer<Value> {
+        let buffer = UnsafeMutableBufferPointer<Value>.allocate(capacity: locals.count)
+        for (index, localType) in locals.enumerated() {
+            buffer[index] = localType.defaultValue
+        }
+        self.buffers.append(UnsafeMutableRawBufferPointer(buffer))
+        return UnsafeBufferPointer(buffer)
+    }
+
+    func allocateInstructions(capacity: Int) -> UnsafeMutableBufferPointer<Instruction> {
+        assert(_isPOD(Instruction.self), "Instruction must be POD")
+        let buffer = UnsafeMutableBufferPointer<Instruction>.allocate(capacity: capacity)
+        self.buffers.append(UnsafeMutableRawBufferPointer(buffer))
+        return buffer
+    }
+
     deinit {
         for buffer in buffers {
             buffer.deallocate()
@@ -404,7 +420,13 @@ struct InstructionTranslator: InstructionVisitor {
         iseqBuilder.assertDanglingLabels()
         #endif
         let instructions = iseqBuilder.finalize()
-        return InstructionSequence(instructions: instructions, maxStackHeight: valueStack.maxHeight)
+        // TODO: Figure out a way to avoid the copy here while keeping the execution performance.
+        let buffer = allocator.allocateInstructions(capacity: instructions.count + 1)
+        for (idx, instruction) in instructions.enumerated() {
+            buffer[idx] = instruction
+        }
+        buffer[instructions.count] = .endOfFunction
+        return InstructionSequence(instructions: UnsafeBufferPointer(buffer), maxStackHeight: valueStack.maxHeight)
     }
 
     // MARK: - Visitor

diff --git a/Sources/WasmKit/Types/Module.swift b/Sources/WasmKit/Types/Module.swift
@@ -80,14 +80,14 @@ typealias LabelIndex = UInt32
 /// > Note:
 /// <https://webassembly.github.io/spec/core/syntax/modules.html#functions>
 struct GuestFunction {
-    init(type: TypeIndex, locals: [WasmParser.ValueType], body: @escaping () throws -> InstructionSequence) {
+    init(
+        type: TypeIndex,
+        locals: [WasmParser.ValueType],
+        allocator: ISeqAllocator,
+        body: @escaping () throws -> InstructionSequence
+    ) {
         self.type = type
-        // TODO: Deallocate const default locals after the module is deallocated
-        let defaultLocals = UnsafeMutableBufferPointer<Value>.allocate(capacity: locals.count)
-        for (index, localType) in locals.enumerated() {
-            defaultLocals[index] = localType.defaultValue
-        }
-        self.defaultLocals = UnsafeBufferPointer(defaultLocals)
+        self.defaultLocals = allocator.allocateDefaultLocals(locals)
         self.materializer = body
     }