1use proc_macro2::{Ident, TokenStream};
2use quote::{format_ident, quote};
3use syn::{Attribute, DataStruct, Field, Fields, FieldsNamed, Type};
4
5#[proc_macro_derive(Table, attributes(index, index_set))]
6pub fn derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
7 let ast = syn::parse_macro_input!(input as syn::DeriveInput);
8 let gen: TokenStream = impl_from_args(&ast);
9 gen.into()
10}
11
12fn impl_from_args(input: &syn::DeriveInput) -> TokenStream {
13 match &input.data {
14 syn::Data::Struct(DataStruct {
15 fields: Fields::Named(FieldsNamed { named: fields, .. }),
16 ..
17 }) => {
18 let vis = &input.vis;
19 let name = &input.ident;
20 let table = format_ident!("{}Table", input.ident);
21 let key = format_ident!("{}Key", input.ident);
22
23 let gen: Vec<_> = fields.iter().filter_map(GenField::parse).collect();
24
25 let indexes = gen.iter().map(|f| f.fields(&key));
26 let getters = gen.iter().map(|f| f.find(name, &key));
27 let removers = gen.iter().map(|f| f.remove());
28 let inserters = gen.iter().map(|f| f.insert());
29
30 let clear = gen.iter().map(|f| f.clear());
31 let with_capacity = gen.iter().map(|f| f.with_capacity());
32
33 quote! {
34 ::edb::slotmap::new_key_type! {
35 #vis struct #key;
36 }
37
38 #[derive(::core::default::Default)]
39 #vis struct #table {
40 storage: ::edb::Storage<#key, #name>,
41 #( #indexes )*
42 }
43
44 impl ::core::ops::Index<#key> for #table {
45 type Output = #name;
46
47 fn index(&self, key: #key) -> &Self::Output {
48 &self.storage[key]
49 }
50 }
51
52 impl ::core::ops::IndexMut<#key> for #table {
53 fn index_mut(&mut self, key: #key) -> &mut Self::Output {
54 &mut self.storage[key]
55 }
56 }
57
58 impl #table {
59 pub fn get(&self, key: #key) -> ::core::option::Option<&#name> {
60 self.storage.get(key)
61 }
62
63 pub fn get_mut(&mut self, key: #key) -> ::core::option::Option<&mut #name> {
64 self.storage.get_mut(key)
65 }
66
67 pub fn insert(&mut self, value: #name) -> #key {
68 self.storage.insert_with_key(|key| {
69 #( #inserters )*
70 value
71 })
72 }
73
74 pub fn remove(&mut self, key: #key) -> ::core::option::Option<#name> {
75 self.storage.remove(key).map(|value| {
76 #( #removers )*
77 value
78 })
79 }
80
81 pub fn clear(&mut self) {
82 self.storage.clear();
83 #( #clear )*
84 }
85
86 pub fn with_capacity(capacity: usize) -> Self {
87 Self {
88 storage: ::edb::Storage::with_capacity_and_key(capacity),
89
90 #( #with_capacity )*
91 }
92 }
93
94 #( #getters )*
95 }
96
97 }
98 }
99 _ => TokenStream::new(),
100 }
101}
102
103fn has_ident(attrs: &[Attribute], ident: &str) -> bool {
104 attrs
105 .iter()
106 .filter(|attr| matches!(attr.style, syn::AttrStyle::Outer))
107 .any(|attr| attr.path().is_ident(ident))
108}
109
110struct GenField {
111 name: Ident,
112 ty: Type,
113 index: GenIndex,
114}
115
116enum GenIndex {
117 Index,
118 IndexSet,
119}
120
121impl GenField {
122 fn new(name: Ident, ty: Type, index: GenIndex) -> Self {
123 Self { name, ty, index }
124 }
125
126 fn parse(
127 Field {
128 ident, ty, attrs, ..
129 }: &Field,
130 ) -> Option<Self> {
131 let name = ident.as_ref()?.clone();
132
133 if has_ident(attrs, "index") {
134 let field = GenField::new(name, ty.clone(), GenIndex::Index);
135 return Some(field);
136 }
137
138 if has_ident(attrs, "index_set") {
139 let field = GenField::new(name, ty.clone(), GenIndex::IndexSet);
140 return Some(field);
141 }
142
143 None
144 }
145
146 fn fields(&self, key: &Ident) -> TokenStream {
147 let index = format_ident!("index_{}", self.name);
148 let ty = &self.ty;
149 match self.index {
150 GenIndex::Index => quote! { #index : ::edb::Index<#ty, #key>, },
151 GenIndex::IndexSet => quote! { #index : ::edb::IndexSet<#ty, #key>, },
152 }
153 }
154
155 fn with_capacity(&self) -> TokenStream {
156 let index = format_ident!("index_{}", self.name);
157 match self.index {
158 GenIndex::Index => quote! { #index : ::edb::Index::with_capacity(capacity), },
159 GenIndex::IndexSet => quote! { #index : ::edb::IndexSet::with_capacity(capacity), },
160 }
161 }
162
163 fn clear(&self) -> TokenStream {
164 let index = format_ident!("index_{}", self.name);
165 match self.index {
166 GenIndex::Index => quote! { self.#index.clear(); },
167 GenIndex::IndexSet => quote! { self.#index.clear(); },
168 }
169 }
170
171 fn insert(&self) -> TokenStream {
172 let name = &self.name;
173 let index = format_ident!("index_{}", self.name);
174
175 match self.index {
176 GenIndex::Index => quote! {
177 self.#index.insert(value.#name.clone(), key);
178 },
179 GenIndex::IndexSet => quote! {
180 self.#index.entry(value.#name.clone()).or_default().insert(key);
181 },
182 }
183 }
184
185 fn remove(&self) -> TokenStream {
186 let name = &self.name;
187 let index = format_ident!("index_{}", self.name);
188
189 match self.index {
190 GenIndex::Index => quote! {
191 self.#index.remove(&value.#name);
192 },
193 GenIndex::IndexSet => quote! {
194 if let ::core::option::Option::Some(set) = self.#index.get_mut(&value.#name) {
195 set.remove(&key);
196 if set.is_empty() {
197 self.#index.remove(&value.#name);
198 }
199 };
200 },
201 }
202 }
203
204 fn find(&self, value: &Ident, key: &Ident) -> TokenStream {
205 let name = &self.name;
206 let ty = &self.ty;
207 let index = format_ident!("index_{}", self.name);
208
209 match self.index {
210 GenIndex::Index => {
211 let find_by = format_ident!("find_by_{}", self.name);
212 let key_by = format_ident!("key_by_{}", self.name);
213 let find_by_mut = format_ident!("find_by_{}_mut", self.name);
214 let remove_by = format_ident!("remove_by_{}", self.name);
215
216 quote! {
217 pub fn #find_by<Q>(&self, #name: &Q) -> ::core::option::Option<&#value>
218 where
219 #ty: ::core::borrow::Borrow<Q>,
220 Q: ::core::hash::Hash + ::core::cmp::Eq + ?::core::marker::Sized,
221 {
222 self.#index.get(#name).map(|&key| &self.storage[key])
223 }
224
225 pub fn #key_by<Q>(&self, #name: &Q) -> ::core::option::Option<#key>
226 where
227 #ty: ::core::borrow::Borrow<Q>,
228 Q: ::core::hash::Hash + ::core::cmp::Eq + ?::core::marker::Sized,
229 {
230 self.#index.get(#name).copied()
231 }
232
233 pub fn #find_by_mut<Q>(&mut self, #name: &Q) -> ::core::option::Option<&mut #value>
234 where
235 #ty: ::core::borrow::Borrow<Q>,
236 Q: ::core::hash::Hash + ::core::cmp::Eq + ?::core::marker::Sized,
237 {
238 self.#index.get(#name).map(|&key| &mut self.storage[key])
239 }
240
241 pub fn #remove_by<Q>(&mut self, #name: &Q) -> ::core::option::Option<#value>
242 where
243 #ty: ::core::borrow::Borrow<Q>,
244 Q: ::core::hash::Hash + ::core::cmp::Eq + ?::core::marker::Sized,
245 {
246 self.#index.get(#name).copied().and_then(|key| self.remove(key))
247 }
248 }
249 }
250 GenIndex::IndexSet => {
251 let contains = format_ident!("contains_{}", self.name);
252 let list_by = format_ident!("list_by_{}", self.name);
253 let keys_by = format_ident!("keys_by_{}", self.name);
254
255 quote! {
256 pub fn #contains<Q>(&self, #name: &Q) -> bool
257 where
258 #ty: ::core::borrow::Borrow<Q>,
259 Q: ::core::hash::Hash + ::core::cmp::Eq + ?::core::marker::Sized,
260 {
261 self.#index.contains_key(#name)
262 }
263
264 pub fn #list_by<Q>(&self, #name: &Q) -> ::core::option::Option<impl ::core::iter::Iterator<Item = &#value>>
265 where
266 #ty: ::core::borrow::Borrow<Q>,
267 Q: ::core::hash::Hash + ::core::cmp::Eq + ?::core::marker::Sized,
268 {
269 self.#index.get(#name).map(|keys| keys.iter().map(|&key| &self.storage[key]))
270 }
271
272 pub fn #keys_by<Q>(&self, #name: &Q) -> ::core::option::Option<impl ::core::iter::Iterator<Item = #key> + '_>
273 where
274 #ty: ::core::borrow::Borrow<Q>,
275 Q: ::core::hash::Hash + ::core::cmp::Eq + ?::core::marker::Sized,
276 {
277 self.#index.get(#name).map(|keys| keys.iter().copied())
278 }
279 }
280 }
281 }
282 }
283}