alloc-once-macros 0.1.0-rc.3

//! Procedural macro backing `alloc_once!`.
//!
//! ```ignore
//! alloc_once!(
//!     n = 1i32;
//!     bytes = [10u8, 20, 30];
//!     s = "hello";
//! );
//! // n, bytes, s are now in scope as &mut into a single arena allocation.
//! ```
//!
//! `alloc_once!` packs the named values into a single tuple, allocates that
//! tuple once inside a freshly created bump arena, and expands to statements
//! that bring each name into scope as a `&mut` into the arena. A hidden guard
//! binding owns the arena and keeps it alive until the end of the scope; the
//! values' destructors run and the arena is freed when it drops.
//!
//! All bumpalo details are hidden: the expansion only names paths under
//! `::alloc_once`, so callers need no bumpalo dependency of their own.

use proc_macro::TokenStream;
use quote::quote;
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::{Expr, Ident, Token};

/// The parsed macro input: a `;`-separated list of named fields.
struct Fields {
    fields: Punctuated<Field, Token![;]>,
}

/// A single `name = value` entry.
struct Field {
    name: Ident,
    value: Expr,
}

impl Parse for Field {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let name = input.parse()?;
        input.parse::<Token![=]>()?;
        let value = input.parse()?;
        Ok(Field { name, value })
    }
}

impl Parse for Fields {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let fields = Punctuated::parse_terminated(input)?;
        if fields.is_empty() {
            return Err(syn::Error::new(
                input.span(),
                "alloc_once! needs at least one `name = value` field to allocate",
            ));
        }
        Ok(Fields { fields })
    }
}

/// Allocate several named (possibly different-typed) values in a single bump
/// allocation and bring each name into scope as a `&mut` into the arena.
///
/// Fields are separated by `;`; a trailing `;` is optional. The macro expands to
/// statements, so invoke it in statement position rather than assigning its
/// result:
///
/// ```ignore
/// alloc_once!(
///     n = 1i32;
///     bytes = [10u8, 20, 30];
///     s = "hello";
/// );
/// *n += 41;      // n:     &mut i32
/// bytes[0] = 99; // bytes: &mut [u8; 3]
/// ```
#[proc_macro]
pub fn alloc_once(input: TokenStream) -> TokenStream {
    let Fields { fields } = syn::parse_macro_input!(input as Fields);

    let names = fields.iter().map(|f| &f.name);
    let values = fields.iter().map(|f| &f.value);

    // The guard binding outlives the generated block and so lands in the caller's
    // scope. A `mixed_site` span gives it the macro's own hygiene context: the two
    // references below resolve to each other, but caller code — even a variable
    // literally named `__alloc_once_guard` — can neither see nor collide with it.
    let guard = Ident::new("__alloc_once_guard", proc_macro2::Span::mixed_site());

    let expanded = quote! {
        // A guard under a hidden name owns the arena for the rest of the scope,
        // keeping the backing memory alive while the `&mut` bindings below borrow
        // into it.
        let mut #guard = {
            // Build the bundle first. The trailing comma keeps the single-field
            // case a 1-tuple `(v,)`.
            let __alloc_once_value = ( #(#values,)* );
            // Size the arena to hold exactly this bundle and nothing more. For a
            // sized tuple this is a compile-time constant, so the arena allocates
            // a single chunk of precisely the bytes needed.
            let __alloc_once_cap = ::core::mem::size_of_val(&__alloc_once_value);
            let __alloc_once_arena =
                ::alloc_once::__private::Bump::with_capacity(__alloc_once_cap);
            let __alloc_once_ptr: *mut _ = __alloc_once_arena.alloc(__alloc_once_value);
            unsafe { ::alloc_once::AllocOnce::__from_parts(__alloc_once_arena, __alloc_once_ptr) }
        };
        // Deref to the tuple and destructure into one `&mut` per field. Match
        // ergonomics makes each binding a `&mut`, and the borrow checker allows
        // it because the fields are disjoint.
        let ( #(#names,)* ) = &mut *#guard;
    };

    expanded.into()
}