Skip to content

Feature request: support context-threaded generation, so Arbitrary can work with arena allocators #240

Description

@freshtonic

Support context-threaded generation, so Arbitrary can work with arena allocators

Summary

Arbitrary::arbitrary has this shape today:

pub trait Arbitrary<'a>: Sized {
    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self>;
    // + arbitrary_take_rest, size_hint, ...
}

The function takes one input: the byte stream u. It has no channel for extra construction context. This works for plain heap types, because they call the global allocator implicitly. It does not work for a type that needs an external handle to construct itself - most importantly, a type that is backed by an arena allocator such as bumpalo::Bump.

We propose a new trait, ArbitraryWithContext<'a, C>, that adds one extra input: a caller-supplied context value C. This closes the gap and lets arbitrary-based generation reach arena-backed data structures, both by hand and through #[derive(Arbitrary)].

We are willing to prototype this and help drive an implementation if there is interest in the direction.

Problem statement

bumpalo provides arena-allocated equivalents of common owned types:

  • bumpalo::boxed::Box<'bump, T>, built with Box::new_in(value, bump)
  • bumpalo::collections::Vec<'bump, T>, built with Vec::new_in(bump)

Both constructors need a &'bump Bump handle at construction time. arbitrary has no way to supply one. The trait method has one parameter, u: &mut Unstructured<'a>, and no route for a second value to flow in. So there is no way to write:

impl<'a, 'bump, T: Arbitrary<'a>> Arbitrary<'a> for bumpalo::boxed::Box<'bump, T> {
    // Where would `&'bump Bump` come from?
}

This is a real barrier, not a theoretical one. We maintain recursa, a Rust parser generator. For large grammars, it generates AST node types that live in a bumpalo::Bump arena instead of the heap. A real PostgreSQL SQL grammar generates over 1,400 distinct node types in this mode, for example:

pub struct SelectStmt<'bump> {
    pub target_list: ArenaVec1<'bump, ResTarget<'bump>>,
    pub from_clause: Option<ArenaBox<'bump, FromClause<'bump>>>,
    // ... hundreds more fields, across hundreds of node types
}

(ArenaBox and ArenaVec1 wrap bumpalo::boxed::Box and bumpalo::collections::Vec respectively.)

We want to generate random, structurally valid SelectStmt<'bump> values - for a grammar explorer and for fuzzer-seed generation, which is exactly the job arbitrary does. But every field that holds an ArenaBox or ArenaVec needs a live &'bump Bump to build. Since Arbitrary::arbitrary cannot receive one, we cannot implement Arbitrary for these container types. Since almost every generated node type contains one of these containers somewhere in its field tree, we cannot derive Arbitrary for almost any node type in a 1,400-type grammar. The trait shape blocks the whole use case at the root.

Why existing options do not work

We looked at three ways to work around the missing channel. Each one has a real cost, and none of them is acceptable in a published, general-purpose library.

Thread-local or task-local scoped allocator. Store &'bump Bump in a thread_local!, and read it from a hand-written Arbitrary impl. This makes a public trait impl depend on invisible, ambiently-scoped setup. Any caller that invokes <T as Arbitrary>::arbitrary(u) directly - which is the normal way fuzzers and property-testers call the trait gets a panic or unsound behavior if the thread-local was not set up first. The approach also breaks down on async runtimes: a plain thread-local gives no per-task isolation when two tasks interleave on one worker thread.

Leak a fresh arena per generated sample. Box::leak a new Bump for each call, then generate into it. This works once, but it leaks memory on every single sample. A fuzzing loop or an interactive tool samples repeatedly and for a long time, so this is not viable for any long-running process.

Nightly #![feature(allocator_api)]. The real fix for the underlying mismatch is core::alloc::Allocator: implement it for &Bump, and use plain std::boxed::Box<T, &Bump> / std::vec::Vec<T, &Bump>. This is clean, but allocator_api has been unstable since 2016 (rust-lang/rust#32838), with no stabilization date. A general-purpose library cannot require its users to run nightly Rust indefinitely.

Edit: a stabilised allocator API will not actually address the issue - context will still be necessary. See #240 (comment)

