nuvec_macros/
lib.rs

1//! Derive macros for nuvec, re-exported as `nuvec::Tree`, `nuvec::Static` and
2//! `nuvec_nn::Module`: use them through those crates.
3//!
4//! One expansion engine generates `Leaves` and `Tree` impls; `derive(Module)` is
5//! `derive(Tree)` plus a lint against bare `Array` fields. The generated code names only
6//! `::nuvec::__derive` (or the module given by `#[tree(crate = "..")]` on the type).
7//!
8//! Field attributes (`#[tree(..)]`; `#[module(..)]` is an alias under `derive(Module)`):
9//!
10//! | attribute | meaning |
11//! |---|---|
12//! | (none) | a child tree (its leaves are visited) |
13//! | `skip` | not a leaf; hashed into the structure (the type must be `Static`); on an `Array` field (last path segment `Array`): a computed constant, neither hashed nor saved |
14//! | `skip, static_unchecked` | not a leaf, not hashed: baked into compiled traces |
15//! | `buffer` | an `Array` leaf that is saved and loaded but never trained |
16//! | `mode` | a `bool` mode flag (`true` = training), part of the structure |
17//! | `rename = ".."` | the path segment (fields and enum variants) |
18//!
19//! **Stability:** no promise when this crate is depended on directly; depend on `nuvec`,
20//! whose re-exports are semver-tracked.
21
22#![forbid(unsafe_code)]
23
24use proc_macro::TokenStream;
25use syn::{DeriveInput, parse_macro_input};
26
27mod attrs;
28mod expand;
29mod lint;
30
31use expand::Kind;
32
33/// Derives `Leaves` and `Tree` for a struct (named, tuple or unit) or an enum: the type's
34/// array leaves get paths from its field names, as MLX's `tree_flatten`.
35///
36/// ```ignore
37/// #[derive(Clone, Tree)]
38/// struct Batch { x: Array, y: Array, #[tree(skip)] len: usize }
39///
40/// #[derive(Clone, Tree)]
41/// enum Either { A(Array), #[tree(rename = "b")] B(Vec<Array>), Empty }
42/// ```
43///
44/// A newtype enum variant is transparent unless renamed; struct variants name their
45/// fields; every variant reports its index, so a leaf-less variant changes the structure.
46/// Type parameters are bounded only where a field uses them. Trait-object fields are
47/// rejected (use an enum of the concrete types).
48#[proc_macro_derive(Tree, attributes(tree))]
49pub fn derive_tree(input: TokenStream) -> TokenStream {
50    let input = parse_macro_input!(input as DeriveInput);
51    expand::derive_tree(&input, Kind::Tree)
52        .unwrap_or_else(syn::Error::into_compile_error)
53        .into()
54}
55
56/// Derives `Tree` for a model (it emits no `impl Module`: `Module` is implemented for every
57/// tree), with one more check: a field typed `Array` (directly or in `Option`, `Vec` or
58/// `Box`) without `buffer` or `skip` is an error suggesting `Param`. The check is
59/// syntactic: a type alias of `Array` is not caught.
60#[proc_macro_derive(Module, attributes(tree, module))]
61pub fn derive_module(input: TokenStream) -> TokenStream {
62    let input = parse_macro_input!(input as DeriveInput);
63    expand::derive_tree(&input, Kind::Module)
64        .unwrap_or_else(syn::Error::into_compile_error)
65        .into()
66}
67
68/// Derives `Static` (the type of a hashed `#[tree(skip)]` field) by hashing every field in
69/// order (enums: the variant index, then its fields).
70#[proc_macro_derive(Static, attributes(tree))]
71pub fn derive_static(input: TokenStream) -> TokenStream {
72    let input = parse_macro_input!(input as DeriveInput);
73    expand::derive_static(&input)
74        .unwrap_or_else(syn::Error::into_compile_error)
75        .into()
76}
nuvec_macros/lib.rs

nuvec_macros/
lib.rs
