Skip to content

Commit d4e7b9d

Browse files
dcodeIOdcodeIO
authored
feat: Slightly optimize shadow stack pass (#1667)
1 parent 4ccda0b commit d4e7b9d

File tree

103 files changed

+170472
-222556
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

103 files changed

+170472
-222556
lines changed

‎src/passes/shadowstack.ts‎

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -154,14 +154,26 @@ type SlotMap = Map<LocalIndex,SlotIndex>;
154154
type TempMap = Map<NativeType,LocalIndex>;
155155

156156
/** Attempts to match the `__tostack(value)` pattern. Returns `value` if a match, otherwise `0`. */
157-
function matchTostack(module: Module, expr: ExpressionRef): ExpressionRef {
157+
function matchPattern(module: Module, expr: ExpressionRef): ExpressionRef {
158158
if (_BinaryenExpressionGetId(expr) == ExpressionId.Call && module.readStringCached(_BinaryenCallGetTarget(expr)) == BuiltinNames.tostack) {
159159
assert(_BinaryenCallGetNumOperands(expr) == 1);
160160
return _BinaryenCallGetOperandAt(expr, 0);
161161
}
162162
return 0;
163163
}
164164

165+
/** Tests whether a `value` matched by `matchTostack` needs a slot. */
166+
function needsSlot(module: Module, value: ExpressionRef): bool {
167+
switch (_BinaryenExpressionGetId(value)) {
168+
// no need to stack null pointers
169+
case ExpressionId.Const: return !isConstZero(value);
170+
// already kept in another slot
171+
case ExpressionId.LocalGet:
172+
case ExpressionId.LocalSet: return false; // tee
173+
}
174+
return true;
175+
}
176+
165177
/** Instruments a module with a shadow stack for precise GC. */
166178
export class ShadowStackPass extends Pass {
167179
/** Stack frame slots, per function. */
@@ -320,9 +332,9 @@ export class ShadowStackPass extends Pass {
320332
var numSlots = 0;
321333
for (let i = 0, k = operands.length; i < k; ++i) {
322334
let operand = operands[i];
323-
let match = matchTostack(module, operand);
335+
let match = matchPattern(module, operand);
324336
if (!match) continue;
325-
if (isConstZero(match)) {
337+
if (!needsSlot(module, match)) {
326338
operands[i] = match;
327339
continue;
328340
}
@@ -366,10 +378,10 @@ export class ShadowStackPass extends Pass {
366378
operands[i] = _BinaryenCallGetOperandAt(call, i);
367379
}
368380
let numSlots = this.updateCallOperands(operands);
381+
for (let i = 0, k = operands.length; i < k; ++i) {
382+
_BinaryenCallSetOperandAt(call, i, operands[i]);
383+
}
369384
if (numSlots) {
370-
for (let i = 0, k = operands.length; i < k; ++i) {
371-
_BinaryenCallSetOperandAt(call, i, operands[i]);
372-
}
373385
// Reserve these slots for us so nested calls use their own
374386
this.callSlotOffset += numSlots;
375387
}
@@ -390,10 +402,10 @@ export class ShadowStackPass extends Pass {
390402
operands[i] = _BinaryenCallIndirectGetOperandAt(callIndirect, i);
391403
}
392404
let numSlots = this.updateCallOperands(operands);
405+
for (let i = 0, k = operands.length; i < k; ++i) {
406+
_BinaryenCallIndirectSetOperandAt(callIndirect, i, operands[i]);
407+
}
393408
if (numSlots) {
394-
for (let i = 0, k = operands.length; i < k; ++i) {
395-
_BinaryenCallIndirectSetOperandAt(callIndirect, i, operands[i]);
396-
}
397409
// Reserve these slots for us so nested calls use their own
398410
this.callSlotOffset += numSlots;
399411
}
@@ -408,14 +420,14 @@ export class ShadowStackPass extends Pass {
408420

409421
/** @override */
410422
visitLocalSet(localSet: ExpressionRef): void {
423+
let module = this.module;
411424
let value = _BinaryenLocalSetGetValue(localSet);
412-
let match = matchTostack(this.module, value);
425+
let match = matchPattern(module, value);
413426
if (!match) return;
414-
if (isConstZero(match)) {
427+
if (!needsSlot(module, match)) {
415428
_BinaryenLocalSetSetValue(localSet, match);
416429
return;
417430
}
418-
var module = this.module;
419431
let index = _BinaryenLocalSetGetIndex(localSet);
420432
let slotIndex = this.noteSlot(this.currentFunction, index);
421433
let stmts = new Array<ExpressionRef>();

‎tests/compiler/assert-nonnull.optimized.wat‎

Lines changed: 13 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -84,17 +84,6 @@
8484
global.get $~lib/memory/__stack_pointer
8585
local.get $0
8686
i32.store
87-
global.get $~lib/memory/__stack_pointer
88-
i32.const 4
89-
i32.sub
90-
global.set $~lib/memory/__stack_pointer
91-
call $~stack_check
92-
global.get $~lib/memory/__stack_pointer
93-
i32.const 0
94-
i32.store
95-
global.get $~lib/memory/__stack_pointer
96-
local.get $0
97-
i32.store
9887
local.get $0
9988
i32.eqz
10089
if
@@ -109,10 +98,6 @@
10998
i32.const 4
11099
i32.add
111100
global.set $~lib/memory/__stack_pointer
112-
global.get $~lib/memory/__stack_pointer
113-
i32.const 4
114-
i32.add
115-
global.set $~lib/memory/__stack_pointer
116101
local.get $0
117102
)
118103
(func $export:assert-nonnull/testObj (param $0 i32) (result i32)
@@ -124,17 +109,6 @@
124109
global.get $~lib/memory/__stack_pointer
125110
local.get $0
126111
i32.store
127-
global.get $~lib/memory/__stack_pointer
128-
i32.const 4
129-
i32.sub
130-
global.set $~lib/memory/__stack_pointer
131-
call $~stack_check
132-
global.get $~lib/memory/__stack_pointer
133-
i32.const 0
134-
i32.store
135-
global.get $~lib/memory/__stack_pointer
136-
local.get $0
137-
i32.store
138112
local.get $0
139113
i32.eqz
140114
if
@@ -151,10 +125,6 @@
151125
i32.const 4
152126
i32.add
153127
global.set $~lib/memory/__stack_pointer
154-
global.get $~lib/memory/__stack_pointer
155-
i32.const 4
156-
i32.add
157-
global.set $~lib/memory/__stack_pointer
158128
)
159129
(func $export:assert-nonnull/testProp (param $0 i32) (result i32)
160130
global.get $~lib/memory/__stack_pointer
@@ -208,16 +178,13 @@
208178
local.get $0
209179
i32.store
210180
global.get $~lib/memory/__stack_pointer
211-
i32.const 8
181+
i32.const 4
212182
i32.sub
213183
global.set $~lib/memory/__stack_pointer
214184
call $~stack_check
215185
global.get $~lib/memory/__stack_pointer
216-
i64.const 0
217-
i64.store
218-
global.get $~lib/memory/__stack_pointer
219-
local.get $0
220-
i32.store offset=4
186+
i32.const 0
187+
i32.store
221188
local.get $0
222189
i32.eqz
223190
if
@@ -271,7 +238,7 @@
271238
i32.add
272239
global.set $~lib/memory/__stack_pointer
273240
global.get $~lib/memory/__stack_pointer
274-
i32.const 8
241+
i32.const 4
275242
i32.add
276243
global.set $~lib/memory/__stack_pointer
277244
global.get $~lib/memory/__stack_pointer
@@ -290,21 +257,18 @@
290257
local.get $0
291258
i32.store
292259
global.get $~lib/memory/__stack_pointer
293-
i32.const 8
260+
i32.const 4
294261
i32.sub
295262
global.set $~lib/memory/__stack_pointer
296263
call $~stack_check
297264
global.get $~lib/memory/__stack_pointer
298-
i64.const 0
299-
i64.store
300-
global.get $~lib/memory/__stack_pointer
301-
local.get $0
265+
i32.const 0
302266
i32.store
303267
global.get $~lib/memory/__stack_pointer
304268
local.get $0
305269
call $~lib/array/Array<assert-nonnull/Foo|null>#__get
306270
local.tee $0
307-
i32.store offset=4
271+
i32.store
308272
local.get $0
309273
i32.eqz
310274
if
@@ -316,7 +280,7 @@
316280
unreachable
317281
end
318282
global.get $~lib/memory/__stack_pointer
319-
i32.const 8
283+
i32.const 4
320284
i32.add
321285
global.set $~lib/memory/__stack_pointer
322286
global.get $~lib/memory/__stack_pointer
@@ -348,9 +312,6 @@
348312
local.set $1
349313
global.get $~lib/memory/__stack_pointer
350314
local.set $2
351-
global.get $~lib/memory/__stack_pointer
352-
local.get $0
353-
i32.store offset=4
354315
block $__inlined_func$assert-nonnull/testAll
355316
block $folding-inner0
356317
local.get $0
@@ -417,9 +378,6 @@
417378
local.set $1
418379
global.get $~lib/memory/__stack_pointer
419380
local.set $2
420-
global.get $~lib/memory/__stack_pointer
421-
local.get $0
422-
i32.store offset=4
423381
block $__inlined_func$assert-nonnull/testAll2
424382
block $folding-inner0
425383
local.get $0
@@ -491,17 +449,14 @@
491449
local.get $0
492450
i32.store
493451
global.get $~lib/memory/__stack_pointer
494-
i32.const 8
452+
i32.const 4
495453
i32.sub
496454
global.set $~lib/memory/__stack_pointer
497455
call $~stack_check
498456
global.get $~lib/memory/__stack_pointer
499-
i64.const 0
500-
i64.store
501-
global.get $~lib/memory/__stack_pointer
502-
global.get $~lib/memory/__stack_pointer
503-
local.get $0
457+
i32.const 0
504458
i32.store
459+
global.get $~lib/memory/__stack_pointer
505460
local.get $0
506461
i32.eqz
507462
if
@@ -513,12 +468,12 @@
513468
unreachable
514469
end
515470
local.get $0
516-
i32.store offset=4
471+
i32.store
517472
local.get $0
518473
i32.load
519474
call_indirect (type $none_=>_i32)
520475
global.get $~lib/memory/__stack_pointer
521-
i32.const 8
476+
i32.const 4
522477
i32.add
523478
global.set $~lib/memory/__stack_pointer
524479
global.get $~lib/memory/__stack_pointer

0 commit comments

Comments
 (0)