Because all three workarounds fail for a normal, non-nightly, long-running process, the fix needs to live in the arbitrary trait itself.

Proposed API change

Add a new trait, ArbitraryWithContext<'a, C>, alongside Arbitrary<'a>.

/// Like `Arbitrary`, but threads caller-supplied context `C` through
/// construction.
pub trait ArbitraryWithContext<'a, C>: Sized {
    fn arbitrary_with(u: &mut Unstructured<'a>, cx: C) -> Result<Self>;

    fn arbitrary_take_rest_with(mut u: Unstructured<'a>, cx: C) -> Result<Self> {
        Self::arbitrary_with(&mut u, cx)
    }

    fn size_hint(depth: usize) -> (usize, Option<usize>) {
        (0, None)
    }
}

A blanket impl that covers every context, not only ()

The natural first attempt is a blanket impl for the unit context only:

impl<'a, T: Arbitrary<'a>> ArbitraryWithContext<'a, ()> for T {
    fn arbitrary_with(u: &mut Unstructured<'a>, _cx: ()) -> Result<Self> {
        T::arbitrary(u)
    }
}

This is not enough on its own. A derived struct with an arena field and a plain field - for example SelectStmt above, which mixes ArenaVec1 with plain Option<...> - needs every field, plain or arena-backed, to answer to the same ArbitraryWithContext<'a, C> call for whatever C the struct carries. u32 or String do not know about &'bump Bump.

So the blanket impl should cover every C, not only ():

impl<'a, C, T: Arbitrary<'a>> ArbitraryWithContext<'a, C> for T {
    fn arbitrary_with(u: &mut Unstructured<'a>, _cx: C) -> Result<Self> {
        T::arbitrary(u)
    }
}

With this in place, a derive macro can call <FieldType as ArbitraryWithContext<'a, C>>::arbitrary_with(u, cx) uniformly on every field, with no branching on whether that field type happens to need cx. Plain fields ignore cx through the blanket impl; arena fields consume it through a hand-written impl (below).

This blanket impl has a real cost, covered under "open questions."

Threading C through nested and derived fields

#[derive(Arbitrary)] needs a way to say "this type requires context C," and to propagate that same C to every field. We propose an opt-in struct or enum attribute:

#[derive(Arbitrary)]
#[arbitrary(context = "&'bump bumpalo::Bump")]
pub struct SelectStmt<'bump> {
    pub target_list: ArenaVec1<'bump, ResTarget<'bump>>,
    pub from_clause: Option<ArenaBox<'bump, FromClause<'bump>>>,
}

With no #[arbitrary(context = ...)] attribute, the derive macro behaves exactly as it does today, and emits a plain Arbitrary impl. This keeps the change fully backward compatible and opt-in.

With the attribute present, the derive macro emits an ArbitraryWithContext impl instead, and calls arbitrary_with(u, cx) on every field, passing the same cx value through unchanged:

impl<'a, 'bump> ArbitraryWithContext<'a, &'bump bumpalo::Bump> for SelectStmt<'bump> {
    fn arbitrary_with(
        u: &mut Unstructured<'a>,
        cx: &'bump bumpalo::Bump,
    ) -> Result<Self> {
        Ok(SelectStmt {
            target_list: ArbitraryWithContext::arbitrary_with(u, cx)?,
            from_clause: ArbitraryWithContext::arbitrary_with(u, cx)?,
        })
    }
}

Because cx is used once per field (and, for collections, once per element), C needs a Copy bound. A reference context such as &'bump Bump satisfies Copy for free, which is why the design threads cx by value rather than by reference. This bound is discussed further below.

C as a type parameter, not an associated type

We propose C as a generic type parameter on the trait (ArbitraryWithContext<'a, C>), not as an associated type (ArbitraryWithContext<'a> { type Context; ... }).

An associated type would give each implementing type exactly one context type. That is simpler in some ways, but it removes the ability to give one type more than one context-shaped generation strategy - one crate might want Vec<T>-like generation keyed on an allocator, another on a length distribution config, and so on. A generic type parameter keeps that door open, the same way From<T> stays generic instead of using an associated type, so that one type can convert from several source types.

