deinterleave_derive/
lib.rs

1use std::collections::HashSet;
2
3use proc_macro::TokenStream;
4use quote::{format_ident, quote};
5use syn::{parse_macro_input, Data, DeriveInput, Error, Index};
6
7#[proc_macro_derive(Deinterleave)]
8pub fn derive_deinterleave(tokens: TokenStream) -> TokenStream {
9  let derive_input = parse_macro_input!(tokens as DeriveInput);
10
11  process_derive_input(derive_input).into()
12}
13
14fn process_derive_input(derive_input: DeriveInput) -> proc_macro2::TokenStream {
15  // we only support structs for now
16  let Data::Struct(data_struct) = derive_input.data else {
17    return Error::new_spanned(derive_input, "expected a struct").into_compile_error();
18  };
19
20  let struct_ident = derive_input.ident;
21
22  let mut seen_field_ty = HashSet::new();
23  let mut field_types = Vec::default();
24  let mut has_fields = Vec::default();
25  let mut push_to = Vec::default();
26  let mut remove_from = Vec::default();
27  let mut swap_remove_from = Vec::default();
28  let mut iter_field_ty = Vec::default();
29  let mut iter_field = Vec::default();
30  let mut iter_mut_field_ty = Vec::default();
31  let mut iter_mut_field = Vec::default();
32
33  for (i, f) in data_struct.fields.into_iter().enumerate() {
34    let index = Index::from(i);
35    let field_ty = f.ty;
36    let field_ident = f.ident;
37
38    if seen_field_ty.contains(&field_ty) {
39      return Error::new_spanned(field_ty.clone(), "type seen more than once (you can only use a type once in the fields of a deinterleaved struct)").to_compile_error();
40    }
41
42    seen_field_ty.insert(field_ty.clone());
43
44    field_types.push(field_ty.clone());
45
46    iter_field_ty.push(quote! {
47      #field_ident: std::slice::Iter<'a, #field_ty>
48    });
49
50    iter_field.push(quote! {
51      #field_ident: data.#index.iter()
52    });
53
54    iter_mut_field_ty.push(quote! {
55      #field_ident: std::slice::IterMut<'a, #field_ty>
56    });
57
58    iter_mut_field.push(quote! {
59      #field_ident: data.#index.iter_mut()
60    });
61
62    has_fields.push(quote! {
63      impl deinterleave::HasField<#field_ty> for #struct_ident {
64        fn get(data: &Self::Data, i: usize) -> Option<&#field_ty> {
65          data.#index.get(i)
66        }
67
68        fn get_mut(data: &mut Self::Data, i: usize) -> Option<&mut #field_ty> {
69          data.#index.get_mut(i)
70        }
71      }
72    });
73
74    push_to.push(quote! {
75      data.#index.push(value.#field_ident)
76    });
77
78    remove_from.push(quote! {
79      #field_ident : data.#index.remove(i)
80    });
81
82    swap_remove_from.push(quote! {
83      #field_ident : data.#index.swap_remove(i)
84    });
85  }
86
87  let iter_ident = format_ident!("{struct_ident}Iter");
88  let iter_mut_ident = format_ident!("{struct_ident}IterMut");
89
90  quote! {
91    #(#has_fields)*
92
93    #[derive(Debug)]
94    pub struct #iter_ident<'a> {
95      // PhantomData required for empty structs
96      _phantom: std::marker::PhantomData<&'a ()>,
97      #(pub #iter_field_ty),*
98    }
99
100    #[derive(Debug)]
101    pub struct #iter_mut_ident<'a> {
102      // PhantomData required for empty structs
103      _phantom: std::marker::PhantomData<&'a mut ()>,
104      #(pub #iter_mut_field_ty),*
105    }
106
107    impl deinterleave::Deinterleave for #struct_ident {
108      type Data = (#(Vec<#field_types>),*);
109      type Iter<'a> = #iter_ident<'a>;
110      type IterMut<'a> = #iter_mut_ident<'a>;
111
112      fn new_data() -> Self::Data {
113        (#(Vec::<#field_types>::default()),*)
114      }
115
116      fn iter(data: &Self::Data) -> Self::Iter<'_> {
117        #iter_ident {
118          _phantom: std::marker::PhantomData,
119          #(#iter_field),*
120        }
121      }
122
123      fn iter_mut(data: &mut Self::Data) -> Self::IterMut<'_> {
124        #iter_mut_ident {
125          _phantom: std::marker::PhantomData,
126          #(#iter_mut_field),*
127        }
128      }
129
130      fn push_to(data: &mut Self::Data, value: Self) {
131        #(#push_to);*
132      }
133
134      fn remove_from(data: &mut Self::Data, i: usize) -> Self {
135        Self {
136          #(#remove_from),*
137        }
138      }
139
140      fn swap_remove_from(data: &mut Self::Data, i: usize) -> Self {
141        Self {
142          #(#swap_remove_from),*
143        }
144      }
145    }
146  }
147}