1mod generator;
2use proc_macro::TokenStream as TokenStream1;
3
4use proc_macro2::Span;
5use proc_macro2::TokenStream;
6use proc_macro_error::{abort, proc_macro_error};
7use std::collections::{HashMap, HashSet};
8use syn::parse::Parse;
9use syn::spanned::Spanned;
10use syn::*;
11use template_quote::{quote, ToTokens};
12
13#[derive(Clone, Debug, PartialEq, Eq, Hash)]
14enum TraitTarget {
15 Default,
16 IterMut,
17 IntoIter,
18 Map,
19}
20
21fn squash_minlens(outs: &[Expr]) -> Expr {
22 if outs.is_empty() {
23 abort!(Span::call_site(), "needs one or more variants");
24 }
25 let mut acc = outs[outs.len() - 1].clone();
26 if outs.len() >= 2 {
27 for out in outs[0..(outs.len() - 1)].iter().rev() {
28 acc = parse_quote! {__parametric_type_min(#out, #acc)};
29 }
30 }
31 parse_quote! {
32 {
33 const fn __parametric_type_min(a: usize, b: usize) -> usize {
34 if a < b { a } else { b }
35 }
36 #acc
37 }
38 }
39}
40fn squash_maxlens(outs: &[Expr]) -> Expr {
41 if outs.is_empty() {
42 abort!(Span::call_site(), "needs one or more variants");
43 }
44 let mut acc = outs[outs.len() - 1].clone();
45 if outs.len() >= 2 {
46 for out in outs[0..(outs.len() - 1)].iter().rev() {
47 acc = parse_quote! {__parametric_type_max(#out, #acc)};
48 }
49 }
50 parse_quote! {
51 {
52 const fn __parametric_type_max(a: Option<usize>, b: Option<usize>) -> Option<usize> {
53 match (a, b) {
54 (Some(a), Some(b)) => if a > b { Some(a) } else { Some(b) }
55 _ => None,
56 }
57 }
58 #acc
59 }
60 }
61}
62
63fn replace_type(mut ty: Type, from: Type, to: Type) -> Type {
64 use syn::visit_mut::VisitMut;
65 struct Visitor(Type, Type);
66 impl VisitMut for Visitor {
67 fn visit_type_mut(&mut self, ty: &mut Type) {
68 if ty == &self.0 {
69 *ty = self.1.clone();
70 } else {
71 syn::visit_mut::visit_type_mut(self, ty)
72 }
73 }
74 }
75 Visitor(from, to).visit_type_mut(&mut ty);
76 ty
77}
78
79impl TraitTarget {
80 fn make_enough(mut set: HashSet<Self>) -> HashSet<Self> {
81 if set.contains(&Self::Map) {
82 set.insert(Self::IntoIter);
83 }
84 if set.contains(&Self::IntoIter) {
85 set.insert(Self::Default);
86 }
87 if set.contains(&Self::IterMut) {
88 set.insert(Self::Default);
89 }
90 set
91 }
92
93 #[allow(clippy::too_many_arguments)]
94 fn emit(
95 &self,
96 krate: &Path,
97 ident: &Ident,
98 generics: &Generics,
99 param_index: usize,
100 replacing_ty: &Type,
101 self_val: &Ident,
102 tys_exprs: &[Vec<(Type, Expr)>],
103 mut f: impl FnMut(&[TokenStream]) -> TokenStream,
104 mut emit_map_f: impl FnMut(&[Vec<Expr>]) -> TokenStream,
105 needs_ref: bool,
106 ) -> Result<TokenStream> {
107 let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
108 match self {
109 Self::Default => {
110 let out_len = tys_exprs
111 .iter()
112 .map(|item| {
113 Ok(generator::EmitContext {
114 kind: generator::EmitLen,
115 krate: krate.clone(),
116 replacing_ty: replacing_ty.clone(),
117 }
118 .emit_for_tys_exprs(item.iter().map(|(a, b)| {
119 (
120 a.clone(),
121 if needs_ref {
122 parse_quote! {&#b}
123 } else {
124 parse_quote! {#b}
125 },
126 )
127 }))?
128 .unwrap_or(parse_quote!(0usize)))
129 })
130 .collect::<Result<Vec<_>>>()?
131 .into_iter()
132 .map(|expr| quote!(#expr))
133 .collect::<Vec<_>>();
134 let out_max_len = tys_exprs
135 .iter()
136 .map(|item| {
137 Ok(generator::EmitContext {
138 kind: generator::EmitMaxLen,
139 krate: krate.clone(),
140 replacing_ty: replacing_ty.clone(),
141 }
142 .emit_for_tys_exprs(item.iter().map(|(a, b)| {
143 (
144 a.clone(),
145 if needs_ref {
146 parse_quote!(&#b)
147 } else {
148 parse_quote!(#b)
149 },
150 )
151 }))?
152 .unwrap_or(parse_quote!(::core::option::Option::Some(0usize))))
153 })
154 .collect::<Result<Vec<_>>>()?;
155 let out_min_len = tys_exprs
156 .iter()
157 .map(|item| {
158 Ok(generator::EmitContext {
159 kind: generator::EmitMinLen,
160 krate: krate.clone(),
161 replacing_ty: replacing_ty.clone(),
162 }
163 .emit_for_tys_exprs(item.iter().map(|(a, b)| {
164 (
165 a.clone(),
166 if needs_ref {
167 parse_quote!(&#b)
168 } else {
169 parse_quote!(#b)
170 },
171 )
172 }))?
173 .unwrap_or(parse_quote!(0usize)))
174 })
175 .collect::<Result<Vec<_>>>()?;
176 let iter_ty_lt: Lifetime = parse_quote!('__parametrized_lt);
177 let out_iter_ty = tys_exprs
178 .iter()
179 .map(|item| {
180 Ok(generator::EmitContext {
181 kind: generator::EmitIterTy(iter_ty_lt.clone(), replacing_ty.clone()),
182 krate: krate.clone(),
183 replacing_ty: replacing_ty.clone(),
184 }
185 .emit_for_tys_exprs(
186 item.iter().map(|(a, _)| (a.clone(), replacing_ty.clone())),
187 )?
188 .unwrap_or(parse_quote!(::core::iter::Empty<&#iter_ty_lt #replacing_ty>)))
189 })
190 .collect::<Result<Vec<_>>>()?;
191 let out_iter = tys_exprs
192 .iter()
193 .map(|item| {
194 Ok(generator::EmitContext {
195 kind: generator::EmitIter,
196 krate: krate.clone(),
197 replacing_ty: replacing_ty.clone(),
198 }
199 .emit_for_tys_exprs(item.iter().map(|(a, b)| {
200 (
201 a.clone(),
202 if needs_ref {
203 parse_quote!(&#b)
204 } else {
205 parse_quote!(#b)
206 },
207 )
208 }))?
209 .unwrap_or(parse_quote!(::core::iter::empty())))
210 })
211 .collect::<Result<Vec<_>>>()?
212 .into_iter()
213 .zip(&out_iter_ty)
214 .map(|(expr, ty)| {
215 if tys_exprs.len() > 1 {
216 quote!(sumtype!(
217 #expr,
218 for<#iter_ty_lt> #ty
219 where
220 #(for tp in &generics.params) {
221 #(if let GenericParam::Type(tp) = tp) {
222 #{&tp.ident}: #iter_ty_lt,
223 }
224 }
225 ))
226 } else {
227 quote!(#expr)
228 }
229 })
230 .collect::<Vec<_>>();
231 Ok(quote! {
232 #(if tys_exprs.len() > 1) {
233 #[#krate::_imp::sumtype::sumtype(#krate::_imp::sumtype::traits::Iterator)]
234 }
235 impl #impl_generics #krate::Parametrized<#param_index> for #ident
236 #ty_generics #where_clause {
237 type Item = #replacing_ty;
238 const MIN_LEN: usize = #{squash_minlens(out_min_len.as_slice())};
239 const MAX_LEN: Option<usize> = #{squash_maxlens(out_max_len.as_slice())};
240 fn param_len(&#self_val) -> usize {
241 #{f(out_len.as_slice())}
242 }
243 #(if tys_exprs.len() > 1) {
244 type Iter<#iter_ty_lt> = sumtype![#iter_ty_lt] where (Self, Self::Item): #iter_ty_lt;
245 } #(else) {
246 type Iter<#iter_ty_lt> = #(#out_iter_ty)* where (Self, Self::Item): #iter_ty_lt;
247 }
248 fn param_iter<'__parametrized_lt>(&'__parametrized_lt #self_val) -> Self::Iter<'__parametrized_lt>
249 where
250 Self::Item: '__parametrized_lt
251 {
252 #{f(out_iter.as_slice())}
253 }
254 }
255 })
256 }
257 Self::IterMut => {
258 let iter_ty_lt: Lifetime = parse_quote!('__parametrized_lt);
259 let out_iter_mut_ty = tys_exprs
260 .iter()
261 .map(|item| {
262 Ok(generator::EmitContext {
263 kind: generator::EmitIterMutTy(
264 iter_ty_lt.clone(),
265 replacing_ty.clone(),
266 ),
267 krate: krate.clone(),
268 replacing_ty: replacing_ty.clone(),
269 }
270 .emit_for_tys_exprs(
271 item.iter().map(|(a, _)| (a.clone(), replacing_ty.clone())),
272 )?
273 .unwrap_or(
274 parse_quote!(::core::iter::Empty<& #iter_ty_lt mut #replacing_ty>),
275 ))
276 })
277 .collect::<Result<Vec<_>>>()?;
278 let out_iter_mut = tys_exprs
279 .iter()
280 .map(|item| {
281 Ok(generator::EmitContext {
282 kind: generator::EmitIterMut,
283 krate: krate.clone(),
284 replacing_ty: replacing_ty.clone(),
285 }
286 .emit_for_tys_exprs(item.iter().map(|(a, b)| {
287 (
288 a.clone(),
289 if needs_ref {
290 parse_quote!(&mut #b)
291 } else {
292 parse_quote!(#b)
293 },
294 )
295 }))?
296 .unwrap_or(parse_quote!(::core::iter::empty())))
297 })
298 .collect::<Result<Vec<_>>>()?
299 .into_iter()
300 .zip(&out_iter_mut_ty)
301 .map(|(expr, ty)| {
302 if tys_exprs.len() > 1 {
303 quote!(sumtype!(
304 #expr,
305 for<#iter_ty_lt> #ty
306 where
307 #(for tp in &generics.params) {
308 #(if let GenericParam::Type(tp) = tp) {
309 #{&tp.ident}: #iter_ty_lt,
310 }
311 }
312 ))
313 } else {
314 quote!(#expr)
315 }
316 })
317 .collect::<Vec<_>>();
318 Ok(quote! {
319 #(if tys_exprs.len() > 1) {
320 #[#krate::_imp::sumtype::sumtype(#krate::_imp::sumtype::traits::Iterator)]
321 }
322 impl #impl_generics #krate::ParametrizedIterMut<#param_index> for #ident #ty_generics #where_clause {
323 #(if tys_exprs.len() > 1) {
324 type IterMut<#iter_ty_lt> = sumtype![#iter_ty_lt] where (Self, Self::Item): #iter_ty_lt;
325 } #(else) {
326 type IterMut<#iter_ty_lt> = #(#out_iter_mut_ty)* where (Self, Self::Item): #iter_ty_lt;
327 }
328 fn param_iter_mut<'__parametrized_lt>(&'__parametrized_lt mut #self_val) -> Self::IterMut<'__parametrized_lt>
329 where
330 Self::Item: '__parametrized_lt
331 {
332 #{f(out_iter_mut.as_slice())}
333 }
334 }
335 })
336 }
337 Self::IntoIter => {
338 let out_into_iter_ty = tys_exprs
339 .iter()
340 .map(|item| {
341 Ok(generator::EmitContext {
342 kind: generator::EmitIntoIterTy(replacing_ty.clone()),
343 krate: krate.clone(),
344 replacing_ty: replacing_ty.clone(),
345 }
346 .emit_for_tys_exprs(
347 item.iter().map(|(a, _)| (a.clone(), replacing_ty.clone())),
348 )?
349 .unwrap_or(parse_quote!(::core::iter::Empty<#replacing_ty>)))
350 })
351 .collect::<Result<Vec<_>>>()?;
352 let out_into_iter = tys_exprs
353 .iter()
354 .map(|item| {
355 Ok(generator::EmitContext {
356 kind: generator::EmitIntoIter,
357 krate: krate.clone(),
358 replacing_ty: replacing_ty.clone(),
359 }
360 .emit_for_tys_exprs(item.iter().map(|(a, b)| (a.clone(), b.clone())))?
361 .unwrap_or(parse_quote!(::core::iter::empty())))
362 })
363 .collect::<Result<Vec<_>>>()?
364 .into_iter()
365 .zip(&out_into_iter_ty)
366 .map(|(expr, ty)| {
367 if tys_exprs.len() > 1 {
368 quote!(sumtype!(#expr, #ty))
369 } else {
370 quote!(#expr)
371 }
372 })
373 .collect::<Vec<_>>();
374 Ok(quote! {
375 #(if tys_exprs.len() > 1) {
376 #[#krate::_imp::sumtype::sumtype(#krate::_imp::sumtype::traits::Iterator)]
377 }
378 impl #impl_generics #krate::ParametrizedIntoIter<#param_index> for #ident #ty_generics #where_clause {
379 #(if tys_exprs.len() > 1) {
380 type IntoIter = sumtype![];
381 } #(else) {
382 type IntoIter = #(#out_into_iter_ty)*;
383 }
384 fn param_into_iter(#self_val) -> Self::IntoIter
385 {
386 #{f(out_into_iter.as_slice())}
387 }
388 }
389 })
390 }
391 Self::Map => {
392 let map_fn: Ident = parse_quote!(__parametrized_map_fn);
393 let mapped_param: Ident = parse_quote!(__PARAMETRIZED_MAP_PARAM);
394 let out_map = tys_exprs
395 .iter()
396 .map(|item| {
397 item.iter()
398 .map(|(a, b)| {
399 Ok(generator::EmitContext {
400 kind: generator::EmitMap(map_fn.clone(), mapped_param.clone()),
401 krate: krate.clone(),
402 replacing_ty: replacing_ty.clone(),
403 }
404 .emit(a, &(b.clone(), a.clone()))?
405 .map(|a| a.0)
406 .unwrap_or(b.clone()))
407 })
408 .collect::<Result<Vec<_>>>()
409 })
410 .collect::<Result<Vec<_>>>()?;
411 let mapped = replace_type(
412 parse_quote!(#ident #ty_generics),
413 replacing_ty.clone(),
414 parse_quote!(#mapped_param),
415 );
416 let mapped_param_impl_generics = generics
417 .params
418 .iter()
419 .filter_map(|p| {
420 if let GenericParam::Type(p) = p {
421 let ident = &p.ident;
422 if &parse_quote!(#ident) as &Type == replacing_ty {
423 let mut p = p.clone();
424 p.ident = mapped_param.clone();
425 Some(quote!(#p))
426 } else {
427 None
428 }
429 } else {
430 None
431 }
432 })
433 .next()
434 .unwrap_or(quote!(#mapped_param));
435 Ok(quote! {
436 impl <
437 #(for p in &generics.params){ #p, }
438 #mapped_param_impl_generics
439 > #krate::ParametrizedMap<#param_index, #mapped_param> for #ident #ty_generics #where_clause {
440 type Mapped = #mapped;
441 fn param_map(#self_val, mut #map_fn: impl FnMut(Self::Item) -> #mapped_param) -> Self::Mapped
442 where
443 Self::Item: ::core::marker::Sized
444 {
445 #{emit_map_f(out_map.as_slice())}
446 }
447 }
448 })
449 }
450 }
451 }
452}
453
454impl Parse for TraitTarget {
455 fn parse(input: parse::ParseStream) -> Result<Self> {
456 let ident: Ident = input.fork().parse()?;
457 let r = match ident.to_string().as_str() {
458 "default" => Self::Default,
459 "iter_mut" => Self::IterMut,
460 "into_iter" => Self::IntoIter,
461 "map" => Self::Map,
462 _ => return Err(input.error("Require one of `iter`, `iter_mut`, `into_iter`, `map`")),
463 };
464 input.parse::<Ident>()?;
465 Ok(r)
466 }
467}
468
469#[derive(Debug, Default)]
470struct Arguments {
471 trait_impls: HashMap<usize, HashSet<TraitTarget>>,
472 krate: Option<Path>,
473}
474
475impl Parse for Arguments {
476 fn parse(input: parse::ParseStream) -> Result<Self> {
477 let mut ret: Self = Default::default();
478 while input.peek(Ident) {
479 if let Ok(tr) = input.parse::<TraitTarget>() {
480 let param_indices = if input.parse::<Token![=]>().is_ok() {
482 if let Ok(expr_arr) = input.parse::<ExprArray>() {
483 expr_arr
484 .elems
485 .into_iter()
486 .map(|expr| {
487 if let Expr::Lit(ExprLit {
488 lit: Lit::Int(lit), ..
489 }) = expr
490 {
491 lit.base10_parse::<usize>()
492 } else {
493 Err(Error::new(expr.span(), "Bad integer"))
494 }
495 })
496 .collect::<Result<Vec<_>>>()?
497 } else {
498 vec![input.parse::<LitInt>()?.base10_parse::<usize>()?]
499 }
500 } else {
501 vec![0]
502 };
503 for param_index in param_indices {
504 ret.trait_impls
505 .entry(param_index)
506 .or_default()
507 .insert(tr.clone());
508 }
509 } else {
510 let ident: Ident = input.parse()?;
511 if &ident == "krate" {
512 input.parse::<Token![=]>()?;
513 ret.krate = Some(input.parse()?);
514 } else {
515 return Err(Error::new(ident.span(), "Bad option"));
516 }
517 }
518 if input.parse::<Token![,]>().is_err() {
519 break;
520 }
521 }
522 if !input.is_empty() {
523 Err(syn::parse::Error::new(input.span(), "Unparsed args"))
524 } else {
525 Ok(ret)
526 }
527 }
528}
529
530trait ImplTarget {
531 fn emit_impl(&self, krate: &Path, tr: &TraitTarget, param_index: usize) -> Result<TokenStream>;
532}
533
534fn get_replacing_ty(generics: &Generics, param_index: usize) -> Type {
535 generics
536 .params
537 .get(param_index)
538 .and_then(|g| {
539 if let GenericParam::Type(TypeParam { ident, .. }) = g {
540 Some(parse_quote!(#ident))
541 } else {
542 None
543 }
544 })
545 .unwrap_or_else(|| {
546 abort!(
547 Span::call_site(),
548 format!("Cannot implement for index {}", param_index)
549 )
550 })
551}
552
553impl ImplTarget for ItemStruct {
554 fn emit_impl(&self, krate: &Path, tr: &TraitTarget, param_index: usize) -> Result<TokenStream> {
555 let self_val = Ident::new("self", Span::call_site());
556 let replacing_ty = get_replacing_ty(&self.generics, param_index);
557 let tys_exprs = self
558 .fields
559 .iter()
560 .enumerate()
561 .map(|(i, field)| {
562 let ty = field.ty.clone();
563 if let Some(ident) = &field.ident {
564 (ty, parse_quote! {#self_val.#ident})
565 } else {
566 let i = Index {
567 index: i as u32,
568 span: Span::call_site(),
569 };
570 (ty, parse_quote! {#self_val.#i})
571 }
572 })
573 .collect::<Vec<_>>();
574 tr.emit(
575 krate,
576 &self.ident,
577 &self.generics,
578 param_index,
579 &replacing_ty,
580 &self_val,
581 &[tys_exprs],
582 |inner| {
583 quote! { #(#inner)* }
584 },
585 |items| {
586 quote! {
587 #[allow(unused)]
588 #{&self.ident}
589 #(if let Fields::Named(_) = &self.fields) {
590 {#(for (inner, field) in items[0].iter().zip(&self.fields)) {
591 #{&field.ident} : #inner,
592 }}
593 } #(else) {
594 ( #(for inner in items[0].iter()), { #inner })
595 }
596 }
597 },
598 true,
599 )
600 }
601}
602
603impl ImplTarget for ItemEnum {
604 fn emit_impl(&self, krate: &Path, tr: &TraitTarget, param_index: usize) -> Result<TokenStream> {
605 let self_val: Ident = Ident::new("self", Span::call_site());
606 let replacing_ty = get_replacing_ty(&self.generics, param_index);
607 let variant_idents = self
608 .variants
609 .iter()
610 .map(|variant| {
611 variant
612 .fields
613 .iter()
614 .enumerate()
615 .map(|(i, field)| {
616 field.ident.clone().unwrap_or(Ident::new(
617 &format!("__parametric_type_id_{}", i),
618 Span::call_site(),
619 ))
620 })
621 .collect::<Vec<_>>()
622 })
623 .collect::<Vec<_>>();
624 let variant_items = self
625 .variants
626 .iter()
627 .zip(&variant_idents)
628 .map(|(var, idents)| {
629 var.fields
630 .iter()
631 .zip(idents)
632 .map(|(field, ident)| {
633 let ty = field.ty.clone();
634 (ty, parse_quote! {#ident})
635 })
636 .collect::<Vec<_>>()
637 })
638 .collect::<Vec<_>>();
639 tr.emit(
640 krate,
641 &self.ident,
642 &self.generics,
643 param_index,
644 &replacing_ty,
645 &self_val,
646 variant_items.as_slice(),
647 |inner| {
648 quote! {
649 #[allow(unused)]
650 match #self_val {
651 #(for ((variant, inner), idents) in self
652 .variants.iter().zip(inner).zip(&variant_idents)
653 ) {
654 #{&self.ident}::#{&variant.ident}
655 #(if let Fields::Named(_) = &variant.fields) {
656 { #(#idents),* }
657 }
658 #(if let Fields::Unnamed(_) = &variant.fields) {
659 ( #(#idents),* )
660 }
661 => { #inner }
662 }
663 }
664 }
665 },
666 |items| {
667 quote! {
668 match #self_val {
669 #(for ((variant, inner), idents) in self
670 .variants.iter().zip(items).zip(&variant_idents)
671 ) {
672 #{&self.ident}::#{&variant.ident}
673 #(if let Fields::Named(_) = &variant.fields) {
674 { #(#idents),* }
675 }
676 #(if let Fields::Unnamed(_) = &variant.fields) {
677 ( #(#idents),* )
678 }
679 => {
680 #{&self.ident}::#{&variant.ident}
681 #(if let Fields::Named(_) = &variant.fields) {
682 {
683 #(for (ident, inner) in idents.iter().zip(inner)) {
684 #ident: #inner,
685 }
686 }
687 }
688 #(if let Fields::Unnamed(_) = &variant.fields) {
689 ( #(#inner),* )
690 }
691 }
692 }
693 }
694 }
695 },
696 false,
697 )
698 }
699}
700
701fn inner_target<T: ImplTarget + ToTokens>(target: &T, arg: Arguments) -> TokenStream {
702 let krate = arg.krate.unwrap_or(parse_quote!(::parametrized));
703 let mut out = quote!(#target);
704 for (param_index, impl_traits) in &arg.trait_impls {
705 let impl_traits = TraitTarget::make_enough(impl_traits.clone());
706 for impl_trait in &impl_traits {
707 let ret = target
708 .emit_impl(&krate, impl_trait, *param_index)
709 .unwrap_or_else(|e| {
710 abort!(
711 e.span(),
712 format!(
713 "Cannot implement {:?} for parameter {}: {}",
714 &impl_trait, param_index, e
715 )
716 )
717 });
718 out.extend(ret);
719 }
720 }
721 out
722}
723
724fn inner(arg: Arguments, input: Item) -> TokenStream {
725 match input {
726 Item::Enum(item_enum) => inner_target(&item_enum, arg),
727 Item::Struct(item_struct) => inner_target(&item_struct, arg),
728 _ => abort!(input.span(), "Bad item"),
729 }
730}
731
732#[proc_macro_error]
733#[proc_macro_attribute]
734pub fn parametrized(attr: TokenStream1, input: TokenStream1) -> TokenStream1 {
735 inner(
736 parse(attr).unwrap_or_else(|e| abort!(e.span(), &format!("{}", e))),
737 parse_macro_input!(input as Item),
738 )
739 .into()
740}