hydro_lang/embedded.rs
1//! Runtime helpers for driving embedded Hydro dataflow graphs from handwritten Rust.
2//!
3//! The embedded backend (`hydro_lang::compile::embedded`, available with the `build`
4//! feature) generates one function per location. Each generated function takes
5//! [`futures::Stream`] inputs and `FnMut` output callbacks, and returns a DFIR graph that
6//! *borrows* from them. Storing the graph together with its inputs and outputs in a single
7//! owned value therefore requires a self-referential struct, which cannot be expressed in
8//! safe Rust. This module provides the building blocks for such a holder, with all `unsafe`
9//! confined to this module so that downstream integration code is safe:
10//!
11//! - [`InputBuffer`] / [`InputStream`]: an address-stable queue that backs a
12//! [`futures::Stream`] input parameter of a generated function.
13//! - [`CallbackSlot`]: an address-stable slot holding a caller-provided `FnMut(T)` for the
14//! duration of a run, backing an output (or network-out) parameter.
15//! - [`AliasedBox`] / [`OwnedErasedBox`]: owned heap allocations that, unlike [`Box`], may
16//! be moved while raw pointers into their contents exist.
17//! - [`DfirRunnable`]: object-safe erasure for the unnameable `Dfir<impl TickClosure>`.
18//! - [`embedded_flow!`](crate::embedded_flow): a macro assembling the above into a safe
19//! holder struct.
20//!
21//! The same primitives support local and networked channels alike: a network-out is just a
22//! `CallbackSlot<(TaglessMemberId, Bytes)>` and a network-in is just an
23//! `InputBuffer<Result<(TaglessMemberId, BytesMut), io::Error>>`.
24//!
25//! # Threading
26//!
27//! All of these types are single-threaded: they are `!Sync`, and `!Send` wherever a raw
28//! pointer may reference them. A holder assembled from them must live and run on one thread.
29//!
30//! # `no_std`
31//!
32//! The core primitives ([`InputBuffer`], [`CallbackSlot`]) only use `core` items, have
33//! `const fn new` constructors, and communicate address stability through [`Pin`], so they
34//! are compatible with future `no_std` targets using statically-allocated (fixed-size)
35//! storage. Only the owning-allocation helpers ([`AliasedBox`], [`OwnedErasedBox`]) and the
36//! [`embedded_flow!`](crate::embedded_flow) macro require `alloc`.
37
38use core::cell::{Cell, UnsafeCell};
39use core::fmt;
40use core::marker::{PhantomData, PhantomPinned};
41use core::ops::Deref;
42use core::pin::Pin;
43use core::ptr::NonNull;
44use core::task::{Context, Poll, Waker};
45use std::collections::VecDeque;
46
47use dfir_rs::scheduled::context::{Dfir, TickClosure};
48
49pub use crate::embedded_flow;
50
51/// Implementation details of [`embedded_flow!`](crate::embedded_flow). Not stable API.
52///
53/// These re-exports exist so that the macro expansion does not depend on the invoking
54/// crate's prelude, and so a future `no_std` version of `hydro_lang` can swap `std` paths
55/// for `alloc` paths without changing (or breaking) the macro itself.
56#[doc(hidden)]
57pub mod __private {
58 pub use core::cell::RefCell;
59 pub use std::boxed::Box;
60 pub use std::vec::Vec;
61}
62
63/// Object-safe trait for running a DFIR flow synchronously.
64///
65/// Generated embedded functions return `Dfir<impl TickClosure + 'a>`, whose type parameter
66/// cannot be named by the caller. Boxing the graph as `Box<dyn DfirRunnable>` erases that
67/// type with a single virtual call per [`run_available_sync`](Self::run_available_sync)
68/// invocation. (This is cheaper than `Dfir::into_erased`, which boxes a future on every
69/// tick, and also works for tick closures that only implement `TickClosure`.)
70pub trait DfirRunnable {
71 /// Runs ticks as long as work is available, then returns.
72 ///
73 /// See `Dfir::run_available_sync`; panics if a tick yields asynchronously.
74 fn run_available_sync(&mut self);
75}
76
77impl<Tick: TickClosure> DfirRunnable for Dfir<Tick> {
78 fn run_available_sync(&mut self) {
79 Dfir::run_available_sync(self)
80 }
81}
82
83/// An owned heap allocation, like [`Box`], but without `Box`'s unique-aliasing guarantee.
84///
85/// Moving a `Box` asserts that it is the *only* pointer to its contents (`noalias`),
86/// invalidating raw pointers previously derived from them (e.g. under Stacked Borrows).
87/// `AliasedBox` stores only a raw pointer, so it may be freely moved — such as into the
88/// holder struct built by [`embedded_flow!`](crate::embedded_flow) — while [`InputStream`]s
89/// and sink closures hold `NonNull` pointers into its contents.
90///
91/// The contents are never moved and never exposed as `&mut` (only [`Deref`]), so their
92/// address is stable from construction until drop; this is what makes
93/// [`as_pin`](Self::as_pin) sound.
94pub struct AliasedBox<T> {
95 ptr: NonNull<T>,
96 /// Declare ownership of a `T` for drop-check purposes.
97 _owns: PhantomData<T>,
98}
99
100impl<T> AliasedBox<T> {
101 /// Heap-allocates `value`.
102 pub fn new(value: T) -> Self {
103 // SAFETY: `Box::into_raw` never returns null.
104 let ptr = unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(value))) };
105 Self {
106 ptr,
107 _owns: PhantomData,
108 }
109 }
110
111 /// Returns the contents as a pinned reference.
112 ///
113 /// This is safe because the contents are heap-allocated, never moved, and never
114 /// exposed mutably, so they remain at the same address until `self` is dropped
115 /// (at which point their destructor runs, upholding the [`Pin`] drop guarantee).
116 pub fn as_pin(&self) -> Pin<&T> {
117 // SAFETY: see doc comment above.
118 unsafe { Pin::new_unchecked(self.ptr.as_ref()) }
119 }
120}
121
122impl<T> Deref for AliasedBox<T> {
123 type Target = T;
124
125 fn deref(&self) -> &T {
126 // SAFETY: allocated in `new`, freed only in `drop`.
127 unsafe { self.ptr.as_ref() }
128 }
129}
130
131impl<T> Drop for AliasedBox<T> {
132 fn drop(&mut self) {
133 // SAFETY: `ptr` came from `Box::into_raw` in `new` and is dropped exactly once.
134 drop(unsafe { Box::from_raw(self.ptr.as_ptr()) });
135 }
136}
137
138impl<T: fmt::Debug> fmt::Debug for AliasedBox<T> {
139 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
140 T::fmt(self, f)
141 }
142}
143
144/// A type-erased owned heap allocation, dropped (running `T`'s destructor) when this value
145/// is dropped.
146///
147/// [`new`](Self::new) hands back a `&'static mut T` into the allocation, which can be lent
148/// to a generated embedded function while the `OwnedErasedBox` — whose type does not mention
149/// `T`, useful when `T` is unnameable (e.g. a struct of `impl FnMut` sinks) — is stored
150/// alongside the flow to keep the allocation alive. Like [`AliasedBox`], only a raw pointer
151/// is stored, so moving this value does not invalidate the outstanding reference.
152pub struct OwnedErasedBox {
153 ptr: NonNull<u8>,
154 /// Drops the allocation as its original `T`.
155 drop_fn: unsafe fn(NonNull<u8>),
156}
157
158impl OwnedErasedBox {
159 /// Heap-allocates `value`, returning the type-erased owner along with a
160 /// `&'static mut T` to the allocation.
161 ///
162 /// # Safety
163 ///
164 /// The returned reference (and anything derived from it) must not be used after the
165 /// returned `OwnedErasedBox` is dropped, despite the `'static` lifetime.
166 pub unsafe fn new<T>(value: T) -> (Self, &'static mut T) {
167 /// SAFETY: must be called at most once, with a pointer originating from
168 /// `Box::into_raw::<T>`.
169 unsafe fn drop_impl<T>(ptr: NonNull<u8>) {
170 // SAFETY: per this function's contract.
171 drop(unsafe { Box::from_raw(ptr.cast::<T>().as_ptr()) });
172 }
173
174 let raw: *mut T = Box::into_raw(Box::new(value));
175 // SAFETY: `Box::into_raw` never returns null.
176 let ptr = unsafe { NonNull::new_unchecked(raw) };
177 (
178 Self {
179 ptr: ptr.cast(),
180 drop_fn: drop_impl::<T>,
181 },
182 // SAFETY: fresh exclusive allocation, freed only by `drop_fn`; the caller
183 // promises not to use the reference after that.
184 unsafe { &mut *raw },
185 )
186 }
187}
188
189impl Drop for OwnedErasedBox {
190 fn drop(&mut self) {
191 // SAFETY: `ptr` and `drop_fn` were created as a pair in `new`, and this runs once.
192 unsafe { (self.drop_fn)(self.ptr) }
193 }
194}
195
196impl fmt::Debug for OwnedErasedBox {
197 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
198 f.debug_struct("OwnedErasedBox").finish_non_exhaustive()
199 }
200}
201
202/// An address-stable FIFO queue that backs a [`futures::Stream`] input parameter of a
203/// generated embedded function.
204///
205/// The queue is filled from application code via [`push`](Self::push) and drained by the
206/// dataflow graph through an [`InputStream`] created with [`stream`](Self::stream). The
207/// stream never terminates: while the queue is empty it returns [`Poll::Pending`], and it
208/// is woken by the next `push`.
209pub struct InputBuffer<T> {
210 /// Queue of pending items.
211 ///
212 /// Interior mutability is required because [`push`](Self::push) is called through
213 /// `&self` (from application code) while an [`InputStream`] holding a raw pointer to
214 /// this buffer is owned by the flow. Every access is scoped to a single expression
215 /// that does not call back into user code, so overlapping borrows cannot occur; all
216 /// accesses happen on one thread because `InputBuffer` is `!Sync` and `!Send`.
217 buf: UnsafeCell<VecDeque<T>>,
218 /// The waker of the task that last polled an empty [`InputStream`], woken on `push`.
219 /// Same access discipline as `buf`, except that foreign waker code is only ever run
220 /// *after* ending the access (by moving the waker out of the cell).
221 waker: UnsafeCell<Option<Waker>>,
222 /// `!Unpin`: [`InputStream`] holds a raw pointer to this buffer, so it must not move.
223 _pin: PhantomPinned,
224 /// `!Send`: an [`InputStream`] on the original thread may access this buffer without
225 /// synchronization, so the buffer must not be pushed to from another thread.
226 _not_send: PhantomData<*mut ()>,
227}
228
229impl<T> InputBuffer<T> {
230 /// Creates an empty buffer. Does not allocate.
231 pub const fn new() -> Self {
232 Self {
233 buf: UnsafeCell::new(VecDeque::new()),
234 waker: UnsafeCell::new(None),
235 _pin: PhantomPinned,
236 _not_send: PhantomData,
237 }
238 }
239
240 /// Pushes an item onto the queue and wakes the [`InputStream`] if it is waiting.
241 ///
242 /// Note that the item is only *queued*; it is processed the next time the flow runs.
243 pub fn push(&self, item: T) {
244 // SAFETY: scoped, single-threaded access; see `Self::buf`.
245 unsafe { (*self.buf.get()).push_back(item) };
246 // Move the waker out of the cell *before* invoking it, so that no access to the
247 // cell is in progress while running arbitrary waker code (which could re-enter
248 // `push`). Consuming the waker also coalesces wakes: subsequent pushes before the
249 // next poll skip the wake, and `poll_next` re-registers on `Pending`.
250 // SAFETY: scoped, single-threaded access; see `Self::waker`.
251 let waker = unsafe { (*self.waker.get()).take() };
252 if let Some(waker) = waker {
253 waker.wake();
254 }
255 }
256
257 /// Returns the number of items currently queued.
258 pub fn len(&self) -> usize {
259 // SAFETY: scoped, single-threaded access; see `Self::buf`.
260 unsafe { (*self.buf.get()).len() }
261 }
262
263 /// Returns `true` if no items are currently queued.
264 pub fn is_empty(&self) -> bool {
265 self.len() == 0
266 }
267
268 /// Creates a never-terminating [`futures::Stream`] that drains this buffer, for
269 /// passing to a generated embedded function.
270 ///
271 /// Multiple streams over one buffer are allowed (though rarely useful): items are
272 /// delivered to whichever stream polls first, and only the most recent waker is woken.
273 ///
274 /// # Safety
275 ///
276 /// The buffer must outlive the returned stream. (The buffer cannot move while the
277 /// stream exists, which is guaranteed by `self` being pinned, and the stream cannot be
278 /// sent to another thread, as it is `!Send`.)
279 pub unsafe fn stream(self: Pin<&Self>) -> InputStream<T> {
280 InputStream {
281 ptr: NonNull::from(self.get_ref()),
282 }
283 }
284}
285
286impl<T> Default for InputBuffer<T> {
287 fn default() -> Self {
288 Self::new()
289 }
290}
291
292impl<T> fmt::Debug for InputBuffer<T> {
293 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
294 f.debug_struct("InputBuffer")
295 .field("len", &self.len())
296 .finish_non_exhaustive()
297 }
298}
299
300/// A never-terminating [`futures::Stream`] that drains an [`InputBuffer`].
301///
302/// Created via [`InputBuffer::stream`]; see the safety requirements there.
303pub struct InputStream<T> {
304 ptr: NonNull<InputBuffer<T>>,
305}
306
307/// The *buffer* is what must stay pinned; the stream itself moves freely.
308impl<T> Unpin for InputStream<T> {}
309
310impl<T> futures::Stream for InputStream<T> {
311 type Item = T;
312
313 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
314 // SAFETY: the `InputBuffer::stream` contract guarantees the buffer is still alive.
315 let buffer = unsafe { self.ptr.as_ref() };
316
317 // Fast path: an item is available, no waker bookkeeping needed.
318 // SAFETY: scoped, single-threaded access; see `InputBuffer::buf`.
319 if let Some(item) = unsafe { (*buffer.buf.get()).pop_front() } {
320 return Poll::Ready(Some(item));
321 }
322
323 // Register (or refresh) the waker. The old waker is moved out of the cell first so
324 // that no access is in progress while `clone_from` runs foreign waker code.
325 // (`Waker::clone_from` skips the clone when the old waker `will_wake` the new one.)
326 // SAFETY: scoped, single-threaded accesses; see `InputBuffer::waker`.
327 let mut waker = unsafe { (*buffer.waker.get()).take() };
328 match &mut waker {
329 Some(waker) => waker.clone_from(cx.waker()),
330 None => waker = Some(cx.waker().clone()),
331 }
332 // SAFETY: as above.
333 unsafe { *buffer.waker.get() = waker };
334
335 // Re-check the queue: a (pathological) re-entrant `push` from inside the waker
336 // clone above would have found the waker cell empty and not woken anyone, so
337 // returning `Pending` without this check could lose a wakeup.
338 // SAFETY: scoped, single-threaded access; see `InputBuffer::buf`.
339 match unsafe { (*buffer.buf.get()).pop_front() } {
340 Some(item) => Poll::Ready(Some(item)),
341 None => Poll::Pending,
342 }
343 }
344
345 fn size_hint(&self) -> (usize, Option<usize>) {
346 // SAFETY: the `InputBuffer::stream` contract guarantees the buffer is still alive.
347 let buffered = unsafe { self.ptr.as_ref() }.len();
348 (buffered, None) // Never terminates, so there is no upper bound.
349 }
350}
351
352impl<T> fmt::Debug for InputStream<T> {
353 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
354 f.debug_struct("InputStream").finish_non_exhaustive()
355 }
356}
357
358/// An address-stable slot holding a caller-provided callback for the duration of a run,
359/// backing an output (or network-out) parameter of a generated embedded function.
360///
361/// The flow side holds a sink closure (created via [`sink`](Self::sink)) that forwards each
362/// emitted item to whatever callback is currently installed. The application side
363/// temporarily installs a callback around each run via [`invoke_with`](Self::invoke_with):
364///
365/// ```ignore
366/// slot.invoke_with(&mut |item| results.push(item), || flow.run_available_sync());
367/// ```
368///
369/// Items emitted while no callback is installed are **silently dropped**.
370pub struct CallbackSlot<T: 'static> {
371 /// The currently-installed callback, if any. A [`Cell`] (the pointer is `Copy`), so no
372 /// borrow of the slot is ever held across a call into user code.
373 cell: Cell<Option<NonNull<dyn FnMut(T)>>>,
374 /// `!Unpin`: sink closures hold a raw pointer to this slot, so it must not move.
375 _pin: PhantomPinned,
376}
377
378impl<T: 'static> CallbackSlot<T> {
379 /// Creates a slot with no callback installed.
380 pub const fn new() -> Self {
381 Self {
382 cell: Cell::new(None),
383 _pin: PhantomPinned,
384 }
385 }
386
387 /// Creates a sink closure that forwards each item to the currently-installed callback,
388 /// for passing to a generated embedded function.
389 ///
390 /// # Safety
391 ///
392 /// The slot must outlive the returned closure. (The slot cannot move while the closure
393 /// exists, which is guaranteed by `self` being pinned, and the closure cannot be sent
394 /// to another thread, as it is `!Send`.)
395 pub unsafe fn sink(self: Pin<&Self>) -> impl FnMut(T) + use<T> {
396 let ptr = NonNull::from(self.get_ref());
397 move |item| {
398 // SAFETY: the `sink` contract guarantees the slot is still alive.
399 let slot = unsafe { ptr.as_ref() };
400 slot.invoke(item);
401 }
402 }
403
404 /// Invokes the currently-installed callback with `item`, or drops `item` if no
405 /// callback is installed.
406 ///
407 /// The callback is moved out of the slot while it runs (and restored afterwards, even
408 /// on panic), so a re-entrant invocation of the same slot finds it empty and drops its
409 /// item, rather than aliasing the already-running callback.
410 pub fn invoke(&self, item: T) {
411 let Some(mut callback) = self.cell.take() else {
412 return;
413 };
414 let restore = RestoreOnDrop {
415 slot: self,
416 prev: Some(callback),
417 };
418 // SAFETY: the pointer was installed by `invoke_with`, which keeps the callback
419 // alive (and its `&mut` borrow active) until it uninstalls the pointer on exit.
420 // Taking the pointer out of the cell above makes this `&mut` reborrow exclusive.
421 unsafe { (callback.as_mut())(item) };
422 drop(restore);
423 }
424
425 /// Installs `callback` for the duration of `f`, then restores whatever was previously
426 /// installed — even if `f` panics. Items the flow emits to this slot while `f` runs
427 /// are passed to `callback`.
428 ///
429 /// Calls may be nested (e.g. to install callbacks on several slots around a single
430 /// run); the innermost installation wins for its duration.
431 pub fn invoke_with<R>(&self, callback: &mut dyn FnMut(T), f: impl FnOnce() -> R) -> R {
432 let ptr = NonNull::from(callback);
433 // SAFETY(transmute): erases the (unnameable) lifetime of the trait object; this
434 // only changes the type-level lifetime, not the pointer. The pointer is only
435 // dereferenced by `invoke` while installed, and the guard below uninstalls it
436 // before `invoke_with` returns or unwinds — i.e., strictly within the lifetime of
437 // the `&mut callback` borrow.
438 let ptr: NonNull<dyn FnMut(T) + 'static> = unsafe { core::mem::transmute(ptr) };
439 let _restore = RestoreOnDrop {
440 slot: self,
441 prev: self.cell.replace(Some(ptr)),
442 };
443 f()
444 }
445}
446
447impl<T: 'static> Default for CallbackSlot<T> {
448 fn default() -> Self {
449 Self::new()
450 }
451}
452
453impl<T: 'static> fmt::Debug for CallbackSlot<T> {
454 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
455 f.debug_struct("CallbackSlot")
456 .field("installed", &self.cell.get().is_some())
457 .finish()
458 }
459}
460
461/// Restores a [`CallbackSlot`]'s previous contents on drop (including during unwinding).
462struct RestoreOnDrop<'a, T: 'static> {
463 slot: &'a CallbackSlot<T>,
464 prev: Option<NonNull<dyn FnMut(T)>>,
465}
466
467impl<T: 'static> Drop for RestoreOnDrop<'_, T> {
468 fn drop(&mut self) {
469 self.slot.cell.set(self.prev);
470 }
471}
472
473/// Implementation detail of [`embedded_flow!`](crate::embedded_flow): nests
474/// [`CallbackSlot::invoke_with`] calls around a body expression.
475#[doc(hidden)]
476#[macro_export]
477macro_rules! __embedded_invoke_nested {
478 ((), $body:expr) => {
479 $body
480 };
481 (($slot:expr => $cb:expr, $($rest:tt)*), $body:expr) => {
482 $slot.invoke_with($cb, move || $crate::__embedded_invoke_nested!(($($rest)*), $body))
483 };
484}
485
486/// Declares a holder struct that owns a generated embedded DFIR flow (see
487/// `hydro_lang::compile::embedded`) *together with* the input buffers and output callback
488/// slots it borrows from, exposing an entirely safe API.
489///
490/// The generated struct has:
491/// - `fn new(...)` — constructs the buffers/slots and the flow (extra `build` parameters
492/// become parameters of `new`; a `self_id` section adds a leading `self_id` parameter).
493/// - one accessor per input / network-in / membership entry returning
494/// [`&InputBuffer<_>`](crate::embedded::InputBuffer), to [`push`](crate::embedded::InputBuffer::push) items into;
495/// - one accessor per output / network-out entry returning
496/// [`&CallbackSlot<_>`](crate::embedded::CallbackSlot);
497/// - `fn run(&self)` — runs the flow until no more work is available (outputs emitted with
498/// no callback installed are dropped);
499/// - `fn run_with(&self, ...)` — runs the flow with a `&mut dyn FnMut(_)` callback
500/// installed for each output, then each network-out, in declaration order.
501///
502/// # Sections
503///
504/// All sections are optional except `build`, but must appear in the order shown. Each
505/// `Path as binder` clause names a struct generated by the embedded backend (e.g.
506/// `my_fn::EmbeddedOutputs`) and binds the value passed to the generated function to a
507/// local usable inside `build`. Input names are bound directly as stream locals.
508///
509/// ```ignore
510/// hydro_lang::embedded_flow! {
511/// /// Holder for my flow.
512/// pub struct MyFlow {
513/// // For cluster locations: the member id of this instance.
514/// self_id(my_id: TaglessMemberId);
515/// // Cluster membership streams; field names are the cluster fn names.
516/// membership(my_fn::EmbeddedMembershipStreams as memberships) { other_cluster }
517/// inputs { events: MyEvent }
518/// network_in(my_fn::EmbeddedNetworkIn as net_in) {
519/// requests: Result<BytesMut, std::io::Error>,
520/// }
521/// outputs(my_fn::EmbeddedOutputs as outputs) { results: MyResult }
522/// network_out(my_fn::EmbeddedNetworkOut as net_out) { responses: Bytes }
523/// build(some_config: usize) {
524/// my_module::my_fn(my_id, memberships, events, outputs, net_in, net_out)
525/// }
526/// }
527/// }
528///
529/// let flow = MyFlow::new(member_id, 42);
530/// flow.events().push(event);
531/// flow.run_with(
532/// &mut |result| println!("{result:?}"),
533/// &mut |response| network.send(response),
534/// );
535/// ```
536///
537/// Item types must match the generated function's parameter types exactly: `network_in`
538/// items are `Result<BytesMut, io::Error>` (tagged: `Result<(TaglessMemberId, BytesMut),
539/// io::Error>`) unless the channel uses external serialization, `network_out` items are
540/// `Bytes` (tagged: `(TaglessMemberId, Bytes)`), and membership items are always
541/// `(TaglessMemberId, MembershipEvent)`.
542///
543/// # Caveats
544///
545/// - The holder is single-threaded (`!Send`); construct and run it on one thread.
546/// - `run` / `run_with` panic if called re-entrantly from within an output callback
547/// (output callbacks *may* push new inputs, which are processed later in the same run).
548/// - Entry names are used as field and accessor names, so they must be distinct and must
549/// not be named `flow` or `run`; `build` parameters must not be named `self_id`,
550/// `membership`, `inputs`, `network_in`, `outputs`, or `network_out`.
551/// - The fields of the generated struct are self-referential; code in the defining module
552/// must never mutate or replace them (use the generated accessors instead).
553#[macro_export]
554macro_rules! embedded_flow {
555 (
556 $(#[$struct_meta:meta])*
557 $vis:vis struct $name:ident {
558 $( self_id($self_ref:ident : $self_ty:ty); )?
559 $( membership($mem_path:path as $membership_ref:ident) {
560 $($mem_name:ident),* $(,)?
561 } )?
562 $( inputs { $($in_name:ident : $in_ty:ty),* $(,)? } )?
563 $( network_in($nin_path:path as $network_in_ref:ident) {
564 $($nin_name:ident : $nin_ty:ty),* $(,)?
565 } )?
566 $( outputs($outputs_path:path as $outputs_ref:ident) {
567 $($out_name:ident : $out_ty:ty),* $(,)?
568 } )?
569 $( network_out($nout_path:path as $network_out_ref:ident) {
570 $($nout_name:ident : $nout_ty:ty),* $(,)?
571 } )?
572 build($($param:ident : $param_ty:ty),* $(,)?) $build_body:block
573 }
574 ) => {
575 $(#[$struct_meta])*
576 $vis struct $name {
577 /// The DFIR flow. Declared first so it is dropped first: it borrows from all
578 /// of the fields below. Wrapped in a `RefCell` so the flow can be run through
579 /// a shared `&self` borrow, letting output callbacks re-enter the holder to
580 /// push more input during the same run.
581 flow: $crate::embedded::__private::RefCell<
582 $crate::embedded::__private::Box<dyn $crate::embedded::DfirRunnable>,
583 >,
584 /// Type-erased allocations lent to the flow (outputs / network-out structs and
585 /// the self id); dropped after the flow.
586 _owned: $crate::embedded::__private::Vec<$crate::embedded::OwnedErasedBox>,
587 $($(
588 $in_name: $crate::embedded::AliasedBox<$crate::embedded::InputBuffer<$in_ty>>,
589 )*)?
590 $($(
591 $nin_name: $crate::embedded::AliasedBox<$crate::embedded::InputBuffer<$nin_ty>>,
592 )*)?
593 $($(
594 $mem_name: $crate::embedded::AliasedBox<$crate::embedded::InputBuffer<(
595 $crate::location::member_id::TaglessMemberId,
596 $crate::location::MembershipEvent,
597 )>>,
598 )*)?
599 $($(
600 $out_name: $crate::embedded::AliasedBox<$crate::embedded::CallbackSlot<$out_ty>>,
601 )*)?
602 $($(
603 $nout_name: $crate::embedded::AliasedBox<$crate::embedded::CallbackSlot<$nout_ty>>,
604 )*)?
605 }
606
607 impl ::core::fmt::Debug for $name {
608 fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
609 f.debug_struct(::core::stringify!($name)).finish_non_exhaustive()
610 }
611 }
612
613 impl $name {
614 /// Constructs the flow along with the buffers and slots it reads from and
615 /// writes to. No ticks are run yet; call
616 #[doc = ::core::concat!("[`run`](", ::core::stringify!($name), "::run) or [`run_with`](", ::core::stringify!($name), "::run_with) to process queued inputs.")]
617 #[allow(
618 unused_unsafe,
619 unused_mut,
620 reason = "usage depends on which macro sections are present"
621 )]
622 $vis fn new($( self_id: $self_ty, )? $($param: $param_ty),*) -> Self {
623 $($(
624 let $in_name = $crate::embedded::AliasedBox::new(
625 $crate::embedded::InputBuffer::<$in_ty>::new(),
626 );
627 )*)?
628 $($(
629 let $nin_name = $crate::embedded::AliasedBox::new(
630 $crate::embedded::InputBuffer::<$nin_ty>::new(),
631 );
632 )*)?
633 $($(
634 let $mem_name = $crate::embedded::AliasedBox::new(
635 $crate::embedded::InputBuffer::<(
636 $crate::location::member_id::TaglessMemberId,
637 $crate::location::MembershipEvent,
638 )>::new(),
639 );
640 )*)?
641 $($(
642 let $out_name = $crate::embedded::AliasedBox::new(
643 $crate::embedded::CallbackSlot::<$out_ty>::new(),
644 );
645 )*)?
646 $($(
647 let $nout_name = $crate::embedded::AliasedBox::new(
648 $crate::embedded::CallbackSlot::<$nout_ty>::new(),
649 );
650 )*)?
651
652 let mut __owned: $crate::embedded::__private::Vec<$crate::embedded::OwnedErasedBox> =
653 $crate::embedded::__private::Vec::new();
654
655 // SAFETY: the buffers and slots are heap allocations owned by the
656 // `AliasedBox` fields of the returned struct (which, unlike `Box`, may be
657 // moved while the flow holds pointers into their contents), and the
658 // outputs / network-out structs and self id are owned by `_owned`. Field
659 // declaration order guarantees the flow is dropped before all of them.
660 // Everything stays on a single thread: none of these types are `Send`.
661 let flow = unsafe {
662 $(
663 let (__self_id_owner, __self_id_mut) =
664 $crate::embedded::OwnedErasedBox::new(self_id);
665 __owned.push(__self_id_owner);
666 let $self_ref: &'static $self_ty = &*__self_id_mut;
667 )?
668 $($(
669 let $in_name = $crate::embedded::InputBuffer::stream($in_name.as_pin());
670 )*)?
671 $(
672 $(
673 let $mem_name =
674 $crate::embedded::InputBuffer::stream($mem_name.as_pin());
675 )*
676 use $mem_path as __EmbeddedMembership;
677 let $membership_ref = __EmbeddedMembership { $($mem_name),* };
678 )?
679 $(
680 $(
681 let $nin_name =
682 $crate::embedded::InputBuffer::stream($nin_name.as_pin());
683 )*
684 use $nin_path as __EmbeddedNetworkIn;
685 let $network_in_ref = __EmbeddedNetworkIn { $($nin_name),* };
686 )?
687 $(
688 $(
689 let $out_name = $crate::embedded::CallbackSlot::sink($out_name.as_pin());
690 )*
691 use $outputs_path as __EmbeddedOutputs;
692 let (__outputs_owner, $outputs_ref) =
693 $crate::embedded::OwnedErasedBox::new(
694 __EmbeddedOutputs { $($out_name),* },
695 );
696 __owned.push(__outputs_owner);
697 )?
698 $(
699 $(
700 let $nout_name =
701 $crate::embedded::CallbackSlot::sink($nout_name.as_pin());
702 )*
703 use $nout_path as __EmbeddedNetworkOut;
704 let (__network_out_owner, $network_out_ref) =
705 $crate::embedded::OwnedErasedBox::new(
706 __EmbeddedNetworkOut { $($nout_name),* },
707 );
708 __owned.push(__network_out_owner);
709 )?
710
711 let __dfir = $build_body;
712 let __flow: $crate::embedded::__private::Box<dyn $crate::embedded::DfirRunnable> =
713 $crate::embedded::__private::Box::new(__dfir);
714 __flow
715 };
716
717 Self {
718 flow: $crate::embedded::__private::RefCell::new(flow),
719 _owned: __owned,
720 $($( $in_name, )*)?
721 $($( $nin_name, )*)?
722 $($( $mem_name, )*)?
723 $($( $out_name, )*)?
724 $($( $nout_name, )*)?
725 }
726 }
727
728 /// Runs the flow until no more work is available.
729 ///
730 /// Items emitted to outputs with no callback installed are dropped; use
731 #[doc = ::core::concat!("[`run_with`](", ::core::stringify!($name), "::run_with) (or `CallbackSlot::invoke_with`) to capture outputs.")]
732 ///
733 /// # Panics
734 ///
735 /// Panics if called re-entrantly from within an output callback.
736 #[inline]
737 $vis fn run(&self) {
738 self.flow.borrow_mut().run_available_sync();
739 }
740
741 /// Runs the flow with the given callbacks installed: one per `outputs` entry,
742 /// then one per `network_out` entry, in declaration order.
743 ///
744 /// # Panics
745 ///
746 /// Panics if called re-entrantly from within an output callback.
747 $vis fn run_with(
748 &self
749 $($(, $out_name: &mut dyn ::core::ops::FnMut($out_ty))*)?
750 $($(, $nout_name: &mut dyn ::core::ops::FnMut($nout_ty))*)?
751 ) {
752 $crate::__embedded_invoke_nested!(
753 (
754 $($( self.$out_name => $out_name, )*)?
755 $($( self.$nout_name => $nout_name, )*)?
756 ),
757 self.run()
758 )
759 }
760
761 $($(
762 /// Input buffer: push items here, then call
763 #[doc = ::core::concat!("[`run_with`](", ::core::stringify!($name), "::run_with) to process them.")]
764 $vis fn $in_name(&self) -> &$crate::embedded::InputBuffer<$in_ty> {
765 &self.$in_name
766 }
767 )*)?
768 $($(
769 /// Network-in buffer: push received network messages here, then call
770 #[doc = ::core::concat!("[`run_with`](", ::core::stringify!($name), "::run_with) to process them.")]
771 $vis fn $nin_name(&self) -> &$crate::embedded::InputBuffer<$nin_ty> {
772 &self.$nin_name
773 }
774 )*)?
775 $($(
776 /// Cluster membership buffer: push membership events here, then call
777 #[doc = ::core::concat!("[`run_with`](", ::core::stringify!($name), "::run_with) to process them.")]
778 $vis fn $mem_name(
779 &self,
780 ) -> &$crate::embedded::InputBuffer<(
781 $crate::location::member_id::TaglessMemberId,
782 $crate::location::MembershipEvent,
783 )> {
784 &self.$mem_name
785 }
786 )*)?
787 $($(
788 /// Output slot: install a callback via `invoke_with` around a run, or use
789 #[doc = ::core::concat!("[`run_with`](", ::core::stringify!($name), "::run_with).")]
790 $vis fn $out_name(&self) -> &$crate::embedded::CallbackSlot<$out_ty> {
791 &self.$out_name
792 }
793 )*)?
794 $($(
795 /// Network-out slot: install a callback via `invoke_with` around a run, or
796 #[doc = ::core::concat!("use [`run_with`](", ::core::stringify!($name), "::run_with).")]
797 $vis fn $nout_name(&self) -> &$crate::embedded::CallbackSlot<$nout_ty> {
798 &self.$nout_name
799 }
800 )*)?
801 }
802 };
803}
804
805#[cfg(test)]
806mod tests {
807 use core::sync::atomic::{AtomicUsize, Ordering};
808 use core::task::Context;
809 use std::rc::Rc;
810 use std::sync::Arc;
811 use std::task::Wake;
812
813 use futures::Stream;
814
815 use super::*;
816
817 /// Waker that counts how many times it has been woken.
818 struct CountingWaker(AtomicUsize);
819
820 impl Wake for CountingWaker {
821 fn wake(self: Arc<Self>) {
822 self.wake_by_ref();
823 }
824
825 fn wake_by_ref(self: &Arc<Self>) {
826 self.0.fetch_add(1, Ordering::Relaxed);
827 }
828 }
829
830 fn poll_next<T>(stream: &mut InputStream<T>, cx: &mut Context<'_>) -> Poll<Option<T>> {
831 Pin::new(stream).poll_next(cx)
832 }
833
834 #[test]
835 fn input_buffer_delivers_in_order_and_wakes() {
836 let buffer = AliasedBox::new(InputBuffer::new());
837 // SAFETY: `buffer` outlives `stream`.
838 let mut stream = unsafe { InputBuffer::stream(buffer.as_pin()) };
839
840 let wake_count = Arc::new(CountingWaker(AtomicUsize::new(0)));
841 let waker = Waker::from(Arc::clone(&wake_count));
842 let mut cx = Context::from_waker(&waker);
843
844 // Empty: pending, waker registered.
845 assert_eq!(poll_next(&mut stream, &mut cx), Poll::Pending);
846 assert_eq!(wake_count.0.load(Ordering::Relaxed), 0);
847
848 // Push wakes exactly once (coalesced), even for multiple pushes.
849 buffer.push(1);
850 buffer.push(2);
851 assert_eq!(wake_count.0.load(Ordering::Relaxed), 1);
852 assert_eq!(buffer.len(), 2);
853 assert_eq!(stream.size_hint(), (2, None));
854
855 // FIFO delivery.
856 assert_eq!(poll_next(&mut stream, &mut cx), Poll::Ready(Some(1)));
857 assert_eq!(poll_next(&mut stream, &mut cx), Poll::Ready(Some(2)));
858 assert_eq!(poll_next(&mut stream, &mut cx), Poll::Pending);
859 assert!(buffer.is_empty());
860
861 // Waker was re-registered after going pending again.
862 buffer.push(3);
863 assert_eq!(wake_count.0.load(Ordering::Relaxed), 2);
864 }
865
866 #[test]
867 fn input_buffer_stable_across_moves() {
868 // Moving the `AliasedBox` (e.g. into a struct) must not invalidate the stream.
869 let buffer = AliasedBox::new(InputBuffer::new());
870 // SAFETY: `buffer` outlives `stream` (both live to the end of this function).
871 let mut stream = unsafe { InputBuffer::stream(buffer.as_pin()) };
872
873 let moved = [buffer];
874 moved[0].push("hello");
875
876 let waker = Waker::noop();
877 let mut cx = Context::from_waker(waker);
878 assert_eq!(poll_next(&mut stream, &mut cx), Poll::Ready(Some("hello")));
879 }
880
881 #[test]
882 fn callback_slot_invoke_with_installs_and_restores() {
883 let slot: CallbackSlot<u32> = CallbackSlot::new();
884
885 // No callback installed: items are dropped.
886 slot.invoke(1);
887
888 let mut seen = Vec::new();
889 slot.invoke_with(&mut |x| seen.push(x), || {
890 slot.invoke(2);
891 slot.invoke(3);
892 });
893 assert_eq!(seen, vec![2, 3]);
894
895 // Uninstalled again after `invoke_with` returns.
896 slot.invoke(4);
897 assert_eq!(seen, vec![2, 3]);
898 }
899
900 #[test]
901 fn callback_slot_nested_invoke_with_restores_outer() {
902 let slot: CallbackSlot<u32> = CallbackSlot::new();
903 let mut outer = Vec::new();
904 let mut inner = Vec::new();
905
906 slot.invoke_with(&mut |x| outer.push(x), || {
907 slot.invoke(1);
908 slot.invoke_with(&mut |x| inner.push(x), || slot.invoke(2));
909 // The outer callback must be restored after the nested installation.
910 slot.invoke(3);
911 });
912
913 assert_eq!(outer, vec![1, 3]);
914 assert_eq!(inner, vec![2]);
915 }
916
917 #[test]
918 fn callback_slot_reentrant_invoke_drops_item() {
919 let slot: Rc<CallbackSlot<u32>> = Rc::new(CallbackSlot::new());
920 let slot2 = Rc::clone(&slot);
921
922 let mut seen = Vec::new();
923 let mut callback = |x: u32| {
924 seen.push(x);
925 // Re-entrant invocation while the callback is running: the callback has been
926 // taken out of the slot, so this must be a no-op rather than aliasing UB.
927 if x == 1 {
928 slot2.invoke(99);
929 }
930 };
931 slot.invoke_with(&mut callback, || {
932 slot.invoke(1);
933 slot.invoke(2);
934 });
935 assert_eq!(seen, vec![1, 2]);
936 }
937
938 #[test]
939 fn callback_slot_uninstalls_on_panic() {
940 let slot: Rc<CallbackSlot<u32>> = Rc::new(CallbackSlot::new());
941 let slot2 = Rc::clone(&slot);
942
943 let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
944 slot.invoke_with(&mut |_| {}, || panic!("boom"));
945 }));
946 assert!(panicked.is_err());
947
948 // The (now-dead) callback must not still be installed after the unwind.
949 let mut seen = Vec::new();
950 slot2.invoke(5); // dropped, not UB
951 slot2.invoke_with(&mut |x| seen.push(x), || slot2.invoke(6));
952 assert_eq!(seen, vec![6]);
953 }
954
955 #[test]
956 fn owned_erased_box_drops_contents() {
957 struct NoteDrop(Rc<Cell<bool>>);
958 impl Drop for NoteDrop {
959 fn drop(&mut self) {
960 self.0.set(true);
961 }
962 }
963
964 let dropped = Rc::new(Cell::new(false));
965 // SAFETY: `reference` is not used after `erased` is dropped.
966 let (erased, reference) = unsafe { OwnedErasedBox::new(NoteDrop(Rc::clone(&dropped))) };
967 assert!(!reference.0.get());
968
969 // Move the erased box around; the allocation must remain valid.
970 let moved = [erased];
971 assert!(!dropped.get());
972 drop(moved);
973 assert!(dropped.get());
974 }
975
976 #[test]
977 fn aliased_box_derefs_and_drops() {
978 struct NoteDrop(Rc<Cell<bool>>);
979 impl Drop for NoteDrop {
980 fn drop(&mut self) {
981 self.0.set(true);
982 }
983 }
984
985 let dropped = Rc::new(Cell::new(false));
986 let boxed = AliasedBox::new(NoteDrop(Rc::clone(&dropped)));
987 assert!(!boxed.0.get());
988 drop(boxed);
989 assert!(dropped.get());
990 }
991}