Skip to content

Commit 869d691

Browse files
committed
perf: keep each value stack's height outside its Vec
The height of a value-stack lane is its own field, over a Vec whose length only grows to the most the lane has ever reserved. Slots above the height hold stale values that are written before they are read again. The tail-call handlers can then pass the height between them in a register and write it back without unsafe code. Entering a function zeroes one or two locals with direct stores, and only the ones in between with memset.
1 parent 3a1723e commit 869d691

1 file changed

Lines changed: 89 additions & 42 deletions

File tree

‎crates/tinywasm/src/interpreter/stack/value_stack.rs‎

Lines changed: 89 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -18,45 +18,54 @@ pub(crate) struct ValueStack {
1818
}
1919

2020
#[cfg_attr(feature = "debug", derive(Debug))]
21+
/// One value lane: a stack of `len` values over `data`.
22+
///
23+
/// `data` only grows. Its length is the most the stack has ever reserved, and the slots above
24+
/// `len` hold stale values that are written before they are read again. Keeping the height outside
25+
/// the `Vec` lets the tail-call handlers carry it in a register and write it back without `unsafe`.
2126
pub(crate) struct Stack<T: Copy + Default> {
2227
data: Vec<T>,
28+
len: usize,
2329
max_size: usize,
2430
dynamic: bool,
2531
}
2632

