1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
use std::cell::Cell;
use std::future::Future;
use std::pin::Pin;
use std::ptr::NonNull;
use std::task::{Context, Poll};
use crate::allocator::AllocBox;
use crate::side_channel::SideChannelFuture;
use crate::Alloc;
pub(crate) struct FrameList<'alloc> {
alloc: &'alloc Alloc,
is_dropping: Cell<bool>,
head: Cell<NonNull<Frame>>,
pushed: Cell<bool>,
}
impl<'alloc> FrameList<'alloc> {
pub fn new(alloc: &'alloc Alloc) -> Self {
Self {
alloc,
is_dropping: Cell::new(false),
head: Cell::new(placeholder(alloc)),
pushed: Cell::new(false),
}
}
/// Returns true if the frame list using this is currently being dropped.
pub fn is_dropping(&self) -> bool {
self.is_dropping.get()
}
/// Consume a frame, turning it into an owned box (that will deallocate the
/// frame when dropped).
///
/// # Safety
///
/// The given `frame` parameter must have originated from a call to
/// [`Self::push`], and this must only be used if `self.is_dropping()` is
/// `false`. Otherwise, `FrameList` is in charge of managing the allocation
/// for the frame.
pub unsafe fn consume_frame(&self, frame: NonNull<Frame>) -> AllocBox<'_, Frame> {
// Safety: `frame` comes from `Self::push`, which created it through the
// allocator stored in `self`.
unsafe { self.alloc.box_from_raw(frame.as_ptr()) }
}
/// Drop the entire call stack.
///
/// # Safety
///
/// The frame list must not be used again after this has been called.
pub unsafe fn drop(&self) {
self.is_dropping.set(true);
drop_list(self.alloc, self.head.get())
}
/// Poll the future at the head of the list. Returns `Poll::Ready(())` only
/// if every future in this list has been polled to completion.
///
/// # Safety
///
/// During the execution of this function, only [`Self::push`] and
/// [`Self::consume_frame`] may be invoked.
pub unsafe fn poll(&self, cx: &mut Context) -> Poll<()> {
loop {
let mut pollee = self.head.get();
// Safety:
// - the reference is alive during the future `poll`
// - the pointer is initialized and valid, since it comes from
// `self.bump.new_box(...)`
let future = unsafe { &mut pollee.as_mut().future };
// Safety: this pointer can be safely pinned since it is a field
// within a box. Although the box will later be handed away with
// `self.consume_frame`, the future itself cannot be moved out.
let future = unsafe { Pin::new_unchecked(future) };
match future.poll(cx) {
Poll::Pending => {
// If a new future has been pushed, continue polling it
// instead of yielding.
if self.pushed.replace(false) {
continue;
} else {
break Poll::Pending;
}
}
Poll::Ready(()) => {
if self.pop() {
continue;
} else {
break Poll::Ready(());
}
}
}
}
}
/// Push a new future `f` to the top of the frame list, returning a
/// pointer to its type erased form. The caller is expected to deallocate
/// the returned frame once it has completed.
///
/// # Safety
///
/// The frame must not be moved out of until [`Frame::is_done`] returns
/// `true`.
pub unsafe fn push<'a>(&self, f: impl SideChannelFuture + 'a) -> NonNull<Frame>
where
Self: 'a,
{
let f = Frame::new(self.alloc, Some(self.head.get()), f);
self.head.set(f);
self.pushed.set(true);
f
}
/// # Safety
///
/// This should only be invoked by a receiver when the top of the stack is
/// its corresponding sender.
pub unsafe fn pop_head(&self) -> NonNull<Frame> {
let head = self.head.get();
let did_pop = self.pop();
debug_assert!(did_pop);
head
}
pub fn set_head<'a>(&self, f: impl Future<Output = ()> + 'a)
where
Self: 'a,
{
let f = Frame::new(self.alloc, None, Bottom(f));
self.is_dropping.set(true);
drop_list(self.alloc, self.head.replace(f));
self.is_dropping.set(false);
}
}
impl<'alloc> FrameList<'alloc> {
/// Pop the topmost frame and return `true` if it had anything to return to.
/// If this returns `false`, then the topmost future was the last, and will
/// keep being the top of the stack.
fn pop(&self) -> bool {
let head = self.head.get();
// Safety: the pointer comes from `Box::new`, so it is non-null and
// properly aligned.
let head_frame = unsafe { &(*head.as_ptr()) };
if let Some(next) = head_frame.return_to {
self.head.set(next);
head_frame.is_done.set(true);
true
} else {
// The last frame is dropped when the frame list is dropped.
false
}
}
}
fn drop_list(alloc: &Alloc, f: NonNull<Frame>) {
// Traverse the call stack and drop frames from top to bottom to ensure
// earlier frames outlive later ones.
let mut next = Some(f);
while let Some(head) = next {
// Safety: Every `NonNull` frame pointer that can be accessed here comes
// from `FrameList::push`. Since `meta.is_dropping` is `true`, it's up
// to us to manage the allocations.
let frame = unsafe { alloc.box_from_raw(head.as_ptr()) };
next = frame.return_to;
drop(frame);
}
}
pub(crate) type Frame = FrameInner<dyn SideChannelFuture>;
pub(crate) struct FrameInner<T: ?Sized> {
is_done: Cell<bool>,
return_to: Option<NonNull<Frame>>,
future: T,
}
impl Frame {
/// Allocate a frame within the given allocator and return a type-erased
/// pointer to it.
fn new<'a>(
alloc: &'a Alloc,
return_to: Option<NonNull<Self>>,
future: impl SideChannelFuture + 'a,
) -> NonNull<Self> {
let is_done = Cell::new(false);
let this = FrameInner {
is_done,
return_to,
future,
};
let this: *mut FrameInner<_> = alloc.new_raw(this);
// Safety: `Box::into_raw` is always properly aligned and non-null
let this: NonNull<FrameInner<_>> = unsafe { NonNull::new_unchecked(this) };
// Unsizing coercion
let this: NonNull<FrameInner<dyn SideChannelFuture + 'a>> = this;
// Coerce to `'static` lifetime
// Safety: only the top future on the stack will ever be polled, and
// they will be dropped in reverse order (top to bottom). Their actual
// lifetime is therefore smaller than `'a`.
let this: NonNull<Self> = unsafe { std::mem::transmute(this) };
this
}
/// Get the [`SideChannelFuture`] stored inside this frame.
pub fn as_future(&self) -> &(dyn SideChannelFuture + 'static) {
&self.future
}
/// Returns `true` if the future stored inside this frame has been polled
/// to completion.
pub fn is_done(&self) -> bool {
self.is_done.get()
}
}
/// Construct a placeholder frame that panics if invoked.
fn placeholder(alloc: &Alloc) -> NonNull<Frame> {
Frame::new(alloc, None, Placeholder)
}
struct Placeholder;
impl Future for Placeholder {
type Output = ();
fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<()> {
panic!("placeholder must not be invoked")
}
}
impl SideChannelFuture for Placeholder {
unsafe fn get_result_pointer(&self) -> NonNull<Cell<Option<()>>> {
unreachable!("write_result_pointer must not be called on the placeholder future")
}
}
struct Bottom<F>(F);
impl<F> Future for Bottom<F>
where
F: Future<Output = ()>,
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// Safety: the inner future is only ever accessed through a pin.
let inner = unsafe { self.map_unchecked_mut(|this| &mut this.0) };
inner.poll(cx)
}
}
impl<F> SideChannelFuture for Bottom<F>
where
F: Future<Output = ()>,
{
unsafe fn get_result_pointer(&self) -> NonNull<Cell<Option<()>>> {
unreachable!("write_result_pointer must not be called on the bottom future")
}
}