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}