Skip to main content

prgpu_macro/
lib.rs

1use proc_macro::TokenStream;
2
3mod diagnostics;
4mod generate;
5mod layout;
6mod parse;
7mod types;
8
9use types::GpuType;
10
11#[proc_macro_attribute]
12pub fn gpu_struct(attr: TokenStream, item: TokenStream) -> TokenStream {
13    let attr_tokens: proc_macro2::TokenStream = attr.into();
14    let item_tokens: proc_macro2::TokenStream = item.into();
15
16    let config = match parse::parse_config(&attr_tokens) {
17        Ok(c) => c,
18        Err(e) => return e.to_compile_error().into(),
19    };
20
21    let mut item_struct: syn::ItemStruct = match syn::parse2(item_tokens.clone()) {
22        Ok(s) => s,
23        Err(e) => return e.to_compile_error().into(),
24    };
25
26    if !item_struct.generics.params.is_empty() {
27        return syn::Error::new(
28            item_struct.ident.span(),
29            "#[gpu_struct] does not support generic structs; \
30             remove generic parameters or use a concrete type",
31        )
32        .to_compile_error()
33        .into();
34    }
35
36    if let Err(e) = diagnostics::validate_repr(&item_struct, &config) {
37        return e.to_compile_error().into();
38    }
39
40    let fields = match &item_struct.fields {
41        syn::Fields::Named(f) => &f.named,
42        syn::Fields::Unit => {
43            return syn::Error::new(
44                item_struct.ident.span(),
45                "#[gpu_struct] cannot be applied to unit structs; add at least one field",
46            )
47            .to_compile_error()
48            .into();
49        }
50        _ => {
51            return syn::Error::new(
52                item_struct.ident.span(),
53                "#[gpu_struct] only supports structs with named fields",
54            )
55            .to_compile_error()
56            .into();
57        }
58    };
59
60    let mut resolved_fields: Vec<(syn::Ident, GpuType, proc_macro2::Span)> = Vec::new();
61    for field in fields {
62        let field_name = field.ident.clone().unwrap();
63        let field_span = syn::spanned::Spanned::span(&field.ty);
64
65        let is_gpu_nested = field.attrs.iter().any(|attr| {
66            attr.path().segments.len() == 1 && attr.path().segments[0].ident == "gpu_nested"
67        });
68
69        match types::resolve_type(&field.ty, &config, is_gpu_nested) {
70            Ok(gpu_type) => resolved_fields.push((field_name, gpu_type, field_span)),
71            Err(e) => return e.to_compile_error().into(),
72        }
73    }
74
75    let field_layouts: Vec<_> = resolved_fields
76        .iter()
77        .map(|(name, gpu_type, _)| (name.clone(), gpu_type.clone()))
78        .collect();
79
80    let struct_layout = layout::compute_layout(&field_layouts, config.align, &config.targets);
81
82    let output = generate::generate(&mut item_struct, &config, &struct_layout, &resolved_fields);
83
84    output.into()
85}