2733
impl<T: Copy + Default> Stack<T> {
2834
pub(crate) fn new(config: StackConfig) -> Self {
29-
Self { data: Vec::with_capacity(config.initial_size), max_size: config.max_size, dynamic: config.dynamic }
35+
Self { data: Vec::with_capacity(config.initial_size), len: 0, max_size: config.max_size, dynamic: config.dynamic }
3036
}
3137

3238
pub(crate) fn clear(&mut self) {
33-
self.data.clear();
39+
self.len = 0;
3440
}
3541

3642
#[inline(always)]
3743
pub(crate) fn len(&self) -> usize {
38-
self.data.len()
44+
self.len
3945
}
4046

4147
/// Pushes a value inside a function body. `enter_locals` reserved the function's whole operand
42-
/// stack, so the stack is never full here. After this check `Vec::push` cannot reach its own
43-
/// growth path, so the instruction handlers make no calls.
48+
/// stack, so the stack is never full here, and the handlers make no calls.
4449
#[inline(always)]
4550
pub(crate) fn push(&mut self, value: T) {
46-
if self.data.len() == self.data.capacity() {
47-
crate::invariant_violated("value stack push beyond the function's reservation");
51+
let len = self.len;
52+
match self.data.get_mut(len) {
53+
Some(slot) => *slot = value,
54+
None => crate::invariant_violated("value stack push beyond the function's reservation"),
4855
}
49-
self.data.push(value);
56+
self.len = len + 1;
5057
}
5158

5259
/// Pushes a value outside a function body (host arguments and results), which no reservation
5360
/// covers, so a dynamic stack grows here if needed.
5461
#[inline(always)]
5562
pub(crate) fn push_or_grow(&mut self, value: T) -> Result<(), Trap> {
56-
if self.data.len() == self.data.capacity() {
57-
return self.push_grow(value);
63+
let len = self.len;
64+
match self.data.get_mut(len) {
65+
Some(slot) => *slot = value,
66+
None => return self.push_grow(value),
5867
}
59-
self.data.push(value);
68+
self.len = len + 1;
6069
Ok(())
6170
}
6271

@@ -69,33 +78,36 @@ impl<T: Copy + Default> Stack<T> {
6978
#[cold]
7079
#[inline(never)]
7180
fn push_grow(&mut self, value: T) -> Result<(), Trap> {
72-
// Check the limit only at capacity to avoid an extra hot-path check. Vec growth may
73-
// intentionally overshoot max_size.
74-
if !self.dynamic || self.data.len() >= self.max_size {
81+
// Only the reserved capacity is free to use. Past it, check the limit, which Vec growth
82+
// may intentionally overshoot.
83+
if self.data.len() == self.data.capacity() && (!self.dynamic || self.data.len() >= self.max_size) {
7584
return Err(Trap::ValueStackOverflow);
7685
}
7786
self.data.push(value);
87+
self.len = self.data.len();
7888
Ok(())
7989
}
8090

91+
/// Pops the top value. On an empty stack the index wraps around, so the bounds check in
92+
/// [`Self::get`] also catches an underflow.
8193
#[inline(always)]
8294
pub(crate) fn pop(&mut self) -> T {
83-
match self.data.pop() {
84-
Some(value) => value,
85-
None => crate::invariant_violated("value stack underflow"),
86-
}
95+
let index = self.len.wrapping_sub(1);
96+
let value = *self.get(index);
97+
self.len = index;
98+
value
8799
}
88100

89101
#[inline(always)]
90102
pub(crate) fn last(&self) -> &T {
91-
match self.data.last() {
92-
Some(value) => value,
93-
None => crate::invariant_violated("value stack underflow"),
94-
}
103+
self.get(self.len.wrapping_sub(1))
95104
}
96105

106+
/// The slot at `index`, which validation keeps below the height. The check is against the
107+
/// slots in use at any height, which is what memory safety needs.
97108
#[inline(always)]
98109
pub(crate) fn get(&self, index: usize) -> &T {
110+
debug_assert!(index < self.len);
99111
match self.data.get(index) {
100112
Some(value) => value,
101113
None => crate::invariant_violated("value stack index out of range"),
@@ -104,6 +116,7 @@ impl<T: Copy + Default> Stack<T> {
104116

105117
#[inline(always)]
106118
pub(crate) fn set(&mut self, index: usize, value: T) {
119+
debug_assert!(index < self.len);
107120
match self.data.get_mut(index) {
108121
Some(slot) => *slot = value,
109122
None => crate::invariant_violated("value stack index out of range"),
@@ -118,29 +131,35 @@ impl<T: Copy + Default> Stack<T> {
118131

119132
#[inline(always)]
120133
pub(crate) fn truncate_keep(&mut self, n: usize, end_keep: usize) {
121-
let len = self.data.len();
134+
let len = self.len;
122135
debug_assert!(n <= len);
123136
if n >= len {
124137
return;
125138
}
126139

127140
let keep = len.wrapping_sub(n).min(end_keep);
128-
self.data.copy_within(len.wrapping_sub(keep)..len, n);
129-
self.data.truncate(n.wrapping_add(keep));
141+
if keep != 0 {
142+
// Copying to the start of the values above `n` leaves copy_within no check to fail.
143+
match self.data.get_mut(n..len) {
144+
Some(above) => above.copy_within(above.len().wrapping_sub(keep).., 0),
145+
None => crate::invariant_violated("value stack index out of range"),
146+
}
147+
}
148+
self.len = n.wrapping_add(keep);
130149
}
131150

132151
#[inline(always)]
133152
pub(crate) fn truncate_to(&mut self, n: usize) {
134-
debug_assert!(n <= self.data.len());
135-
self.data.truncate(n);
153+
debug_assert!(n <= self.len);
154+
self.len = n;
136155
}
137156

138157
#[inline(always)]
139158
pub(crate) fn truncate_to_one_tail(&mut self, n: usize) {
140-
debug_assert!(n < self.data.len());
159+
debug_assert!(n < self.len);
141160
let last = self.pop();
142-
self.data.truncate(n);
143-
self.data.push(last);
161+
self.len = n;
162+
self.push(last);
144163
}
145164

146165
/// Enters a function: turns its parameters into the first locals, zeroes the rest, and reserves
@@ -154,27 +173,52 @@ impl<T: Copy + Default> Stack<T> {
154173
max_stack: usize,
155174
) -> Result<u32, Trap> {
156175
debug_assert!(param_count <= local_count);
157-
debug_assert!(param_count <= self.data.len());
176+
debug_assert!(param_count <= self.len);
158177

159-
let len = self.data.len();
178+
let len = self.len;
160179
let start = len - param_count;
161180
let end = start + local_count;
162181
let reserve = end + max_stack;
163182

164-
if reserve > self.data.capacity() {
183+
if reserve > self.data.len() {
165184
core::hint::cold_path();
185+
self.grow_to(reserve)?;
186+
}
187+
188+
// Most functions have no or few locals in a lane. Store the first and last directly and
189+
// fill only what lies between: `fill` becomes a memset call even for one value.
190+
if end > len {
191+
match self.data.get_mut(len..end) {
192+
Some([a]) => *a = T::default(),
193+
Some([a, middle @ .., b]) => {
194+
if !middle.is_empty() {
195+
middle.fill(T::default());
196+
}
197+
(*a, *b) = (T::default(), T::default());
198+
}
199+
_ => crate::invariant_violated("value stack index out of range"),
200+
}
201+
}
202+
self.len = end;
203+
Ok(start as u32)
204+
}
205+
206+
/// Makes `reserve` slots usable, growing the allocation if a dynamic stack allows it.
207+
#[cold]
208+
#[inline(never)]
209+
fn grow_to(&mut self, reserve: usize) -> Result<(), Trap> {
210+
let cap = self.data.capacity();
211+
if reserve > cap {
166212
if reserve > self.max_size || !self.dynamic {
167213
return Err(Trap::ValueStackOverflow);
168214
}
169-
let cap = self.data.capacity();
170215
let target = reserve.max(cap.max(1).saturating_mul(2)).min(self.max_size);
171-
if self.data.try_reserve(target - len).is_err() {
216+
if self.data.try_reserve(target - self.data.len()).is_err() {
172217
return Err(Trap::ValueStackOverflow);
173218
}
174219
}
175-
176-
self.data.resize(end, T::default());
177-
Ok(start as u32)
220+
self.data.resize(reserve, T::default());
221+
Ok(())
178222
}
179223

180224
#[inline(always)]
@@ -183,7 +227,7 @@ impl<T: Copy + Default> Stack<T> {
183227
return;
184228
}
185229

186-
let len = self.data.len();
230+
let len = self.len;
187231
let needed = count.wrapping_mul(2);
188232

189233
if len < needed {
@@ -193,10 +237,13 @@ impl<T: Copy + Default> Stack<T> {
193237
if !condition {
194238
let dst = len.wrapping_sub(needed);
195239
let src = len.wrapping_sub(count);
196-
self.data.copy_within(src..len, dst);
240+
match self.data.get_mut(..len) {
241+
Some(live) => live.copy_within(src..len, dst),
242+
None => crate::invariant_violated("value stack index out of range"),
243+
}
197244
}
198245

199-
self.data.truncate(len.wrapping_sub(count));
246+
self.len = len.wrapping_sub(count);
200247
}
201248
}
202249

@@ -205,7 +252,7 @@ impl<'a, T: Copy + Default> IntoIterator for &'a Stack<T> {
205252
type IntoIter = core::slice::Iter<'a, T>;
206253

207254
fn into_iter(self) -> Self::IntoIter {
208-
self.data.iter()
255+
self.data[..self.len].iter()
209256
}
210257
}
211258
impl ValueStack {

0 commit comments

Comments
 (0)