Skip to main content

alloc_once_macros/
lib.rs

1//! Procedural macro backing `alloc_once!`.
2//!
3//! ```ignore
4//! alloc_once!(
5//!     n = 1i32;
6//!     bytes = [10u8, 20, 30];
7//!     s = "hello";
8//! );
9//! // n, bytes, s are now in scope as &mut into a single arena allocation.
10//! ```
11//!
12//! `alloc_once!` packs the named values into a single tuple, allocates that
13//! tuple once inside a freshly created bump arena, and expands to statements
14//! that bring each name into scope as a `&mut` into the arena. A hidden guard
15//! binding owns the arena and keeps it alive until the end of the scope; the
16//! values' destructors run and the arena is freed when it drops.
17//!
18//! All bumpalo details are hidden: the expansion only names paths under
19//! `::alloc_once`, so callers need no bumpalo dependency of their own.
20
21use proc_macro::TokenStream;
22use quote::quote;
23use syn::parse::{Parse, ParseStream};
24use syn::punctuated::Punctuated;
25use syn::{Expr, Ident, Token};
26
27/// The parsed macro input: a `;`-separated list of named fields.
28struct Fields {
29    fields: Punctuated<Field, Token![;]>,
30}
31
32/// A single `name = value` entry.
33struct Field {
34    name: Ident,
35    value: Expr,
36}
37
38impl Parse for Field {
39    fn parse(input: ParseStream) -> syn::Result<Self> {
40        let name = input.parse()?;
41        input.parse::<Token![=]>()?;
42        let value = input.parse()?;
43        Ok(Field { name, value })
44    }
45}
46
47impl Parse for Fields {
48    fn parse(input: ParseStream) -> syn::Result<Self> {
49        let fields = Punctuated::parse_terminated(input)?;
50        if fields.is_empty() {
51            return Err(syn::Error::new(
52                input.span(),
53                "alloc_once! needs at least one `name = value` field to allocate",
54            ));
55        }
56        Ok(Fields { fields })
57    }
58}
59
60/// Allocate several named (possibly different-typed) values in a single bump
61/// allocation and bring each name into scope as a `&mut` into the arena.
62///
63/// Fields are separated by `;`; a trailing `;` is optional. The macro expands to
64/// statements, so invoke it in statement position rather than assigning its
65/// result:
66///
67/// ```ignore
68/// alloc_once!(
69///     n = 1i32;
70///     bytes = [10u8, 20, 30];
71///     s = "hello";
72/// );
73/// *n += 41;      // n:     &mut i32
74/// bytes[0] = 99; // bytes: &mut [u8; 3]
75/// ```
76#[proc_macro]
77pub fn alloc_once(input: TokenStream) -> TokenStream {
78    let Fields { fields } = syn::parse_macro_input!(input as Fields);
79
80    let names = fields.iter().map(|f| &f.name);
81    let values = fields.iter().map(|f| &f.value);
82
83    // The guard binding outlives the generated block and so lands in the caller's
84    // scope. A `mixed_site` span gives it the macro's own hygiene context: the two
85    // references below resolve to each other, but caller code — even a variable
86    // literally named `__alloc_once_guard` — can neither see nor collide with it.
87    let guard = Ident::new("__alloc_once_guard", proc_macro2::Span::mixed_site());
88
89    let expanded = quote! {
90        // A guard under a hidden name owns the arena for the rest of the scope,
91        // keeping the backing memory alive while the `&mut` bindings below borrow
92        // into it.
93        let mut #guard = {
94            // Build the bundle first. The trailing comma keeps the single-field
95            // case a 1-tuple `(v,)`.
96            let __alloc_once_value = ( #(#values,)* );
97            // Size the arena to hold exactly this bundle and nothing more. For a
98            // sized tuple this is a compile-time constant, so the arena allocates
99            // a single chunk of precisely the bytes needed.
100            let __alloc_once_cap = ::core::mem::size_of_val(&__alloc_once_value);
101            let __alloc_once_arena =
102                ::alloc_once::__private::Bump::with_capacity(__alloc_once_cap);
103            let __alloc_once_ptr: *mut _ = __alloc_once_arena.alloc(__alloc_once_value);
104            unsafe { ::alloc_once::AllocOnce::__from_parts(__alloc_once_arena, __alloc_once_ptr) }
105        };
106        // Deref to the tuple and destructure into one `&mut` per field. Match
107        // ergonomics makes each binding a `&mut`, and the borrow checker allows
108        // it because the fields are disjoint.
109        let ( #(#names,)* ) = &mut *#guard;
110    };
111
112    expanded.into()
113}