vortex_array_macros/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Proc macros for `vortex-array`.
5
6use proc_macro::TokenStream;
7use proc_macro2::Span;
8use quote::format_ident;
9use quote::quote;
10use syn::Field;
11use syn::Fields;
12use syn::Ident;
13use syn::ItemStruct;
14use syn::LitInt;
15use syn::Path;
16use syn::Token;
17use syn::Type;
18use syn::Visibility;
19use syn::parse::Parse;
20use syn::parse::ParseStream;
21use syn::parse_macro_input;
22use syn::spanned::Spanned;
23
24/// Name of the per-field attribute that pins a field to a slot index.
25const SLOT_ATTR: &str = "slot";
26
27/// Generate slot index constants, a borrowed view struct, and a typed ext trait
28/// from a slot struct definition.
29///
30/// Fields must be `ArrayRef` (required slot), `Option<ArrayRef>` (optional slot), or
31/// `Vec<ArrayRef>` (variadic tail of required slots).
32///
33/// Every field must carry a `#[slot(..)]` attribute naming the exact slot index it maps
34/// to. The attribute — not the declaration order — defines the storage layout, so fields
35/// may be reordered, grouped, or documented in any order without changing the slot
36/// indices an array is built from or read back with.
37///
38/// # Example
39///
40/// ```ignore
41/// #[array_slots(Patched)]
42/// pub struct PatchedSlots {
43///     #[slot(0)]
44///     pub inner: ArrayRef,
45///     #[slot(1)]
46///     pub lane_offsets: ArrayRef,
47///     #[slot(2)]
48///     pub patch_indices: ArrayRef,
49///     #[slot(3)]
50///     pub patch_values: ArrayRef,
51/// }
52/// ```
53///
54/// # Generated output
55///
56/// Given the above, the macro generates:
57///
58/// ```ignore
59/// // --- The original struct, minus the consumed `#[slot(..)]` attributes ---
60/// pub struct PatchedSlots { ... }
61///
62/// // --- Slot index constants and conversion methods on the struct ---
63/// impl PatchedSlots {
64///     pub const INNER: usize = 0;
65///     pub const LANE_OFFSETS: usize = 1;
66///     pub const PATCH_INDICES: usize = 2;
67///     pub const PATCH_VALUES: usize = 3;
68///     pub const COUNT: usize = 4;
69///     pub const NAMES: [&'static str; 4] = ["inner", "lane_offsets", "patch_indices", "patch_values"];
70///
71///     /// Take ownership of slots from an `ArraySlots`.
72///     pub fn from_slots(slots: ArraySlots) -> Self { ... }
73///
74///     /// Convert back into storage order.
75///     pub fn into_slots(self) -> ArraySlots { ... }
76/// }
77///
78/// // --- Borrowed view with &ArrayRef / Option<&ArrayRef> fields ---
79/// pub struct PatchedSlotsView<'a> {
80///     pub inner: &'a ArrayRef,
81///     pub lane_offsets: &'a ArrayRef,
82///     pub patch_indices: &'a ArrayRef,
83///     pub patch_values: &'a ArrayRef,
84/// }
85///
86/// impl<'a> PatchedSlotsView<'a> {
87///     pub fn from_slots(slots: &'a [Option<ArrayRef>]) -> Self { ... }
88///     pub fn to_owned(&self) -> PatchedSlots { ... }
89/// }
90///
91/// // --- Ext trait with per-field accessors + slots_view() ---
92/// pub trait PatchedArraySlotsExt: TypedArrayRef<Patched> {
93///     fn inner(&self) -> &ArrayRef { ... }         // indexes slots directly
94///     fn lane_offsets(&self) -> &ArrayRef { ... }
95///     fn patch_indices(&self) -> &ArrayRef { ... }
96///     fn patch_values(&self) -> &ArrayRef { ... }
97///     fn slots_view(&self) -> PatchedSlotsView<'_> { ... }
98/// }
99///
100/// impl<T: TypedArrayRef<Patched>> PatchedArraySlotsExt for T {}
101/// ```
102///
103/// # Slot index annotations
104///
105/// - Fixed fields use `#[slot(N)]`, where `N` is the exact index of the slot.
106/// - A variadic tail uses `#[slot(N..)]`, where `N` is the index its first slot occupies.
107///
108/// The annotations are validated at compile time: the fixed indices must cover
109/// `0..FIXED_COUNT` exactly — no duplicates and no gaps — and a variadic tail must start
110/// immediately after the last fixed slot. A field without a `#[slot(..)]` attribute is a
111/// compile error, so the layout can never silently fall back to declaration order.
112///
113/// # Required vs optional slots
114///
115/// - `ArrayRef` — the slot must be present. `from_slots()` panics if `None`.
116///   The ext trait accessor returns `&ArrayRef`. The view field is `&'a ArrayRef`.
117///
118/// - `Option<ArrayRef>` — the slot may be absent. `from_slots()` preserves `None`.
119///   The ext trait accessor returns `Option<&ArrayRef>`. The view field is
120///   `Option<&'a ArrayRef>`.
121///
122/// The underlying storage is always `ArraySlots` — the field type only
123/// controls whether the macro inserts a `.vortex_expect()` unwrap or not.
124///
125/// # Variadic tail slots
126///
127/// One field may be `Vec<ArrayRef>`, declaring that every slot from its index onward
128/// belongs to a homogeneous, variable-length run of required slots. This supports
129/// encodings like `Chunked` (`[chunk_offsets, chunks...]`), `Struct`
130/// (`[validity?, fields...]`), and `Union` (`[type_ids, children...]`).
131///
132/// ```ignore
133/// #[array_slots(Chunked)]
134/// pub struct ChunkedSlots {
135///     #[slot(0)]
136///     pub chunk_offsets: ArrayRef,
137///     #[slot(1..)]
138///     pub chunks: Vec<ArrayRef>,
139/// }
140/// ```
141///
142/// For a struct with a variadic tail, the macro generates a different set of
143/// constants — slot count is no longer a compile-time constant:
144///
145/// ```ignore
146/// impl ChunkedSlots {
147///     pub const CHUNK_OFFSETS: usize = 0;
148///     /// Offset at which the `chunks` slots begin.
149///     pub const CHUNKS_OFFSET: usize = 1;
150///     /// Number of fixed (non-variadic) slots.
151///     pub const FIXED_COUNT: usize = 1;
152///     /// Names of the fixed slots in storage order.
153///     pub const FIXED_NAMES: [&'static str; 1] = ["chunk_offsets"];
154///
155///     /// Name of the slot at `idx`, e.g. "chunk_offsets" or "chunks[3]".
156///     pub fn slot_name(idx: usize) -> String { ... }
157///
158///     pub fn from_slots(slots: ArraySlots) -> Self { ... }
159///     pub fn into_slots(self) -> ArraySlots { ... }
160/// }
161/// ```
162///
163/// The view field and ext trait accessor for the tail are a `vortex_array::SlotSlice`,
164/// a borrowed run of required slots supporting `len()`, `get()`, `iter()`, and
165/// indexing:
166///
167/// ```ignore
168/// pub struct ChunkedSlotsView<'a> {
169///     pub chunk_offsets: &'a ArrayRef,
170///     pub chunks: SlotSlice<'a>,
171/// }
172///
173/// pub trait ChunkedArraySlotsExt: TypedArrayRef<Chunked> {
174///     fn chunk_offsets(&self) -> &ArrayRef { ... }
175///     fn chunks(&self) -> SlotSlice<'_> { ... }
176///     fn slots_view(&self) -> ChunkedSlotsView<'_> { ... }
177/// }
178/// ```
179#[proc_macro_attribute]
180pub fn array_slots(attr: TokenStream, item: TokenStream) -> TokenStream {
181    let encoding = parse_macro_input!(attr as Path);
182    let item_struct = parse_macro_input!(item as ItemStruct);
183
184    match expand_array_slots(encoding, item_struct) {
185        Ok(tokens) => tokens.into(),
186        Err(err) => err.to_compile_error().into(),
187    }
188}
189
190fn expand_array_slots(
191    encoding: Path,
192    item_struct: ItemStruct,
193) -> syn::Result<proc_macro2::TokenStream> {
194    if !item_struct.generics.params.is_empty() || item_struct.generics.where_clause.is_some() {
195        return Err(syn::Error::new(
196            item_struct.generics.span(),
197            "#[array_slots] does not support generic slot structs",
198        ));
199    }
200
201    let fields = match &item_struct.fields {
202        Fields::Named(fields) => &fields.named,
203        _ => {
204            return Err(syn::Error::new(
205                item_struct.span(),
206                "#[array_slots] requires a struct with named fields",
207            ));
208        }
209    };
210
211    let encoding_ident = encoding
212        .segments
213        .last()
214        .map(|segment| &segment.ident)
215        .ok_or_else(|| syn::Error::new(encoding.span(), "missing encoding type"))?;
216
217    let struct_ident = item_struct.ident.clone();
218    let struct_vis = item_struct.vis.clone();
219    let view_ident = format_ident!("{}View", ident_name(&struct_ident));
220    let ext_ident = format_ident!("{}ArraySlotsExt", ident_name(encoding_ident));
221
222    let field_specs = fields
223        .iter()
224        .map(|field| SlotField::new(field, &struct_ident))
225        .collect::<syn::Result<Vec<_>>>()?;
226
227    // The `#[slot(..)]` annotations, not the declaration order, define storage order.
228    let (fixed_specs, tail_spec) = partition_by_slot_index(&field_specs)?;
229
230    let idx_consts = fixed_specs.iter().copied().map(SlotField::idx_const);
231    let view_fields = field_specs.iter().map(SlotField::view_field);
232    let view_from_slots = field_specs.iter().map(SlotField::view_from_slots);
233    let view_to_owned = field_specs.iter().map(SlotField::view_to_owned);
234    let ext_methods = field_specs.iter().map(SlotField::ext_method);
235
236    let counts = gen_counts(&fixed_specs, tail_spec);
237    let from_slots = gen_from_slots(&fixed_specs, tail_spec);
238    let into_slots = gen_into_slots(&fixed_specs, tail_spec);
239
240    // `#[slot(..)]` is inert helper syntax consumed here; strip it from the emitted struct.
241    let item_struct = strip_slot_attrs(item_struct);
242
243    Ok(quote! {
244        #item_struct
245
246        impl #struct_ident {
247            #(#idx_consts)*
248
249            #counts
250
251            #[doc = "Convert owned slot storage into an owned slot struct."]
252            #from_slots
253
254            #[doc = "Convert this slot struct into storage order."]
255            #into_slots
256        }
257
258        #[derive(Clone, Copy, Debug)]
259        #[doc = concat!("Borrowed view of `", stringify!(#struct_ident), "`.")]
260        #struct_vis struct #view_ident<'a> {
261            #(#view_fields,)*
262        }
263
264        impl<'a> #view_ident<'a> {
265            #[doc = "Borrow a slot slice as a typed view."]
266            pub fn from_slots(slots: &'a [Option<::vortex_array::ArrayRef>]) -> Self {
267                Self {
268                    #(#view_from_slots,)*
269                }
270            }
271
272            #[doc = "Clone all referenced slots into an owned slot struct."]
273            pub fn to_owned(&self) -> #struct_ident {
274                #struct_ident {
275                    #(#view_to_owned,)*
276                }
277            }
278        }
279
280        #[doc = concat!("Typed array accessors for `", stringify!(#encoding_ident), "`.")]
281        #struct_vis trait #ext_ident: ::vortex_array::TypedArrayRef<#encoding> {
282            #(#ext_methods)*
283
284            #[doc = "Returns a borrowed view of all slots."]
285            fn slots_view(&self) -> #view_ident<'_> {
286                #view_ident::from_slots(self.as_ref().slots())
287            }
288        }
289
290        impl<T: ::vortex_array::TypedArrayRef<#encoding>> #ext_ident for T {}
291    })
292}
293
294/// Split the fields into the fixed slots (sorted by slot index) and the optional variadic
295/// tail, validating that the annotated indices describe a complete, gap-free layout.
296fn partition_by_slot_index(
297    field_specs: &[SlotField],
298) -> syn::Result<(Vec<&SlotField>, Option<&SlotField>)> {
299    let mut fixed_specs = Vec::with_capacity(field_specs.len());
300    let mut tail_spec: Option<&SlotField> = None;
301
302    for spec in field_specs {
303        if matches!(spec.slot_type, SlotFieldType::VariadicTail) {
304            if let Some(previous) = tail_spec {
305                return Err(syn::Error::new(
306                    spec.index_span,
307                    format!(
308                        "#[array_slots] allows at most one variadic tail, but `{}` is already \
309                         declared as one",
310                        previous.slot_name
311                    ),
312                ));
313            }
314            tail_spec = Some(spec);
315        } else {
316            fixed_specs.push(spec);
317        }
318    }
319
320    fixed_specs.sort_by_key(|spec| spec.index);
321
322    for (expected, spec) in fixed_specs.iter().enumerate() {
323        if spec.index == expected {
324            continue;
325        }
326        // `fixed_specs` is sorted, so a smaller index than expected means the previous
327        // field claimed the same slot, and a larger one means nothing claimed `expected`.
328        return Err(if spec.index < expected {
329            syn::Error::new(
330                spec.index_span,
331                format!(
332                    "#[array_slots] slot index {} is claimed by both `{}` and `{}`",
333                    spec.index,
334                    fixed_specs[expected - 1].slot_name,
335                    spec.slot_name
336                ),
337            )
338        } else {
339            syn::Error::new(
340                spec.index_span,
341                format!(
342                    "#[array_slots] no field claims slot index {expected}; fixed slot indices \
343                     must cover 0..{} without gaps",
344                    fixed_specs.len()
345                ),
346            )
347        });
348    }
349
350    if let Some(tail) = tail_spec
351        && tail.index != fixed_specs.len()
352    {
353        return Err(syn::Error::new(
354            tail.index_span,
355            format!(
356                "#[array_slots] variadic tail `{}` must start at slot index {}, immediately after \
357                 the {} fixed slot(s), but is annotated `#[slot({}..)]`",
358                tail.slot_name,
359                fixed_specs.len(),
360                fixed_specs.len(),
361                tail.index
362            ),
363        ));
364    }
365
366    Ok((fixed_specs, tail_spec))
367}
368
369/// Remove the inert `#[slot(..)]` helper attributes so the struct can be re-emitted.
370fn strip_slot_attrs(mut item_struct: ItemStruct) -> ItemStruct {
371    if let Fields::Named(fields) = &mut item_struct.fields {
372        for field in &mut fields.named {
373            field.attrs.retain(|attr| !attr.path().is_ident(SLOT_ATTR));
374        }
375    }
376    item_struct
377}
378
379fn gen_counts(
380    fixed_specs: &[&SlotField],
381    tail_spec: Option<&SlotField>,
382) -> proc_macro2::TokenStream {
383    let names = fixed_specs.iter().map(|field| field.slot_name.as_str());
384    let fixed_count = fixed_specs.len();
385
386    match tail_spec {
387        None => quote! {
388            #[doc = "Total number of slots."]
389            pub const COUNT: usize = #fixed_count;
390
391            #[doc = "Slot names in storage order."]
392            pub const NAMES: [&'static str; #fixed_count] = [#(#names),*];
393        },
394        Some(tail) => {
395            let offset_const = &tail.const_ident;
396            let tail_name = &tail.slot_name;
397            quote! {
398                #[doc = concat!("Offset at which the `", #tail_name, "` slots begin.")]
399                pub const #offset_const: usize = #fixed_count;
400
401                #[doc = "Number of fixed (non-variadic) slots."]
402                pub const FIXED_COUNT: usize = #fixed_count;
403
404                #[doc = "Names of the fixed slots in storage order."]
405                pub const FIXED_NAMES: [&'static str; #fixed_count] = [#(#names),*];
406
407                #[doc = "Name of the slot at the given index."]
408                pub fn slot_name(idx: usize) -> String {
409                    if idx < Self::FIXED_COUNT {
410                        Self::FIXED_NAMES[idx].to_string()
411                    } else {
412                        format!(concat!(#tail_name, "[{}]"), idx - Self::#offset_const)
413                    }
414                }
415            }
416        }
417    }
418}
419
420fn gen_from_slots(
421    fixed_specs: &[&SlotField],
422    tail_spec: Option<&SlotField>,
423) -> proc_macro2::TokenStream {
424    let owned_from_slots = fixed_specs.iter().copied().map(SlotField::owned_from_slots);
425
426    match tail_spec {
427        None => quote! {
428            pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self {
429                Self {
430                    #(#owned_from_slots,)*
431                }
432            }
433        },
434        Some(tail) => {
435            let tail_ident = &tail.field_ident;
436            let offset_const = &tail.const_ident;
437            let expect_message = &tail.expect_message;
438            quote! {
439                pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self {
440                    let __variadic_tail: ::std::vec::Vec<::vortex_array::ArrayRef> = slots
441                        .drain(Self::#offset_const..)
442                        .map(|slot| ::vortex_error::VortexExpect::vortex_expect(
443                            slot,
444                            #expect_message,
445                        ))
446                        .collect();
447                    Self {
448                        #(#owned_from_slots,)*
449                        #tail_ident: __variadic_tail,
450                    }
451                }
452            }
453        }
454    }
455}
456
457fn gen_into_slots(
458    fixed_specs: &[&SlotField],
459    tail_spec: Option<&SlotField>,
460) -> proc_macro2::TokenStream {
461    let fixed_into_slots = fixed_specs.iter().copied().map(SlotField::storage_slot);
462
463    match tail_spec {
464        None => quote! {
465            pub fn into_slots(self) -> ::vortex_array::ArraySlots {
466                ::vortex_array::smallvec::smallvec![#(#fixed_into_slots),*]
467            }
468        },
469        Some(tail) => {
470            let tail_ident = &tail.field_ident;
471            quote! {
472                pub fn into_slots(self) -> ::vortex_array::ArraySlots {
473                    let mut slots: ::vortex_array::ArraySlots =
474                        ::vortex_array::smallvec::smallvec![#(#fixed_into_slots),*];
475                    slots.extend(self.#tail_ident.into_iter().map(Some));
476                    slots
477                }
478            }
479        }
480    }
481}
482
483struct SlotField {
484    field_ident: Ident,
485    field_vis: Visibility,
486    const_ident: Ident,
487    slot_name: String,
488    slot_type: SlotFieldType,
489    index: usize,
490    index_span: Span,
491    expect_message: syn::LitStr,
492    struct_ident: Ident,
493}
494
495impl SlotField {
496    fn new(field: &Field, struct_ident: &Ident) -> syn::Result<Self> {
497        let field_ident = field
498            .ident
499            .clone()
500            .ok_or_else(|| syn::Error::new(field.span(), "slot fields must be named"))?;
501        let field_name = ident_name(&field_ident);
502        let slot_type = SlotFieldType::from_syn_type(&field.ty)?;
503        let annotation = SlotIndexAttr::from_field(field, &field_name)?;
504
505        match (slot_type, annotation.variadic) {
506            (SlotFieldType::VariadicTail, false) => {
507                return Err(syn::Error::new(
508                    annotation.span,
509                    format!(
510                        "`{field_name}` is a variadic `Vec<ArrayRef>` tail, so it must be \
511                         annotated `#[slot({}..)]`",
512                        annotation.index
513                    ),
514                ));
515            }
516            (SlotFieldType::Required | SlotFieldType::Optional, true) => {
517                return Err(syn::Error::new(
518                    annotation.span,
519                    format!(
520                        "`#[slot(N..)]` declares a variadic `Vec<ArrayRef>` tail; `{field_name}` \
521                         occupies a single slot, so annotate it `#[slot({})]`",
522                        annotation.index
523                    ),
524                ));
525            }
526            _ => {}
527        }
528
529        let const_ident = match slot_type {
530            SlotFieldType::VariadicTail => {
531                format_ident!("{}_OFFSET", to_screaming_snake_case(&field_name))
532            }
533            _ => format_ident!("{}", to_screaming_snake_case(&field_name)),
534        };
535        let expect_message = syn::LitStr::new(
536            &format!("{} {} slot", ident_name(struct_ident), field_name),
537            field.span(),
538        );
539
540        Ok(Self {
541            field_ident,
542            field_vis: field.vis.clone(),
543            const_ident,
544            slot_name: field_name,
545            slot_type,
546            index: annotation.index,
547            index_span: annotation.span,
548            expect_message,
549            struct_ident: struct_ident.clone(),
550        })
551    }
552
553    fn idx_const(&self) -> proc_macro2::TokenStream {
554        let const_ident = &self.const_ident;
555        let index = self.index;
556        let slot_name = &self.slot_name;
557
558        quote! {
559            #[doc = concat!("Slot index for `", #slot_name, "`.")]
560            pub const #const_ident: usize = #index;
561        }
562    }
563
564    fn view_field(&self) -> proc_macro2::TokenStream {
565        let field_ident = &self.field_ident;
566        let field_vis = &self.field_vis;
567        let ty = self.slot_type.view_field_ty();
568
569        quote! {
570            #field_vis #field_ident: #ty
571        }
572    }
573
574    fn view_from_slots(&self) -> proc_macro2::TokenStream {
575        let field_ident = &self.field_ident;
576        let struct_ident = &self.struct_ident;
577        let const_ident = &self.const_ident;
578        let expect_message = &self.expect_message;
579
580        match self.slot_type {
581            SlotFieldType::Required => quote! {
582                #field_ident: ::vortex_error::VortexExpect::vortex_expect(
583                    slots[#struct_ident::#const_ident].as_ref(),
584                    #expect_message,
585                )
586            },
587            SlotFieldType::Optional => quote! {
588                #field_ident: slots[#struct_ident::#const_ident].as_ref()
589            },
590            SlotFieldType::VariadicTail => quote! {
591                #field_ident: ::vortex_array::SlotSlice::new(
592                    &slots[#struct_ident::#const_ident..],
593                    #expect_message,
594                )
595            },
596        }
597    }
598
599    fn view_to_owned(&self) -> proc_macro2::TokenStream {
600        let field_ident = &self.field_ident;
601
602        match self.slot_type {
603            SlotFieldType::Required => quote! {
604                #field_ident: ::std::clone::Clone::clone(self.#field_ident)
605            },
606            SlotFieldType::Optional => quote! {
607                #field_ident: self.#field_ident.cloned()
608            },
609            SlotFieldType::VariadicTail => quote! {
610                #field_ident: self.#field_ident.to_vec()
611            },
612        }
613    }
614
615    fn owned_from_slots(&self) -> proc_macro2::TokenStream {
616        let field_ident = &self.field_ident;
617        let struct_ident = &self.struct_ident;
618        let const_ident = &self.const_ident;
619        let expect_message = &self.expect_message;
620
621        match self.slot_type {
622            SlotFieldType::Required => quote! {
623                #field_ident: ::vortex_error::VortexExpect::vortex_expect(
624                    slots[#struct_ident::#const_ident].take(),
625                    #expect_message,
626                )
627            },
628            SlotFieldType::Optional => quote! {
629                #field_ident: slots[#struct_ident::#const_ident].take()
630            },
631            SlotFieldType::VariadicTail => {
632                unreachable!("variadic tail is drained before fixed fields")
633            }
634        }
635    }
636
637    fn storage_slot(&self) -> proc_macro2::TokenStream {
638        let field_ident = &self.field_ident;
639
640        match self.slot_type {
641            SlotFieldType::Required => quote! {
642                Some(self.#field_ident)
643            },
644            SlotFieldType::Optional => quote! {
645                self.#field_ident
646            },
647            SlotFieldType::VariadicTail => {
648                unreachable!("variadic tail is appended after fixed fields")
649            }
650        }
651    }
652
653    fn ext_method(&self) -> proc_macro2::TokenStream {
654        let field_ident = &self.field_ident;
655        let struct_ident = &self.struct_ident;
656        let const_ident = &self.const_ident;
657        let expect_message = &self.expect_message;
658
659        match self.slot_type {
660            SlotFieldType::Required => quote! {
661                #[inline]
662                fn #field_ident(&self) -> &::vortex_array::ArrayRef {
663                    ::vortex_error::VortexExpect::vortex_expect(
664                        self.as_ref().slots()[#struct_ident::#const_ident].as_ref(),
665                        #expect_message,
666                    )
667                }
668            },
669            SlotFieldType::Optional => quote! {
670                #[inline]
671                fn #field_ident(&self) -> Option<&::vortex_array::ArrayRef> {
672                    self.as_ref().slots()[#struct_ident::#const_ident].as_ref()
673                }
674            },
675            SlotFieldType::VariadicTail => quote! {
676                #[inline]
677                fn #field_ident(&self) -> ::vortex_array::SlotSlice<'_> {
678                    ::vortex_array::SlotSlice::new(
679                        &self.as_ref().slots()[#struct_ident::#const_ident..],
680                        #expect_message,
681                    )
682                }
683            },
684        }
685    }
686}
687
688/// A parsed `#[slot(N)]` or `#[slot(N..)]` field annotation.
689struct SlotIndexAttr {
690    index: usize,
691    /// Whether the trailing `..` was present, marking a variadic tail.
692    variadic: bool,
693    span: Span,
694}
695
696impl SlotIndexAttr {
697    fn from_field(field: &Field, field_name: &str) -> syn::Result<Self> {
698        let mut annotation = None;
699        for attr in &field.attrs {
700            if !attr.path().is_ident(SLOT_ATTR) {
701                continue;
702            }
703            if annotation.is_some() {
704                return Err(syn::Error::new(
705                    attr.span(),
706                    format!("`{field_name}` has more than one `#[slot(..)]` attribute"),
707                ));
708            }
709            annotation = Some(attr.parse_args::<Self>()?);
710        }
711
712        annotation.ok_or_else(|| {
713            syn::Error::new(
714                field.span(),
715                format!(
716                    "`{field_name}` is missing a `#[slot(N)]` attribute; every field of an \
717                     `#[array_slots]` struct must pin itself to a slot index so that reordering \
718                     field declarations cannot change the slot layout"
719                ),
720            )
721        })
722    }
723}
724
725impl Parse for SlotIndexAttr {
726    fn parse(input: ParseStream) -> syn::Result<Self> {
727        let literal: LitInt = input.parse()?;
728        let index = literal.base10_parse::<usize>()?;
729        let variadic = input.peek(Token![..]);
730        if variadic {
731            input.parse::<Token![..]>()?;
732        }
733        if !input.is_empty() {
734            return Err(input.error("expected `#[slot(N)]` or `#[slot(N..)]`"));
735        }
736
737        Ok(Self {
738            index,
739            variadic,
740            span: literal.span(),
741        })
742    }
743}
744
745#[derive(Clone, Copy)]
746enum SlotFieldType {
747    Required,
748    Optional,
749    VariadicTail,
750}
751
752impl SlotFieldType {
753    fn from_syn_type(ty: &Type) -> syn::Result<Self> {
754        if is_array_ref_type(ty) {
755            return Ok(Self::Required);
756        }
757
758        if let Some(inner_ty) = wrapper_inner_type(ty, "Option")
759            && is_array_ref_type(inner_ty)
760        {
761            return Ok(Self::Optional);
762        }
763
764        if let Some(inner_ty) = wrapper_inner_type(ty, "Vec")
765            && is_array_ref_type(inner_ty)
766        {
767            return Ok(Self::VariadicTail);
768        }
769
770        Err(syn::Error::new(
771            ty.span(),
772            "#[array_slots] fields must be ArrayRef, Option<ArrayRef>, or Vec<ArrayRef>",
773        ))
774    }
775
776    fn view_field_ty(self) -> proc_macro2::TokenStream {
777        match self {
778            Self::Required => quote! { &'a ::vortex_array::ArrayRef },
779            Self::Optional => quote! { Option<&'a ::vortex_array::ArrayRef> },
780            Self::VariadicTail => quote! { ::vortex_array::SlotSlice<'a> },
781        }
782    }
783}
784
785fn is_array_ref_type(ty: &Type) -> bool {
786    matches!(
787        ty,
788        Type::Path(type_path)
789            if type_path.qself.is_none()
790                && type_path
791                    .path
792                    .segments
793                    .last()
794                    .is_some_and(|segment| segment.ident == "ArrayRef")
795    )
796}
797
798fn wrapper_inner_type<'a>(ty: &'a Type, wrapper: &str) -> Option<&'a Type> {
799    let Type::Path(type_path) = ty else {
800        return None;
801    };
802    let segment = type_path.path.segments.last()?;
803    if segment.ident != wrapper {
804        return None;
805    }
806
807    let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
808        return None;
809    };
810
811    match args.args.first()? {
812        syn::GenericArgument::Type(inner_ty) => Some(inner_ty),
813        _ => None,
814    }
815}
816
817fn ident_name(ident: &Ident) -> String {
818    ident.to_string().trim_start_matches("r#").to_owned()
819}
820
821fn to_screaming_snake_case(name: &str) -> String {
822    let mut result = String::with_capacity(name.len());
823    let mut prev_is_lower_or_digit = false;
824
825    for ch in name.chars() {
826        if ch.is_ascii_uppercase() && prev_is_lower_or_digit {
827            result.push('_');
828        }
829        result.push(ch.to_ascii_uppercase());
830        prev_is_lower_or_digit = ch.is_ascii_lowercase() || ch.is_ascii_digit();
831    }
832
833    result
834}
vortex_array_macros/lib.rs

vortex_array_macros/
lib.rs
