1use proc_macro::TokenStream;
39use proc_macro2::TokenStream as TS;
40use quote::{quote, ToTokens};
41use syn::{parse_macro_input, Data, DeriveInput, Fields, Ident, Index, Type};
42
43#[proc_macro_derive(Alga, attributes(alga))]
45pub fn derive_alga(input: TokenStream) -> TokenStream {
46 let input = parse_macro_input!(input as DeriveInput);
47 match expand(&input) {
48 Ok(ts) => ts.into(),
49 Err(e) => e.into_compile_error().into(),
50 }
51}
52
53#[derive(PartialEq, Clone, Copy)]
55enum Level {
56 Monoid,
57 Group,
58 Ring,
59 Field,
60}
61
62impl Level {
63 fn parse(input: &DeriveInput) -> syn::Result<Self> {
64 let mut level = Level::Ring;
65 for attr in &input.attrs {
66 if attr.path().is_ident("alga") {
67 attr.parse_nested_meta(|meta| {
68 if meta.path.is_ident("level") {
69 let value = meta.value()?;
70 let s: syn::LitStr = value.parse()?;
71 level = match s.value().as_str() {
72 "Monoid" => Level::Monoid,
73 "Group" => Level::Group,
74 "Ring" => Level::Ring,
75 "Field" => Level::Field,
76 other => {
77 return Err(meta.error(format!(
78 "unknown alga level `{other}` (expected Monoid, Group, Ring or Field)"
79 )))
80 }
81 };
82 Ok(())
83 } else {
84 Err(meta.error("unknown `#[alga(...)]` option (only `level` is supported)"))
85 }
86 })?;
87 }
88 }
89 Ok(level)
90 }
91}
92
93struct FieldInfo {
95 ty: Type,
96 suffix: TS,
97 name: Option<Ident>,
98}
99
100fn expand(input: &DeriveInput) -> syn::Result<TS> {
101 let name = &input.ident;
102 let generics = &input.generics;
103 let (impl_g, ty_g, where_g) = generics.split_for_impl();
104
105 let fields = match &input.data {
106 Data::Struct(s) => fields_of(&s.fields)?,
107 Data::Enum(_) | Data::Union(_) => {
108 return Err(syn::Error::new_spanned(input, "#[derive(Alga)] supports structs only"))
109 }
110 };
111
112 let level = Level::parse(input)?;
113 let g = Gen { name, impl_g: &impl_g, ty_g: &ty_g, where_g };
114 let mut out = TS::new();
115 let add_max = if level == Level::Monoid { "Monoid" } else { "Group" };
119 out.extend(g.single_op_chain(&fields, "Additive", add_max));
120 out.extend(g.single_op_chain(&fields, "Multiplicative", "Monoid"));
121 if matches!(level, Level::Ring | Level::Field) {
122 out.extend(g.two_op_chain(&fields, level));
123 out.extend(g.module_level(&fields));
124 }
125 Ok(out)
126}
127
128fn fields_of(fields: &Fields) -> syn::Result<Vec<FieldInfo>> {
129 let mut out = vec![];
130 match fields {
131 Fields::Named(n) => {
132 for f in &n.named {
133 let id = f.ident.clone().unwrap();
134 out.push(FieldInfo { ty: f.ty.clone(), suffix: id.to_token_stream(), name: Some(id) });
135 }
136 }
137 Fields::Unnamed(u) => {
138 for (i, f) in u.unnamed.iter().enumerate() {
139 let idx = Index::from(i);
140 out.push(FieldInfo { ty: f.ty.clone(), suffix: idx.to_token_stream(), name: None });
141 }
142 }
143 Fields::Unit => {
144 return Err(syn::Error::new_spanned(fields, "#[derive(Alga)] needs at least one field"))
145 }
146 }
147 if out.is_empty() {
148 return Err(syn::Error::new_spanned(fields, "#[derive(Alga)] needs at least one field"));
149 }
150 Ok(out)
151}
152
153fn self_access(f: &FieldInfo) -> TS {
155 let s = &f.suffix;
156 quote!(self.#s)
157}
158
159fn rhs_access(f: &FieldInfo) -> TS {
161 let s = &f.suffix;
162 quote!(rhs.#s)
163}
164
165fn construct(fields: &[FieldInfo], exprs: &[TS]) -> TS {
167 match fields.first().map(|f| f.name.is_some()) {
168 Some(true) => {
169 let pairs = fields.iter().zip(exprs).map(|(f, e)| {
170 let n = f.name.as_ref().unwrap();
171 quote!(#n: #e)
172 });
173 quote!(Self { #(#pairs),* })
174 }
175 _ => quote!(Self( #(#exprs),* )),
176 }
177}
178
179fn field_types(fields: &[FieldInfo]) -> Vec<&Type> {
182 let mut seen: Vec<String> = vec![];
183 fields
184 .iter()
185 .map(|f| &f.ty)
186 .filter(|t| {
187 let key = t.to_token_stream().to_string();
188 if seen.contains(&key) {
189 false
190 } else {
191 seen.push(key);
192 true
193 }
194 })
195 .collect()
196}
197
198fn op_path(op: &str) -> TS {
199 let op = Ident::new(op, proc_macro2::Span::call_site());
200 quote!(::alga2::op::#op)
201}
202
203fn trait_path(name: &str) -> TS {
204 let name = Ident::new(name, proc_macro2::Span::call_site());
205 quote!(::alga2::tower::#name)
206}
207
208struct Gen<'a> {
209 name: &'a Ident,
210 impl_g: &'a syn::ImplGenerics<'a>,
211 ty_g: &'a syn::TypeGenerics<'a>,
212 where_g: Option<&'a syn::WhereClause>,
213}
214
215impl Gen<'_> {
216 fn where_clause(&self, extra: &[TS]) -> TS {
219 let mut preds: Vec<syn::WherePredicate> = Vec::new();
220 if let Some(w) = self.where_g {
221 preds.extend(w.predicates.iter().cloned());
222 }
223 for e in extra {
224 preds.push(syn::parse_quote!(#e));
225 }
226 if preds.is_empty() {
227 quote!()
228 } else {
229 quote!(where #(#preds),*)
230 }
231 }
232
233 fn bounds(&self, fields: &[FieldInfo], trait_name: &str, op: &str) -> Vec<TS> {
235 let tr = trait_path(trait_name);
236 let op = op_path(op);
237 field_types(fields).iter().map(|ty| quote!(#ty: #tr<#op>)).collect()
238 }
239
240 fn single_op_chain(&self, fields: &[FieldInfo], op: &str, max_level: &str) -> TS {
242 let n = self.name;
243 let impl_g = self.impl_g;
244 let ty_g = self.ty_g;
245 let op_ts = op_path(op);
246 let mag = trait_path("Magma");
247 let semi = trait_path("Semigroup");
248 let quasi = trait_path("Quasigroup");
249 let mon = trait_path("Monoid");
250 let loop_t = trait_path("Loop");
251 let grp = trait_path("Group");
252 let abel = trait_path("AbelianGroup");
253
254 let combine_exprs = fields.iter().map(|f| {
255 let a = self_access(f);
256 let b = rhs_access(f);
257 let ty = &f.ty;
258 quote!(<#ty as #mag<#op_ts>>::combine(&#a, &#b))
259 }).collect::<Vec<_>>();
260 let identity_exprs = fields.iter().map(|f| {
261 let ty = &f.ty;
262 quote!(<#ty as #mon<#op_ts>>::identity())
263 }).collect::<Vec<_>>();
264 let inverse_exprs = fields.iter().map(|f| {
265 let a = self_access(f);
266 quote!(#a.inverse())
267 }).collect::<Vec<_>>();
268
269 let magma_b = self.bounds(fields, "Magma", op);
270 let monoid_b = self.bounds(fields, "Monoid", op);
271 let group_b = self.bounds(fields, "Group", op);
272 let w_magma = self.where_clause(&magma_b);
273 let w_monoid = self.where_clause(&monoid_b);
274 let w_group = self.where_clause(&group_b);
275
276 let c_combine = construct(fields, &combine_exprs);
277 let c_identity = construct(fields, &identity_exprs);
278 let c_inverse = construct(fields, &inverse_exprs);
279
280 let quasigroup_ladder = (op == "Additive").then(|| quote! {
283 impl #impl_g #quasi<#op_ts> for #n #ty_g
284 #w_magma
285 {}
286 impl #impl_g #loop_t<#op_ts> for #n #ty_g
287 #w_monoid
288 {}
289 });
290
291 let group_ladder = (max_level == "Group").then(|| quote! {
293 impl #impl_g #grp<#op_ts> for #n #ty_g
294 #w_group
295 {
296 fn inverse(&self) -> Self { #c_inverse }
297 }
298 impl #impl_g #abel<#op_ts> for #n #ty_g
299 #w_group
300 {}
301 });
302
303 quote! {
304 impl #impl_g #mag<#op_ts> for #n #ty_g
305 #w_magma
306 {
307 fn combine(&self, rhs: &Self) -> Self { #c_combine }
308 }
309 impl #impl_g #semi<#op_ts> for #n #ty_g
310 #w_magma
311 {}
312 #quasigroup_ladder
313 impl #impl_g #mon<#op_ts> for #n #ty_g
314 #w_monoid
315 {
316 fn identity() -> Self { #c_identity }
317 }
318 #group_ladder
319 }
320 }
321
322 fn two_op_chain(&self, fields: &[FieldInfo], level: Level) -> TS {
325 let n = self.name;
326 let impl_g = self.impl_g;
327 let ty_g = self.ty_g;
328
329 let semiring_b = field_types(fields).iter().map(|ty| {
330 quote!(#ty: ::alga2::tower::Semiring<::alga2::op::Additive, ::alga2::op::Multiplicative>)
331 }).collect::<Vec<_>>();
332 let field_b = field_types(fields).iter().map(|ty| {
333 quote!(#ty: ::alga2::tower::Field<::alga2::op::Additive, ::alga2::op::Multiplicative>)
334 }).collect::<Vec<_>>();
335
336 let comm_b = if level == Level::Field {
339 &field_b
340 } else {
341 &semiring_b
342 };
343 let w_semiring = self.where_clause(&semiring_b);
344 let w_comm = self.where_clause(comm_b);
345 let w_field = self.where_clause(&field_b);
346
347 let inv_exprs = fields.iter().map(|f| {
348 let a = self_access(f);
349 quote!(#a.inv())
350 }).collect::<Vec<_>>();
351 let c_inv = construct(fields, &inv_exprs);
352
353 let field_ladder = (level == Level::Field).then(|| quote! {
355 impl #impl_g ::alga2::tower::DivisionRing<::alga2::op::Additive, ::alga2::op::Multiplicative> for #n #ty_g
356 #w_field
357 {
358 fn inv(&self) -> Self { #c_inv }
359 }
360 impl #impl_g ::alga2::tower::Field<::alga2::op::Additive, ::alga2::op::Multiplicative> for #n #ty_g
361 #w_field
362 {}
363 });
364
365 quote! {
366 impl #impl_g ::alga2::tower::Semiring<::alga2::op::Additive, ::alga2::op::Multiplicative> for #n #ty_g
367 #w_semiring
368 {}
369 impl #impl_g ::alga2::tower::Ring<::alga2::op::Additive, ::alga2::op::Multiplicative> for #n #ty_g
370 #w_semiring
371 {}
372 impl #impl_g ::alga2::tower::CommutativeRing<::alga2::op::Additive, ::alga2::op::Multiplicative> for #n #ty_g
373 #w_comm
374 {}
375 #field_ladder
376 }
377 }
378
379 fn module_level(&self, fields: &[FieldInfo]) -> TS {
381 let n = self.name;
382 let impl_g = self.impl_g;
383 let ty_g = self.ty_g;
384 let first = fields.first().unwrap();
385
386 let module_b = field_types(fields).iter().map(|ty| {
387 quote!(#ty: ::alga2::tower::Module<::alga2::op::Additive, ::alga2::op::Multiplicative>)
388 }).collect::<Vec<_>>();
389 let w_module = self.where_clause(&module_b);
390
391 let scale_exprs = fields.iter().map(|f| {
392 let s_ = &f.suffix;
393 let a = quote!(v.#s_);
394 quote!(::alga2::tower::Module::<::alga2::op::Additive, ::alga2::op::Multiplicative>::scale(s, #a))
395 }).collect::<Vec<_>>();
396 let c_scale = construct(fields, &scale_exprs);
397
398 let ft = &first.ty;
399 quote! {
400 impl #impl_g ::alga2::tower::Module<::alga2::op::Additive, ::alga2::op::Multiplicative> for #n #ty_g
401 #w_module
402 {
403 type Scalar = <#ft as ::alga2::tower::Module<::alga2::op::Additive, ::alga2::op::Multiplicative>>::Scalar;
404 fn scale(s: &Self::Scalar, v: Self) -> Self { #c_scale }
405 }
406 }
407 }
408}