The cost is that a struct whose fields need two unrelated context types (say, an allocator on one field and an unrelated RNG seed on another) cannot be handled by a single C without bundling both into one context struct or tuple, and teaching the derive macro how to project each field's slice out of that bundle. We consider that a real but separate problem, out of scope for the first version of this proposal. The arena use case only ever needs one context value threaded unchanged to every field, which the design above covers directly.

Closing the loop: the arena-allocator side

Note: we're using bumpalo as an example only - the design should work with any arena allocation crate.

Downstream crates, such as bumpalo itself could add their own support:

impl<'a, 'bump, T> ArbitraryWithContext<'a, &'bump Bump> for bumpalo::boxed::Box<'bump, T>
where
    T: ArbitraryWithContext<'a, &'bump Bump>,
{
    fn arbitrary_with(u: &mut Unstructured<'a>, cx: &'bump Bump) -> Result<Self> {
        let value = T::arbitrary_with(u, cx)?;
        Ok(bumpalo::boxed::Box::new_in(value, cx))
    }

    fn size_hint(depth: usize) -> (usize, Option<usize>) {
        T::size_hint(depth)
    }
}

impl<'a, 'bump, T> ArbitraryWithContext<'a, &'bump Bump> for bumpalo::collections::Vec<'bump, T>
where
    T: ArbitraryWithContext<'a, &'bump Bump>,
{
    fn arbitrary_with(u: &mut Unstructured<'a>, cx: &'bump Bump) -> Result<Self> {
        let len = u.arbitrary_len::<T>()?;
        let mut v = bumpalo::collections::Vec::with_capacity_in(len, cx);
        for _ in 0..len {
            v.push(T::arbitrary_with(u, cx)?);
        }
        Ok(v)
    }
}

With these two impls, plus the derive attribute above, a call site like:

let bump = bumpalo::Bump::new();
let mut u = Unstructured::new(&raw_bytes);
let stmt = SelectStmt::arbitrary_with(&mut u, &bump)?;

produces a fully arena-backed SelectStmt<'_>, recursively, through ordinary arbitrary_with calls, with no thread-local, no leak, and no nightly feature. This is the shape that actually removes the blocker described above.

Rust's orphan rule blocks the direct route: ArbitraryWithContext is foreign (defined in arbitrary) and bumpalo::boxed::Box is foreign (defined in bumpalo), so a downstream crate cannot impl ArbitraryWithContext<'a, &'bump Bump> for bumpalo::boxed::Box<'bump, T> itself - neither the trait nor the type is local to that crate. But the orphan rule only blocks foreign-for- foreign impls. A local newtype wrapping the foreign arena type is a local type, and the impl becomes legal immediately, with no upstream change at all:

/// A local, arena-backed box. Not a type alias - a real newtype, so the
/// orphan rule lets us implement foreign traits for it.
pub struct ArenaBox<'bump, T>(pub bumpalo::boxed::Box<'bump, T>);

impl<'bump, T> std::ops::Deref for ArenaBox<'bump, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.0
    }
}

impl<'a, 'bump, T> ArbitraryWithContext<'a, &'bump bumpalo::Bump> for ArenaBox<'bump, T>
where
    T: ArbitraryWithContext<'a, &'bump bumpalo::Bump>,
{
    fn arbitrary_with(u: &mut Unstructured<'a>, cx: &'bump bumpalo::Bump) -> Result<Self> {
        let value = T::arbitrary_with(u, cx)?;
        Ok(ArenaBox(bumpalo::boxed::Box::new_in(value, cx)))
    }

    fn size_hint(depth: usize) -> (usize, Option<usize>) {
        T::size_hint(depth)
    }
}

/// Same pattern for the collection side.
pub struct ArenaVec<'bump, T>(pub bumpalo::collections::Vec<'bump, T>);

impl<'a, 'bump, T> ArbitraryWithContext<'a, &'bump bumpalo::Bump> for ArenaVec<'bump, T>
where
    T: ArbitraryWithContext<'a, &'bump bumpalo::Bump>,
{
    fn arbitrary_with(u: &mut Unstructured<'a>, cx: &'bump bumpalo::Bump) -> Result<Self> {
        let len = u.arbitrary_len::<T>()?;
        let mut v = bumpalo::collections::Vec::with_capacity_in(len, cx);
        for _ in 0..len {
            v.push(T::arbitrary_with(u, cx)?);
        }
        Ok(ArenaVec(v))
    }
}

