Skip to content

Commit 4ae4f76

Browse files
kateinoigakukunkateinoigakukun
committed
Fix leak of iseq buffer and default locals
1 parent 967e7e7 commit 4ae4f76

File tree

6 files changed

+90
-48
lines changed

6 files changed

+90
-48
lines changed

‎Sources/WasmKit/Execution/Instructions/Expression.swift‎

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,12 @@ struct InstructionSequence: Equatable {
66
/// This height does not count the locals.
77
let maxStackHeight: Int
88

9-
init(instructions: [Instruction], maxStackHeight: Int) {
10-
assert(_isPOD(Instruction.self))
11-
let buffer = UnsafeMutableBufferPointer<Instruction>.allocate(capacity: instructions.count + 1)
12-
for (idx, instruction) in instructions.enumerated() {
13-
buffer[idx] = instruction
14-
}
15-
buffer[instructions.count] = .endOfFunction
16-
self.instructions = UnsafeBufferPointer(buffer)
9+
init(instructions: UnsafeBufferPointer<Instruction>, maxStackHeight: Int) {
10+
self.instructions = instructions
11+
assert(self.instructions.last == .endOfFunction)
1712
self.maxStackHeight = maxStackHeight
1813
}
1914

20-
func deallocate() {
21-
instructions.deallocate()
22-
}
23-
2415
var baseAddress: UnsafePointer<Instruction> {
2516
self.instructions.baseAddress!
2617
}

‎Sources/WasmKit/Execution/Runtime/Runtime.swift‎

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -87,18 +87,21 @@ extension Runtime {
8787
.numericConst(.i32(UInt32(element.initializer.count))),
8888
.tableInit(tableIndex, elementIndex),
8989
.tableElementDrop(elementIndex),
90+
.endOfFunction,
9091
])
91-
let initIseq = InstructionSequence(instructions: instructions, maxStackHeight: 2)
92-
defer { initIseq.deallocate() }
93-
try evaluateConstExpr(initIseq, instance: instance)
92+
try instructions.withUnsafeBufferPointer {
93+
let initIseq = InstructionSequence(instructions: $0, maxStackHeight: 2)
94+
try evaluateConstExpr(initIseq, instance: instance)
95+
}
9496

9597
case .declarative:
96-
let initIseq = InstructionSequence(
97-
instructions: [.tableElementDrop(elementIndex)],
98-
maxStackHeight: 0
99-
)
100-
defer { initIseq.deallocate() }
101-
try evaluateConstExpr(initIseq, instance: instance)
98+
let instructions: [Instruction] = [.tableElementDrop(elementIndex), .endOfFunction]
99+
try instructions.withUnsafeBufferPointer {
100+
let initIseq = InstructionSequence(
101+
instructions: $0, maxStackHeight: 0
102+
)
103+
try evaluateConstExpr(initIseq, instance: instance)
104+
}
102105

103106
case .passive:
104107
continue
@@ -125,14 +128,17 @@ extension Runtime {
125128
default:
126129
throw InstantiationError.unsupported("init expr in data section \(data.offset)")
127130
}
128-
let iseq = InstructionSequence(instructions: instructions + [
131+
instructions.append(contentsOf: [
129132
.numericConst(.i32(0)),
130133
.numericConst(.i32(UInt32(data.initializer.count))),
131134
.memoryInit(UInt32(dataIndex)),
132135
.memoryDataDrop(UInt32(dataIndex)),
133-
], maxStackHeight: 2)
134-
defer { iseq.deallocate() }
135-
try evaluateConstExpr(iseq, instance: instance)
136+
.endOfFunction,
137+
])
138+
try instructions.withUnsafeBufferPointer {
139+
let iseq = InstructionSequence(instructions: $0, maxStackHeight: 2)
140+
try evaluateConstExpr(iseq, instance: instance)
141+
}
136142
}
137143
} catch Trap.outOfBoundsMemoryAccess {
138144
throw InstantiationError.outOfBoundsMemoryAccess
@@ -193,10 +199,12 @@ extension Runtime {
193199
default:
194200
throw InstantiationError.unsupported("init expr in global section \(global.initializer)")
195201
}
196-
let iseq = InstructionSequence(instructions: instructions, maxStackHeight: 1)
197-
defer { iseq.deallocate() }
198-
return try evaluateConstExpr(iseq, instance: globalModuleInstance, arity: 1) { _, stack in
199-
return stack.popValue()
202+
instructions.append(.endOfFunction)
203+
return try instructions.withUnsafeBufferPointer {
204+
let iseq = InstructionSequence(instructions: $0, maxStackHeight: 1)
205+
return try evaluateConstExpr(iseq, instance: globalModuleInstance, arity: 1) { _, stack in
206+
return stack.popValue()
207+
}
200208
}
201209
}
202210

‎Sources/WasmKit/ModuleParser.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ func parseModule<Stream: ByteStream>(stream: Stream, features: WasmFeatureSet =
129129
let funcTypeIndex = typeIndices[index]
130130
let funcType = module.types[Int(funcTypeIndex)]
131131
return GuestFunction(
132-
type: typeIndices[index], locals: code.locals,
132+
type: typeIndices[index], locals: code.locals, allocator: module.allocator,
133133
body: {
134134
let enableAssert: Bool
135135
#if ASSERT

‎Sources/WasmKit/Translator.swift‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,22 @@ class ISeqAllocator {
1111
return buffer
1212
}
1313

14+
func allocateDefaultLocals(_ locals: [ValueType]) -> UnsafeBufferPointer<Value> {
15+
let buffer = UnsafeMutableBufferPointer<Value>.allocate(capacity: locals.count)
16+
for (index, localType) in locals.enumerated() {
17+
buffer[index] = localType.defaultValue
18+
}
19+
self.buffers.append(UnsafeMutableRawBufferPointer(buffer))
20+
return UnsafeBufferPointer(buffer)
21+
}
22+
23+
func allocateInstructions(capacity: Int) -> UnsafeMutableBufferPointer<Instruction> {
24+
assert(_isPOD(Instruction.self), "Instruction must be POD")
25+
let buffer = UnsafeMutableBufferPointer<Instruction>.allocate(capacity: capacity)
26+
self.buffers.append(UnsafeMutableRawBufferPointer(buffer))
27+
return buffer
28+
}
29+
1430
deinit {
1531
for buffer in buffers {
1632
buffer.deallocate()
@@ -404,7 +420,13 @@ struct InstructionTranslator: InstructionVisitor {
404420
iseqBuilder.assertDanglingLabels()
405421
#endif
406422
let instructions = iseqBuilder.finalize()
407-
return InstructionSequence(instructions: instructions, maxStackHeight: valueStack.maxHeight)
423+
// TODO: Figure out a way to avoid the copy here while keeping the execution performance.
424+
let buffer = allocator.allocateInstructions(capacity: instructions.count + 1)
425+
for (idx, instruction) in instructions.enumerated() {
426+
buffer[idx] = instruction
427+
}
428+
buffer[instructions.count] = .endOfFunction
429+
return InstructionSequence(instructions: UnsafeBufferPointer(buffer), maxStackHeight: valueStack.maxHeight)
408430
}
409431

410432
// MARK: - Visitor

‎Sources/WasmKit/Types/Module.swift‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -80,14 +80,14 @@ typealias LabelIndex = UInt32
8080
/// > Note:
8181
/// <https://webassembly.github.io/spec/core/syntax/modules.html#functions>
8282
struct GuestFunction {
83-
init(type: TypeIndex, locals: [WasmParser.ValueType], body: @escaping () throws -> InstructionSequence) {
83+
init(
84+
type: TypeIndex,
85+
locals: [WasmParser.ValueType],
86+
allocator: ISeqAllocator,
87+
body: @escaping () throws -> InstructionSequence
88+
) {
8489
self.type = type
85-
// TODO: Deallocate const default locals after the module is deallocated
86-
let defaultLocals = UnsafeMutableBufferPointer<Value>.allocate(capacity: locals.count)
87-
for (index, localType) in locals.enumerated() {
88-
defaultLocals[index] = localType.defaultValue
89-
}
90-
self.defaultLocals = UnsafeBufferPointer(defaultLocals)
90+
self.defaultLocals = allocator.allocateDefaultLocals(locals)
9191
self.materializer = body
9292
}
9393

‎Tests/WasmKitTests/Execution/HostModuleTests.swift‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,28 +23,49 @@ final class HostModuleTests: XCTestCase {
2323
func testReentrancy() throws {
2424
let runtime = Runtime()
2525
let voidSignature = WasmParser.FunctionType(parameters: [], results: [])
26+
let allocator = ISeqAllocator()
27+
func compile(_ instructions: [WasmKit.Instruction], maxStackHeight: Int) -> InstructionSequence {
28+
let buffer = allocator.allocateInstructions(capacity: instructions.count + 1)
29+
for (i, instruction) in instructions.enumerated() {
30+
buffer[i] = instruction
31+
}
32+
buffer[instructions.count] = .endOfFunction
33+
return InstructionSequence(
34+
instructions: UnsafeBufferPointer(buffer), maxStackHeight: maxStackHeight
35+
)
36+
}
2637
let module = Module(
2738
types: [voidSignature],
2839
functions: [
2940
// [0] (import "env" "bar" func)
3041
// [1] (import "env" "qux" func)
3142
// [2] "foo"
3243
GuestFunction(
33-
type: 0, locals: [],
44+
type: 0,
45+
locals: [],
46+
allocator: allocator,
3447
body: {
35-
[
36-
.call(functionIndex: 0),
37-
.call(functionIndex: 0),
38-
.call(functionIndex: 0),
39-
]
48+
compile(
49+
[
50+
.call(functionIndex: 0),
51+
.call(functionIndex: 0),
52+
.call(functionIndex: 0),
53+
],
54+
maxStackHeight: 0
55+
)
4056
}),
4157
// [3] "bar"
4258
GuestFunction(
43-
type: 0, locals: [],
59+
type: 0,
60+
locals: [],
61+
allocator: allocator,
4462
body: {
45-
[
46-
.call(functionIndex: 1)
47-
]
63+
compile(
64+
[
65+
.call(functionIndex: 1)
66+
],
67+
maxStackHeight: 0
68+
)
4869
}),
4970
],
5071
imports: [

0 commit comments

Comments
 (0)