Skip to main content

miniextendr_macros/
typed_dataframe.rs

1//! Parser and expansion for the `typed_dataframe!` macro.
2//!
3//! Generates a struct that wraps an R `data.frame`, validates declared
4//! columns via `TryFromSexp`, and exposes borrowed per-column accessors.
5//!
6//! # Syntax
7//!
8//! ```ignore
9//! typed_dataframe! {
10//!     /// The shape we accept for the Theoph PK dataset.
11//!     pub TheophDf {
12//!         subject: i32,
13//!         weight: f64,
14//!         dose: f64,
15//!         flag: Option<i32>,   // optional column
16//!     }
17//! }
18//! ```
19//!
20//! Expands to a struct `TheophDf` with `TryFromSexp` and per-column
21//! accessors (`subject() -> &[i32]`, `flag() -> Option<&[i32]>`).
22
23use proc_macro2::TokenStream;
24use quote::{format_ident, quote};
25use syn::parse::{Parse, ParseStream};
26use syn::punctuated::Punctuated;
27use syn::{Attribute, Ident, Token, Type, TypePath, Visibility};
28
29/// Parsed input for `typed_dataframe!`.
30pub struct TypedDataframeInput {
31    /// Outer attributes (doc comments etc.) on the struct.
32    pub attrs: Vec<Attribute>,
33    /// Visibility (e.g. `pub`).
34    pub vis: Visibility,
35    /// Struct name.
36    pub name: Ident,
37    /// Declared columns.
38    pub fields: Vec<TypedDataframeField>,
39    /// If `false`, the input data.frame may not have extra (un-declared) columns.
40    /// Default is `true`. Toggle with leading `@exact;` prefix.
41    pub allow_extra: bool,
42}
43
44/// A single column declaration.
45pub struct TypedDataframeField {
46    /// Field attributes (doc comments etc.).
47    pub attrs: Vec<Attribute>,
48    /// Column name (as used in R `names(df)`).
49    pub name: Ident,
50    /// Element type (e.g. `i32`, `f64`).
51    pub elem_ty: Type,
52    /// True if the field was declared `Option<T>` (column may be absent).
53    pub optional: bool,
54}
55
56impl Parse for TypedDataframeInput {
57    fn parse(input: ParseStream) -> syn::Result<Self> {
58        // Optional @exact; prefix for strict mode.
59        let allow_extra = if input.peek(Token![@]) {
60            input.parse::<Token![@]>()?;
61            let mode: Ident = input.parse()?;
62            if mode != "exact" {
63                return Err(syn::Error::new_spanned(
64                    mode,
65                    "expected `exact` after @; use `@exact;` for strict mode",
66                ));
67            }
68            input.parse::<Token![;]>()?;
69            false
70        } else {
71            true
72        };
73
74        // Parse outer attributes (doc comments etc.).
75        let attrs = input.call(Attribute::parse_outer)?;
76
77        // Visibility (pub, pub(crate), or nothing).
78        let vis: Visibility = input.parse()?;
79
80        // Struct name.
81        let name: Ident = input.parse()?;
82
83        // Brace-delimited field list.
84        let content;
85        syn::braced!(content in input);
86
87        let fields_punct: Punctuated<TypedDataframeField, Token![,]> =
88            Punctuated::parse_terminated(&content)?;
89        let fields: Vec<TypedDataframeField> = fields_punct.into_iter().collect();
90
91        if fields.is_empty() {
92            return Err(syn::Error::new(
93                name.span(),
94                "typed_dataframe! requires at least one column declaration",
95            ));
96        }
97
98        Ok(TypedDataframeInput {
99            attrs,
100            vis,
101            name,
102            fields,
103            allow_extra,
104        })
105    }
106}
107
108impl Parse for TypedDataframeField {
109    fn parse(input: ParseStream) -> syn::Result<Self> {
110        let attrs = input.call(Attribute::parse_outer)?;
111        let name: Ident = input.parse()?;
112        input.parse::<Token![:]>()?;
113        let ty: Type = input.parse()?;
114
115        // Detect `Option<T>` and unwrap to the inner type.
116        let (elem_ty, optional) = match unwrap_option(&ty) {
117            Some(inner) => (inner, true),
118            None => (ty, false),
119        };
120
121        Ok(TypedDataframeField {
122            attrs,
123            name,
124            elem_ty,
125            optional,
126        })
127    }
128}
129
130/// If `ty` is `Option<T>`, return `T`; otherwise `None`.
131fn unwrap_option(ty: &Type) -> Option<Type> {
132    let Type::Path(TypePath { qself: None, path }) = ty else {
133        return None;
134    };
135    // Must be exactly `Option<...>` (last segment) — we don't try to verify
136    // the full path because users may write `std::option::Option<T>`.
137    let segment = path.segments.last()?;
138    if segment.ident != "Option" {
139        return None;
140    }
141    let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
142        return None;
143    };
144    if args.args.len() != 1 {
145        return None;
146    }
147    let syn::GenericArgument::Type(inner) = args.args.first()? else {
148        return None;
149    };
150    Some(inner.clone())
151}
152
153/// Expand a parsed `typed_dataframe!` into the generated struct, impls, and
154/// per-column accessors.
155pub fn expand_typed_dataframe(input: TypedDataframeInput) -> TokenStream {
156    let TypedDataframeInput {
157        attrs,
158        vis,
159        name,
160        fields,
161        allow_extra,
162    } = input;
163
164    let struct_name_str = name.to_string();
165
166    // Per-field generated bits.
167    let col_fields = fields.iter().map(|f| {
168        let col_ident = format_ident!("{}_col", f.name);
169        if f.optional {
170            quote! { #col_ident: ::std::option::Option<::miniextendr_api::SEXP> }
171        } else {
172            quote! { #col_ident: ::miniextendr_api::SEXP }
173        }
174    });
175
176    let accessors = fields.iter().map(|f| {
177        let method_ident = &f.name;
178        let col_ident = format_ident!("{}_col", f.name);
179        let elem_ty = &f.elem_ty;
180        let method_attrs = &f.attrs;
181        if f.optional {
182            quote! {
183                #(#method_attrs)*
184                #[inline]
185                pub fn #method_ident(&self) -> ::std::option::Option<&[#elem_ty]> {
186                    self.#col_ident.map(|col| unsafe {
187                        let len = self.nrow;
188                        let ptr = <#elem_ty as ::miniextendr_api::RNativeType>::dataptr_mut(col)
189                            as *const #elem_ty;
190                        if len == 0 { &[] } else { ::std::slice::from_raw_parts(ptr, len) }
191                    })
192                }
193            }
194        } else {
195            quote! {
196                #(#method_attrs)*
197                #[inline]
198                pub fn #method_ident(&self) -> &[#elem_ty] {
199                    unsafe {
200                        let len = self.nrow;
201                        let ptr = <#elem_ty as ::miniextendr_api::RNativeType>::dataptr_mut(self.#col_ident)
202                            as *const #elem_ty;
203                        if len == 0 { &[] } else { ::std::slice::from_raw_parts(ptr, len) }
204                    }
205                }
206            }
207        }
208    });
209
210    // Per-field validation: walk each declared column, batch errors.
211    let validation_locals = fields.iter().map(|f| {
212        let col_ident = format_ident!("{}_col", f.name);
213        let name_str = f.name.to_string();
214        let elem_ty = &f.elem_ty;
215        if f.optional {
216            quote! {
217                let #col_ident: ::std::option::Option<::miniextendr_api::SEXP> = {
218                    match view.column_raw(#name_str) {
219                        ::std::option::Option::None => ::std::option::Option::None,
220                        ::std::option::Option::Some(col) => {
221                            let actual = ::miniextendr_api::SexpExt::type_of(&col);
222                            let expected = <#elem_ty as ::miniextendr_api::RNativeType>::SEXP_TYPE;
223                            if actual != expected {
224                                __errs.push(::std::format!(
225                                    "column `{}`: expected {} ({:?}), got {:?}",
226                                    #name_str,
227                                    ::std::stringify!(#elem_ty),
228                                    expected,
229                                    actual,
230                                ));
231                                ::std::option::Option::None
232                            } else {
233                                ::std::option::Option::Some(col)
234                            }
235                        }
236                    }
237                };
238            }
239        } else {
240            quote! {
241                let #col_ident: ::miniextendr_api::SEXP = {
242                    match view.column_raw(#name_str) {
243                        ::std::option::Option::None => {
244                            __errs.push(::std::format!(
245                                "missing required column `{}` (expected {} / {:?})",
246                                #name_str,
247                                ::std::stringify!(#elem_ty),
248                                <#elem_ty as ::miniextendr_api::RNativeType>::SEXP_TYPE,
249                            ));
250                            ::miniextendr_api::SEXP::nil()
251                        }
252                        ::std::option::Option::Some(col) => {
253                            let actual = ::miniextendr_api::SexpExt::type_of(&col);
254                            let expected = <#elem_ty as ::miniextendr_api::RNativeType>::SEXP_TYPE;
255                            if actual != expected {
256                                __errs.push(::std::format!(
257                                    "column `{}`: expected {} ({:?}), got {:?}",
258                                    #name_str,
259                                    ::std::stringify!(#elem_ty),
260                                    expected,
261                                    actual,
262                                ));
263                            }
264                            col
265                        }
266                    }
267                };
268            }
269        }
270    });
271
272    let col_idents = fields.iter().map(|f| format_ident!("{}_col", f.name));
273    let col_idents_for_struct = col_idents.clone();
274
275    let declared_names: Vec<String> = fields.iter().map(|f| f.name.to_string()).collect();
276    let declared_count = fields.len();
277    let allow_extra_check = if allow_extra {
278        quote! {}
279    } else {
280        quote! {
281            // Strict mode: reject any column not declared.
282            const __DECLARED: &[&str] = &[#( #declared_names ),*];
283            let mut __extras: ::std::vec::Vec<::std::string::String> = ::std::vec::Vec::new();
284            for col_name in view.names() {
285                if !__DECLARED.contains(&col_name.as_str()) {
286                    __extras.push(col_name.to_string());
287                }
288            }
289            if !__extras.is_empty() {
290                __errs.push(::std::format!(
291                    "unexpected extra columns: {:?}",
292                    __extras
293                ));
294            }
295        }
296    };
297
298    quote! {
299        #(#attrs)*
300        #[allow(non_snake_case)]
301        #vis struct #name {
302            sexp: ::miniextendr_api::SEXP,
303            nrow: usize,
304            #( #col_fields, )*
305        }
306
307        impl #name {
308            /// Number of rows in the underlying data.frame.
309            #[inline]
310            pub fn nrow(&self) -> usize { self.nrow }
311
312            /// Number of declared columns.
313            #[inline]
314            pub fn ncol(&self) -> usize { #declared_count }
315
316            /// The backing SEXP (preserves the data.frame class attribute).
317            #[inline]
318            pub fn as_sexp(&self) -> ::miniextendr_api::SEXP { self.sexp }
319
320            #( #accessors )*
321        }
322
323        impl ::miniextendr_api::from_r::TryFromSexp for #name {
324            type Error = ::miniextendr_api::from_r::SexpError;
325
326            fn try_from_sexp(
327                sexp: ::miniextendr_api::SEXP,
328            ) -> ::std::result::Result<Self, Self::Error> {
329                use ::miniextendr_api::SexpExt as _;
330
331                // 1. Must be a data.frame.
332                if !sexp.is_data_frame() {
333                    return ::std::result::Result::Err(
334                        ::miniextendr_api::from_r::SexpError::InvalidValue(
335                            ::std::format!(
336                                "{}: expected a data.frame, got {:?}",
337                                #struct_name_str,
338                                sexp.type_of(),
339                            )
340                        )
341                    );
342                }
343
344                // 2. Build the view for O(1) name lookup + nrow extraction.
345                let view = ::miniextendr_api::dataframe::DataFrame::from_sexp(sexp)
346                    .map_err(|e| ::miniextendr_api::from_r::SexpError::InvalidValue(
347                        ::std::format!("{}: {}", #struct_name_str, e)
348                    ))?;
349                let nrow = view.nrow();
350
351                // 3. Walk each declared column, batching all errors.
352                let mut __errs: ::std::vec::Vec<::std::string::String> =
353                    ::std::vec::Vec::new();
354
355                #( #validation_locals )*
356
357                #allow_extra_check
358
359                if !__errs.is_empty() {
360                    return ::std::result::Result::Err(
361                        ::miniextendr_api::from_r::SexpError::InvalidValue(
362                            ::std::format!(
363                                "{}: {}",
364                                #struct_name_str,
365                                __errs.join("; ")
366                            )
367                        )
368                    );
369                }
370
371                ::std::result::Result::Ok(#name {
372                    sexp,
373                    nrow,
374                    #( #col_idents_for_struct, )*
375                })
376            }
377        }
378    }
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384
385    #[test]
386    fn test_parse_basic() {
387        let input: TypedDataframeInput = syn::parse_quote! {
388            pub TheophDf {
389                subject: i32,
390                weight: f64,
391            }
392        };
393        assert_eq!(input.name.to_string(), "TheophDf");
394        assert_eq!(input.fields.len(), 2);
395        assert_eq!(input.fields[0].name.to_string(), "subject");
396        assert!(!input.fields[0].optional);
397        assert!(input.allow_extra);
398    }
399
400    #[test]
401    fn test_parse_optional() {
402        let input: TypedDataframeInput = syn::parse_quote! {
403            pub Df { x: i32, y: Option<f64> }
404        };
405        assert!(!input.fields[0].optional);
406        assert!(input.fields[1].optional);
407    }
408
409    #[test]
410    fn test_parse_exact_mode() {
411        let input: TypedDataframeInput = syn::parse_quote! {
412            @exact;
413            pub Df { x: i32 }
414        };
415        assert!(!input.allow_extra);
416    }
417
418    #[test]
419    fn test_parse_with_doc_attrs() {
420        let input: TypedDataframeInput = syn::parse_quote! {
421            /// Theoph PK shape.
422            pub TheophDf {
423                /// Subject id.
424                subject: i32,
425            }
426        };
427        assert!(!input.attrs.is_empty());
428        assert!(!input.fields[0].attrs.is_empty());
429    }
430
431    #[test]
432    fn test_parse_empty_fails() {
433        let result: syn::Result<TypedDataframeInput> = syn::parse2(quote! {
434            pub Df {}
435        });
436        assert!(result.is_err());
437    }
438}