This is a completely general technique - it works the same way for any foreign arena or allocator crate, not just bumpalo, and needs nothing from that crate's maintainers beyond the types it already exposes publicly. The cost is the wrapper itself: every method the inner type offers has to be re-exposed (Deref covers read access cheaply, as above; anything that needs to return the wrapped type, such as a push that internally reallocates, needs a small forwarding method). For a generated-code consumer - a derive macro emitting these wrapper types and their forwarding methods mechanically - that cost is close to free, since nothing is hand-written.

It is also worth being direct about why this section exists at all: it is the fallback we would actually build today if this proposal stalls, and it is already the shape our own ArenaBox/ArenaVec type aliases would need to become to make any of this work - a plain pub use bumpalo::boxed::Box as ArenaBox re-export is exactly as foreign as bumpalo::boxed::Box itself, so even our own crate cannot implement a foreign trait against it without first turning it into a newtype like the one above. We are noting this not to suggest the newtype pattern makes the trait-level proposal unnecessary - it still requires ArbitraryWithContext to exist in arbitrary in the first place, and it still pushes real, mechanical boilerplate onto every downstream crate that adopts it - but to show that adopting this proposal does not require the rest of the ecosystem to move first.

Open questions and tradeoffs

We do not present the proposed design as a free lunch.

  • The blanket impl forecloses per-C specialization for existing Arbitrary types. Once impl<'a, C, T: Arbitrary<'a>> ArbitraryWithContext<'a, C> for T exists, no other crate can write a specific ArbitraryWithContext<'a, SomeC> impl for a type that already implements plain Arbitrary - Rust's coherence rules forbid the overlap. In practice this is not a problem for the motivating case, because bumpalo's arena types are new types that do not implement plain Arbitrary at all. It could matter later if a crate wants to give an existing Arbitrary type a second, context-aware generation strategy for a specific C.

  • The ecosystem splits into two families of impl. A type ends up either "always ignores context," through the blanket impl, or "context-aware for one or more specific C," through a hand-written impl - never both for the same C. This mirrors the bumpalo design already (separate arena-aware collection types, rather than one Vec that works both ways), so it should not be a surprise to users of arena crates, but it is worth naming explicitly.

  • The C: Copy bound is a real restriction. It is free for reference contexts (&'bump Bump), which cover the arena case cleanly. It is a genuine limit for a context that is expensive or impossible to Copy (a mutable RNG, a non-Copy interner handle). We think this is an acceptable first-pass restriction, since the motivating use case never needs it, but a full RFC-level design should say explicitly whether Copy is required, or whether the trait should instead standardize on threading cx as &C in places where the extra indirection does not break lifetime relationships the way it would for Box<'bump, T>.

  • Derive macro complexity. The per-field call pattern is not much more complex than today's derive, because the general blanket impl means every field answers to the same arbitrary_with(u, cx) call, with no per-field branching on whether that field needs cx. The added complexity is mostly in attribute parsing (recognizing #[arbitrary(context = "...")] and wiring the named type into the generated impl header) and in enum support (each variant's fields need the same treatment as a struct's).

  • Backward compatibility. This is additive only. Arbitrary is untouched. Existing manual impls and existing derives keep working exactly as they do today, with no attribute added. The general blanket impl is what gives every existing Arbitrary type automatic, free access to ArbitraryWithContext for any C, so downstream code that starts using context-threaded generation does not need to touch types it does not own.

  • Naming and attribute syntax. ArbitraryWithContext, arbitrary_with, and #[arbitrary(context = "...")] are working names for this proposal, not a final answer. We expect the maintainers have views here and are happy to defer to them.

Offer to help

We hit this gap directly while building arena-backed AST generation for a 1,400-node-type grammar, so we have a real test case and a real motivation to see this land. We are willing to prototype this design - the trait, the blanket impl, the derive-macro attribute, and a small bumpalo bridging impl - and to iterate on the API shape with maintainer feedback, if there is interest in taking this direction.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions