1use std::collections::hash_map::DefaultHasher;
12use std::hash::{Hash, Hasher};
13
14use proc_macro::TokenStream;
15use proc_macro2::{Group, Punct, Spacing, Span, TokenStream as TokenStream2, TokenTree};
16use proc_macro_crate::FoundCrate;
17use quote::{format_ident, quote, ToTokens};
18use syn::parse::{Parse, ParseStream};
19use syn::punctuated::Punctuated;
20use syn::visit::Visit;
21use syn::visit_mut::{self, VisitMut};
22use syn::{
23 braced, parenthesized, parse_macro_input, Attribute, Block, Error, Expr, Fields, FnArg,
24 GenericArgument, GenericParam, Ident, ImplItem, Item, ItemEnum, ItemImpl, ItemTrait, Lifetime,
25 LitBool, Meta, Pat, PatIdent, Path, PathArguments, PathSegment, QSelf, Result, ReturnType,
26 Signature, Token, TraitItem, TraitItemFn, Type, TypeParamBound, TypePath, Variant, Visibility,
27};
28
29const FACADE: &str = "whichever";
32
33fn is_dispatch_attribute(attribute: &Attribute) -> bool {
36 attribute
37 .path()
38 .segments
39 .last()
40 .is_some_and(|segment| segment.ident == "dispatch")
41}
42
43#[proc_macro_attribute]
68pub fn dispatchable(attr: TokenStream, item: TokenStream) -> TokenStream {
69 if !attr.is_empty() {
70 return Error::new(Span::call_site(), "`#[dispatchable]` takes no arguments")
71 .to_compile_error()
72 .into();
73 }
74 let item_trait = parse_macro_input!(item as ItemTrait);
75 let site = proc_macro::Span::call_site();
76 expand_dispatchable(item_trait, site.line(), site.column())
77 .unwrap_or_else(|error| error.to_compile_error())
78 .into()
79}
80
81fn expand_dispatchable(item_trait: ItemTrait, line: usize, column: usize) -> Result<TokenStream2> {
82 let captured = item_trait.clone();
83 let mut emitted = item_trait;
84 for item in &mut emitted.items {
85 if let TraitItem::Fn(function) = item {
86 function
87 .attrs
88 .retain(|attribute| !is_dispatch_attribute(attribute));
89 }
90 }
91 let trait_ident = emitted.ident.clone();
92 let vis = emitted.vis.clone();
93 let mut hasher = DefaultHasher::new();
96 captured.to_token_stream().to_string().hash(&mut hasher);
97 let macro_name = format_ident!(
98 "__whichever_{}_{:016x}_{}_{}",
99 trait_ident,
100 hasher.finish(),
101 line,
102 column
103 );
104 let captured = rewrite_crate_paths(captured.to_token_stream());
107 Ok(quote! {
108 #emitted
109
110 #[doc(hidden)]
111 #[macro_export]
112 macro_rules! #macro_name {
113 (@v1 $facade:ident $($rest:tt)*) => {
114 ::$facade::__private::implement! { @v1 trait { #captured } $($rest)* }
115 };
116 }
117
118 #[doc(hidden)]
119 #[allow(unused_imports)]
120 #vis use #macro_name as #trait_ident;
121 })
122}
123
124fn rewrite_crate_paths(tokens: TokenStream2) -> TokenStream2 {
126 let mut output = TokenStream2::new();
127 let mut iter = tokens.into_iter().peekable();
128 while let Some(token) = iter.next() {
129 match token {
130 TokenTree::Group(group) => {
131 let mut rewritten =
132 Group::new(group.delimiter(), rewrite_crate_paths(group.stream()));
133 rewritten.set_span(group.span());
134 output.extend([TokenTree::Group(rewritten)]);
135 }
136 TokenTree::Ident(ident)
137 if ident == "crate"
138 && matches!(iter.peek(), Some(TokenTree::Punct(punct)) if punct.as_char() == ':') =>
139 {
140 output.extend([
141 TokenTree::Punct(Punct::new('$', Spacing::Alone)),
142 TokenTree::Ident(ident),
143 ]);
144 }
145 other => output.extend([other]),
146 }
147 }
148 output
149}
150
151struct TraitSpec {
159 path: Option<Path>,
160 assoc_overrides: Vec<(Ident, Type)>,
161 from: Option<bool>,
162 try_from: Option<bool>,
163 mismatch: Vec<(Ident, Expr)>,
164}
165
166impl TraitSpec {
167 fn assoc_override(&self, name: &Ident) -> Option<&Type> {
168 self.assoc_overrides
169 .iter()
170 .find(|(key, _)| key == name)
171 .map(|(_, ty)| ty)
172 }
173
174 fn mismatch_for(&self, method: &Ident) -> Option<&Expr> {
175 self.mismatch
176 .iter()
177 .find(|(key, _)| key == method)
178 .map(|(_, expr)| expr)
179 }
180}
181
182struct KeyedExpr {
183 key: Ident,
184 value: Expr,
185}
186
187impl Parse for KeyedExpr {
188 fn parse(input: ParseStream) -> Result<Self> {
189 let key = input.parse()?;
190 input.parse::<Token![=]>()?;
191 let value = input.parse()?;
192 Ok(Self { key, value })
193 }
194}
195
196impl Parse for TraitSpec {
197 fn parse(input: ParseStream) -> Result<Self> {
198 let starts_with_option =
199 input.peek(Ident) && (input.peek2(Token![=]) || input.peek2(syn::token::Paren));
200 let path = if input.is_empty() || starts_with_option {
201 None
202 } else {
203 Some(input.parse::<Path>()?)
204 };
205 let mut spec = TraitSpec {
206 path,
207 assoc_overrides: vec![],
208 from: None,
209 try_from: None,
210 mismatch: vec![],
211 };
212 let mut expect_comma = spec.path.is_some();
213 while !input.is_empty() {
214 if expect_comma {
215 input.parse::<Token![,]>()?;
216 if input.is_empty() {
217 break;
218 }
219 }
220 expect_comma = true;
221 let key: Ident = input.parse()?;
222 if key == "mismatch" {
223 let content;
224 parenthesized!(content in input);
225 let pairs = Punctuated::<KeyedExpr, Token![,]>::parse_terminated(&content)?;
226 spec.mismatch
227 .extend(pairs.into_iter().map(|pair| (pair.key, pair.value)));
228 continue;
229 }
230 if !input.peek(Token![=]) {
231 return Err(Error::new(
232 key.span(),
233 format!("unknown option `{key}`; the options are `Name = Type` for an associated type, `mismatch(method = expr)`, `from = false` and `try_from = true`"),
234 ));
235 }
236 input.parse::<Token![=]>()?;
237 match key.to_string().as_str() {
238 "from" => spec.from = Some(input.parse::<LitBool>()?.value),
239 "try_from" => spec.try_from = Some(input.parse::<LitBool>()?.value),
240 _ => spec.assoc_overrides.push((key, input.parse()?)),
241 }
242 }
243 Ok(spec)
244 }
245}
246
247#[proc_macro_attribute]
308pub fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream {
309 let site = proc_macro::Span::call_site();
310 let (line, column) = (site.line(), site.column());
311 let item = parse_macro_input!(item as Item);
312 let expanded = match item {
313 Item::Enum(item_enum) => expand_dispatch(attr.into(), item_enum, line, column),
314 Item::Impl(item_impl) => expand_inherent_invocation(attr.into(), item_impl),
315 other => Err(Error::new_spanned(
316 other,
317 "`#[dispatch]` goes on an enum, or on an inherent impl block of a dispatch enum",
318 )),
319 };
320 expanded
321 .unwrap_or_else(|error| error.to_compile_error())
322 .into()
323}
324
325fn expand_dispatch(
326 first: TokenStream2,
327 mut item_enum: ItemEnum,
328 line: usize,
329 column: usize,
330) -> Result<TokenStream2> {
331 let mut specs = Vec::new();
332 if !first.is_empty() {
333 specs.push(first);
334 }
335 let mut kept = vec![];
336 for attribute in std::mem::take(&mut item_enum.attrs) {
337 if !is_dispatch_attribute(&attribute) {
338 kept.push(attribute);
339 continue;
340 }
341 match &attribute.meta {
342 Meta::List(list) => specs.push(list.tokens.clone()),
343 Meta::Path(_) => {}
345 Meta::NameValue(_) => {
346 return Err(Error::new_spanned(
347 attribute,
348 "expected `#[dispatch(Trait, ...)]`",
349 ));
350 }
351 }
352 }
353 item_enum.attrs = kept;
354 let variants = item_enum
355 .variants
356 .iter()
357 .map(variant_info)
358 .collect::<Result<Vec<_>>>()?;
359 let captured = item_enum.clone();
360 let mut emitted = item_enum;
361 for variant in &mut emitted.variants {
362 variant
363 .attrs
364 .retain(|attribute| !is_dispatch_attribute(attribute));
365 }
366 let facade = facade_ident()?;
367 let parsed_specs = specs
368 .iter()
369 .map(|tokens| syn::parse2::<TraitSpec>(tokens.clone()))
370 .collect::<Result<Vec<_>>>()?;
371 let generate_from = !parsed_specs.iter().any(|spec| spec.from == Some(false));
372 let generate_try_from = parsed_specs.iter().any(|spec| spec.try_from == Some(true));
373 let conversions = conversion_impls(&emitted, &variants, generate_from, generate_try_from);
374 let mut invocations = TokenStream2::new();
375 for (spec, tokens) in parsed_specs.iter().zip(&specs) {
376 let Some(path) = &spec.path else { continue };
377 let mut macro_path = path.clone();
378 if let Some(last) = macro_path.segments.last_mut() {
379 last.arguments = PathArguments::None;
380 }
381 invocations.extend(quote! {
382 #macro_path! { @v1 #facade enum { #captured } spec { #tokens } }
383 });
384 }
385 let enum_ident = &emitted.ident;
389 let mut hasher = DefaultHasher::new();
390 emitted.to_token_stream().to_string().hash(&mut hasher);
391 let enum_macro = format_ident!(
392 "__whichever_enum_{}_{:016x}_{}_{}",
393 enum_ident,
394 hasher.finish(),
395 line,
396 column
397 );
398 Ok(quote! {
399 #emitted
400 #conversions
401
402 #[allow(unused_macros)]
403 macro_rules! #enum_macro {
404 (@v1 $facade:ident $($rest:tt)*) => {
405 ::$facade::__private::implement_inherent! { @v1 enum { #emitted } $($rest)* }
406 };
407 }
408
409 #[doc(hidden)]
410 #[allow(unused_imports)]
411 pub(crate) use #enum_macro as #enum_ident;
412
413 #invocations
414 })
415}
416
417fn expand_inherent_invocation(attr: TokenStream2, item_impl: ItemImpl) -> Result<TokenStream2> {
420 if !attr.is_empty() {
421 return Err(Error::new_spanned(
422 attr,
423 "`#[dispatch]` on an impl block takes no arguments; the variants are named on each method",
424 ));
425 }
426 if let Some((_, trait_path, _)) = &item_impl.trait_ {
427 return Err(Error::new_spanned(
428 trait_path,
429 "`#[dispatch]` forwards inherent methods; a trait is dispatched with `#[dispatch(Trait)]` on the enum",
430 ));
431 }
432 let Type::Path(self_type) = &*item_impl.self_ty else {
433 return Err(Error::new_spanned(
434 &item_impl.self_ty,
435 "the impl must name a dispatch enum by path",
436 ));
437 };
438 let mut macro_path = self_type.path.clone();
439 if let Some(last) = macro_path.segments.last_mut() {
440 last.arguments = PathArguments::None;
441 }
442 let facade = facade_ident()?;
443 Ok(quote! {
444 #macro_path! { @v1 #facade impl { #item_impl } }
445 })
446}
447
448fn facade_ident() -> Result<Ident> {
449 match proc_macro_crate::crate_name(FACADE) {
450 Ok(FoundCrate::Itself) => Ok(format_ident!("{FACADE}")),
451 Ok(FoundCrate::Name(name)) => Ok(Ident::new(&name, Span::call_site())),
452 Err(error) => Err(Error::new(
453 Span::call_site(),
454 format!("`{FACADE}` must be a dependency of this crate: {error}"),
455 )),
456 }
457}
458
459struct VariantInfo {
460 ident: Ident,
461 payload: Type,
462 cfgs: Vec<Attribute>,
463 deref: bool,
465}
466
467fn variant_info(variant: &Variant) -> Result<VariantInfo> {
468 let Fields::Unnamed(fields) = &variant.fields else {
469 return Err(Error::new_spanned(
470 variant,
471 "a dispatch variant must be `Name(Payload)`",
472 ));
473 };
474 if fields.unnamed.len() != 1 {
475 return Err(Error::new_spanned(
476 variant,
477 "a dispatch variant must hold exactly one payload",
478 ));
479 }
480 let mut deref = false;
481 for attribute in variant
482 .attrs
483 .iter()
484 .filter(|attribute| is_dispatch_attribute(attribute))
485 {
486 let option: Ident = attribute.parse_args()?;
487 if option != "deref" {
488 return Err(Error::new_spanned(
489 &option,
490 format!("unknown variant option `{option}`; the only one is `deref`"),
491 ));
492 }
493 deref = true;
494 }
495 Ok(VariantInfo {
496 ident: variant.ident.clone(),
497 payload: fields.unnamed[0].ty.clone(),
498 cfgs: cfg_attributes(&variant.attrs),
499 deref,
500 })
501}
502
503fn target_type(variant: &VariantInfo) -> TokenStream2 {
506 let payload = &variant.payload;
507 if variant.deref {
508 quote!(<#payload as ::core::ops::Deref>::Target)
509 } else {
510 quote!(#payload)
511 }
512}
513
514fn cfg_attributes(attributes: &[Attribute]) -> Vec<Attribute> {
515 attributes
516 .iter()
517 .filter(|attribute| attribute.path().is_ident("cfg"))
518 .cloned()
519 .collect()
520}
521
522fn conversion_impls(
524 item_enum: &ItemEnum,
525 variants: &[VariantInfo],
526 from: bool,
527 try_from: bool,
528) -> TokenStream2 {
529 let (impl_generics, ty_generics, where_clause) = item_enum.generics.split_for_impl();
530 let enum_ident = &item_enum.ident;
531 let mut payload_counts = std::collections::HashMap::<String, usize>::new();
534 for variant in variants {
535 *payload_counts
536 .entry(variant.payload.to_token_stream().to_string())
537 .or_default() += 1;
538 }
539 let value = Ident::new("__value", Span::mixed_site());
540 let other = Ident::new("__other", Span::mixed_site());
541 let mut output = TokenStream2::new();
542 for variant in variants {
543 let VariantInfo {
544 ident,
545 payload,
546 cfgs,
547 ..
548 } = variant;
549 if payload_counts[&payload.to_token_stream().to_string()] > 1 {
550 continue;
551 }
552 if from {
553 output.extend(quote! {
554 #(#cfgs)*
555 #[automatically_derived]
556 impl #impl_generics ::core::convert::From<#payload> for #enum_ident #ty_generics #where_clause {
557 #[inline]
558 fn from(#value: #payload) -> Self {
559 Self::#ident(#value)
560 }
561 }
562 });
563 }
564 if try_from {
565 output.extend(quote! {
566 #(#cfgs)*
567 #[automatically_derived]
568 impl #impl_generics ::core::convert::TryFrom<#enum_ident #ty_generics> for #payload #where_clause {
569 type Error = #enum_ident #ty_generics;
570 #[inline]
571 #[allow(unreachable_patterns)]
572 fn try_from(#value: #enum_ident #ty_generics) -> ::core::result::Result<Self, Self::Error> {
573 match #value {
574 #enum_ident::#ident(#value) => ::core::result::Result::Ok(#value),
575 #other => ::core::result::Result::Err(#other),
576 }
577 }
578 }
579 });
580 }
581 }
582 output
583}
584
585struct ImplementInput {
590 item_trait: ItemTrait,
591 item_enum: ItemEnum,
592 spec: TraitSpec,
593}
594
595impl Parse for ImplementInput {
596 fn parse(input: ParseStream) -> Result<Self> {
597 input.parse::<Token![@]>()?;
598 let version: Ident = input.parse()?;
599 if version != "v1" {
600 return Err(Error::new_spanned(
601 version,
602 "unsupported `whichever` protocol version",
603 ));
604 }
605 input.parse::<Token![trait]>()?;
606 let trait_tokens;
607 braced!(trait_tokens in input);
608 let item_trait: ItemTrait = trait_tokens.parse()?;
609 input.parse::<Token![enum]>()?;
610 let enum_tokens;
611 braced!(enum_tokens in input);
612 let item_enum: ItemEnum = enum_tokens.parse()?;
613 let tag: Ident = input.parse()?;
614 if tag != "spec" {
615 return Err(Error::new_spanned(tag, "expected `spec`"));
616 }
617 let spec_tokens;
618 braced!(spec_tokens in input);
619 let spec: TraitSpec = spec_tokens.parse()?;
620 Ok(Self {
621 item_trait,
622 item_enum,
623 spec,
624 })
625 }
626}
627
628#[doc(hidden)]
630#[proc_macro]
631pub fn implement(input: TokenStream) -> TokenStream {
632 let input = parse_macro_input!(input as ImplementInput);
633 expand_implement(input)
634 .unwrap_or_else(|error| error.to_compile_error())
635 .into()
636}
637
638struct Substitution {
640 types: Vec<(Ident, Type, Option<Path>)>,
641 lifetimes: Vec<(Lifetime, Lifetime)>,
642 consts: Vec<(Ident, Expr)>,
643}
644
645impl Substitution {
646 fn new(item_trait: &ItemTrait, trait_path: &Path) -> Result<Self> {
647 let params: Vec<&GenericParam> = item_trait.generics.params.iter().collect();
648 let last = trait_path
649 .segments
650 .last()
651 .expect("a path has at least one segment");
652 let args: Vec<&GenericArgument> = match &last.arguments {
653 PathArguments::AngleBracketed(arguments) => arguments
654 .args
655 .iter()
656 .filter(|argument| {
657 !matches!(
658 argument,
659 GenericArgument::AssocType(_)
660 | GenericArgument::AssocConst(_)
661 | GenericArgument::Constraint(_)
662 )
663 })
664 .collect(),
665 _ => vec![],
666 };
667 if args.len() != params.len() {
668 return Err(Error::new_spanned(
669 trait_path,
670 format!(
671 "`{}` takes {} generic argument(s); name them, e.g. `Trait<Arg>`",
672 item_trait.ident,
673 params.len()
674 ),
675 ));
676 }
677 let mut substitution = Self {
678 types: vec![],
679 lifetimes: vec![],
680 consts: vec![],
681 };
682 for (param, argument) in params.into_iter().zip(args) {
683 match (param, argument) {
684 (GenericParam::Type(param), GenericArgument::Type(ty)) => {
685 let single_bound = match param.bounds.iter().collect::<Vec<_>>().as_slice() {
686 [TypeParamBound::Trait(bound)] => Some(bound.path.clone()),
687 _ => None,
688 };
689 substitution
690 .types
691 .push((param.ident.clone(), ty.clone(), single_bound));
692 }
693 (GenericParam::Lifetime(param), GenericArgument::Lifetime(lifetime)) => {
694 substitution
695 .lifetimes
696 .push((param.lifetime.clone(), lifetime.clone()));
697 }
698 (GenericParam::Const(param), GenericArgument::Const(expr)) => {
699 substitution
700 .consts
701 .push((param.ident.clone(), expr.clone()));
702 }
703 (GenericParam::Const(param), GenericArgument::Type(Type::Path(path)))
705 if path.qself.is_none() =>
706 {
707 let expr = Expr::Path(syn::ExprPath {
708 attrs: vec![],
709 qself: None,
710 path: path.path.clone(),
711 });
712 substitution.consts.push((param.ident.clone(), expr));
713 }
714 _ => {
715 return Err(Error::new_spanned(
716 argument,
717 "generic argument does not match the trait's parameter",
718 ))
719 }
720 }
721 }
722 Ok(substitution)
723 }
724
725 fn const_for(&self, ident: &Ident) -> Option<&Expr> {
726 self.consts
727 .iter()
728 .find(|(name, _)| name == ident)
729 .map(|(_, expr)| expr)
730 }
731}
732
733fn bare_ident(ty: &Type) -> Option<&Ident> {
735 match ty {
736 Type::Path(TypePath { qself: None, path }) if path.segments.len() == 1 => {
737 Some(&path.segments[0].ident)
738 }
739 _ => None,
740 }
741}
742
743impl VisitMut for Substitution {
744 fn visit_type_mut(&mut self, ty: &mut Type) {
745 if let Type::Path(type_path) = ty {
746 if type_path.qself.is_none() {
747 let first = &type_path.path.segments[0];
748 let replacement = self
749 .types
750 .iter()
751 .find(|(ident, _, _)| *ident == first.ident);
752 if let Some((_, replacement, bound)) = replacement {
753 if type_path.path.segments.len() == 1 {
754 *ty = replacement.clone();
755 return;
756 }
757 if let Some(bound) = bound {
760 let rest: Vec<PathSegment> =
761 type_path.path.segments.iter().skip(1).cloned().collect();
762 let mut path = bound.clone();
763 let position = path.segments.len();
764 path.segments.extend(rest);
765 *ty = Type::Path(TypePath {
766 qself: Some(QSelf {
767 lt_token: Default::default(),
768 ty: Box::new(replacement.clone()),
769 position,
770 as_token: Some(Default::default()),
771 gt_token: Default::default(),
772 }),
773 path,
774 });
775 return;
776 }
777 }
778 }
779 }
780 visit_mut::visit_type_mut(self, ty);
781 }
782
783 fn visit_lifetime_mut(&mut self, lifetime: &mut Lifetime) {
784 if let Some((_, replacement)) = self.lifetimes.iter().find(|(from, _)| from == lifetime) {
785 *lifetime = replacement.clone();
786 }
787 }
788
789 fn visit_expr_mut(&mut self, expr: &mut Expr) {
790 if let Expr::Path(path) = expr {
791 if path.qself.is_none() && path.path.segments.len() == 1 {
792 if let Some(replacement) = self.const_for(&path.path.segments[0].ident) {
793 *expr = replacement.clone();
794 return;
795 }
796 }
797 }
798 visit_mut::visit_expr_mut(self, expr);
799 }
800
801 fn visit_generic_argument_mut(&mut self, argument: &mut GenericArgument) {
804 if let GenericArgument::Type(ty) = argument {
805 if let Some(replacement) = bare_ident(ty).and_then(|ident| self.const_for(ident)) {
806 *argument = GenericArgument::Const(replacement.clone());
807 return;
808 }
809 }
810 visit_mut::visit_generic_argument_mut(self, argument);
811 }
812}
813
814fn expand_implement(input: ImplementInput) -> Result<TokenStream2> {
815 let ImplementInput {
816 item_trait,
817 item_enum,
818 spec,
819 } = input;
820 let enum_ident = &item_enum.ident;
821 let (impl_generics, ty_generics, where_clause) = item_enum.generics.split_for_impl();
822 let Some(trait_path) = &spec.path else {
823 return Err(Error::new(
824 Span::call_site(),
825 "the trait to implement is missing",
826 ));
827 };
828 let variants = item_enum
829 .variants
830 .iter()
831 .map(variant_info)
832 .collect::<Result<Vec<_>>>()?;
833 let Some(first) = variants.first() else {
834 return Err(Error::new_spanned(
835 &item_enum,
836 "a dispatch enum needs at least one variant",
837 ));
838 };
839 let first_target = target_type(first);
840 let first_payload = &first_target;
841 let mut substitution = Substitution::new(&item_trait, trait_path)?;
842 let mut items = TokenStream2::new();
843 let mut agreed_types: Vec<&Ident> = vec![];
844 for item in &item_trait.items {
845 match item {
846 TraitItem::Type(assoc) => {
847 if !assoc.generics.params.is_empty() {
848 return Err(Error::new_spanned(
849 assoc,
850 "generic associated types are not supported",
851 ));
852 }
853 let name = &assoc.ident;
854 let ty = match spec.assoc_override(name) {
855 Some(ty) => quote!(#ty),
856 None => {
857 agreed_types.push(name);
858 quote!(<#first_payload as #trait_path>::#name)
859 }
860 };
861 items.extend(quote! { type #name = #ty; });
862 }
863 TraitItem::Const(constant) => {
864 let name = &constant.ident;
865 let mut ty = constant.ty.clone();
866 substitution.visit_type_mut(&mut ty);
867 items.extend(quote! {
868 const #name: #ty = <#first_payload as #trait_path>::#name;
869 });
870 }
871 TraitItem::Fn(function) => {
872 items.extend(expand_method(
873 function,
874 &variants,
875 enum_ident,
876 trait_path,
877 &spec,
878 &mut substitution,
879 )?);
880 }
881 TraitItem::Macro(item) => {
882 return Err(Error::new_spanned(
883 item,
884 "macro invocations inside the trait are not supported",
885 ));
886 }
887 _ => {}
888 }
889 }
890 let agreement_check = associated_type_check(
891 &agreed_types,
892 &variants,
893 trait_path,
894 first_payload,
895 &item_enum.generics,
896 );
897 Ok(quote! {
898 #[automatically_derived]
899 impl #impl_generics #trait_path for #enum_ident #ty_generics #where_clause {
900 #items
901 }
902 #agreement_check
903 })
904}
905
906fn associated_type_check(
910 names: &[&Ident],
911 variants: &[VariantInfo],
912 trait_path: &Path,
913 first_payload: &TokenStream2,
914 generics: &syn::Generics,
915) -> TokenStream2 {
916 if names.is_empty() || variants.len() < 2 {
917 return TokenStream2::new();
918 }
919 let bindings: Vec<GenericArgument> = names
920 .iter()
921 .map(|name| syn::parse_quote!(#name = <#first_payload as #trait_path>::#name))
922 .collect();
923 let mut bound = trait_path.clone();
924 let last = bound
925 .segments
926 .last_mut()
927 .expect("a path has at least one segment");
928 match &mut last.arguments {
929 PathArguments::AngleBracketed(arguments) => arguments.args.extend(bindings),
930 arguments => {
931 *arguments = PathArguments::AngleBracketed(syn::parse_quote!(<#(#bindings),*>))
932 }
933 }
934 let (impl_generics, _, where_clause) = generics.split_for_impl();
935 let checks = variants.iter().skip(1).map(|variant| {
936 let payload = target_type(variant);
937 let cfgs = &variant.cfgs;
938 quote! { #(#cfgs)* __same_as_first_variant::<#payload>(); }
939 });
940 quote! {
941 const _: () = {
942 #[allow(dead_code, non_snake_case)]
943 fn __whichever_associated_types_agree #impl_generics () #where_clause {
944 fn __same_as_first_variant<__Variant: ?Sized + #bound>() {}
945 #(#checks)*
946 }
947 };
948 }
949}
950
951enum ReturnShape {
952 Plain,
953 RewrapSelf,
954 MapSelf,
955 IntoOverridden,
956 MapErrOverridden,
957}
958
959fn single_segment(ty: &Type) -> Option<&PathSegment> {
960 match ty {
961 Type::Path(TypePath { qself: None, path }) if path.segments.len() == 1 => {
962 Some(&path.segments[0])
963 }
964 _ => None,
965 }
966}
967
968fn is_self(ty: &Type) -> bool {
969 single_segment(ty).is_some_and(|segment| segment.ident == "Self" && segment.arguments.is_none())
970}
971
972fn self_assoc(ty: &Type) -> Option<&Ident> {
974 match ty {
975 Type::Path(TypePath { qself: None, path })
976 if path.segments.len() == 2 && path.segments[0].ident == "Self" =>
977 {
978 Some(&path.segments[1].ident)
979 }
980 _ => None,
981 }
982}
983
984fn generic_type_args(segment: &PathSegment) -> Vec<&Type> {
985 match &segment.arguments {
986 PathArguments::AngleBracketed(arguments) => arguments
987 .args
988 .iter()
989 .filter_map(|argument| match argument {
990 GenericArgument::Type(ty) => Some(ty),
991 _ => None,
992 })
993 .collect(),
994 _ => vec![],
995 }
996}
997
998fn classify_return(output: &ReturnType, spec: &TraitSpec) -> Result<ReturnShape> {
999 let ReturnType::Type(_, ty) = output else {
1000 return Ok(ReturnShape::Plain);
1001 };
1002 if let Type::ImplTrait(_) = **ty {
1003 return Err(Error::new_spanned(
1004 ty,
1005 "a return-position `impl Trait` cannot be dispatched by a `match` because every variant returns a different opaque type; return a concrete type or a boxed trait object",
1006 ));
1007 }
1008 if is_self(ty) {
1009 return Ok(ReturnShape::RewrapSelf);
1010 }
1011 if let Some(name) = self_assoc(ty) {
1012 return Ok(if spec.assoc_override(name).is_some() {
1013 ReturnShape::IntoOverridden
1014 } else {
1015 ReturnShape::Plain
1016 });
1017 }
1018 if let Some(segment) = single_segment(ty) {
1019 let args = generic_type_args(segment);
1020 let wraps_self = args.first().is_some_and(|first| is_self(first));
1021 if (segment.ident == "Option" || segment.ident == "Result") && wraps_self {
1022 return Ok(ReturnShape::MapSelf);
1023 }
1024 if segment.ident == "Result" && args.len() == 2 {
1025 if let Some(name) = self_assoc(args[1]) {
1026 if spec.assoc_override(name).is_some() {
1027 return Ok(ReturnShape::MapErrOverridden);
1028 }
1029 }
1030 }
1031 }
1032 Ok(ReturnShape::Plain)
1033}
1034
1035struct SelfFinder(bool);
1036
1037impl<'ast> Visit<'ast> for SelfFinder {
1038 fn visit_type_path(&mut self, type_path: &'ast TypePath) {
1039 if type_path.qself.is_none()
1040 && type_path.path.segments.len() == 1
1041 && type_path.path.segments[0].ident == "Self"
1042 {
1043 self.0 = true;
1044 }
1045 syn::visit::visit_type_path(self, type_path);
1046 }
1047}
1048
1049#[derive(Clone, Copy)]
1051enum Access {
1052 Shared,
1053 Exclusive,
1054 Owned,
1055}
1056
1057enum ArgumentKind {
1058 Plain,
1059 SameVariant(Access),
1060}
1061
1062fn argument_kind(ty: &Type) -> Result<ArgumentKind> {
1063 let (inner, access) = match ty {
1064 Type::Reference(reference) if reference.mutability.is_some() => {
1065 (&*reference.elem, Access::Exclusive)
1066 }
1067 Type::Reference(reference) => (&*reference.elem, Access::Shared),
1068 other => (other, Access::Owned),
1069 };
1070 if is_self(inner) {
1071 return Ok(ArgumentKind::SameVariant(access));
1072 }
1073 let mut finder = SelfFinder(false);
1074 finder.visit_type(ty);
1075 if finder.0 {
1076 return Err(Error::new_spanned(
1077 ty,
1078 "only `Self`, `&Self` and `&mut Self` arguments can be dispatched; this type mentions `Self` elsewhere",
1079 ));
1080 }
1081 Ok(ArgumentKind::Plain)
1082}
1083
1084fn payload_expression(deref: bool, access: Access, binding: &Ident) -> TokenStream2 {
1088 if !deref {
1089 return quote!(#binding);
1090 }
1091 match access {
1092 Access::Shared => quote!(&**#binding),
1093 Access::Exclusive => quote!(&mut **#binding),
1094 Access::Owned => quote!(*#binding),
1095 }
1096}
1097
1098fn method_mismatch(function: &TraitItemFn) -> Result<Option<Expr>> {
1100 for attribute in &function.attrs {
1101 if !is_dispatch_attribute(attribute) {
1102 continue;
1103 }
1104 let pair: KeyedExpr = attribute.parse_args()?;
1105 if pair.key != "mismatch" {
1106 return Err(Error::new_spanned(
1107 pair.key,
1108 "expected `#[dispatch(mismatch = <expr>)]`",
1109 ));
1110 }
1111 return Ok(Some(pair.value));
1112 }
1113 Ok(None)
1114}
1115
1116fn expand_method(
1117 function: &TraitItemFn,
1118 variants: &[VariantInfo],
1119 enum_ident: &Ident,
1120 trait_path: &Path,
1121 spec: &TraitSpec,
1122 substitution: &mut Substitution,
1123) -> Result<TokenStream2> {
1124 let mut signature: Signature = function.sig.clone();
1125 substitution.visit_signature_mut(&mut signature);
1126 let method = signature.ident.clone();
1127 let Some(FnArg::Receiver(receiver)) = signature.inputs.first() else {
1128 if function.default.is_some() {
1129 return Ok(TokenStream2::new());
1130 }
1131 return Err(Error::new_spanned(
1132 &function.sig,
1133 format!("`{method}` has no receiver, so there is no variant to dispatch on; give it a default body or implement it by hand"),
1134 ));
1135 };
1136 if receiver.colon_token.is_some() {
1137 return Err(Error::new_spanned(
1138 receiver,
1139 "explicit `self: Type` receivers are not supported",
1140 ));
1141 }
1142 let receiver_access = match (&receiver.reference, &receiver.mutability) {
1143 (None, _) => Access::Owned,
1144 (Some(_), None) => Access::Shared,
1145 (Some(_), Some(_)) => Access::Exclusive,
1146 };
1147 let self_token = receiver.self_token;
1150 let hygienic = |name: &str| Ident::new(name, Span::mixed_site());
1153 let inner = hygienic("__inner");
1154 let mut argument_idents = vec![];
1155 let mut same_variant: Vec<(usize, Access)> = vec![];
1156 for (index, input) in signature.inputs.iter_mut().skip(1).enumerate() {
1157 let FnArg::Typed(pat_type) = input else {
1158 unreachable!("receiver handled above")
1159 };
1160 let ident = hygienic(&format!("__argument{index}"));
1161 *pat_type.pat = Pat::Ident(PatIdent {
1162 attrs: vec![],
1163 by_ref: None,
1164 mutability: None,
1165 ident: ident.clone(),
1166 subpat: None,
1167 });
1168 if let ArgumentKind::SameVariant(access) = argument_kind(&pat_type.ty)? {
1169 same_variant.push((index, access));
1170 }
1171 argument_idents.push(ident);
1172 }
1173 let return_shape = classify_return(&signature.output, spec)?;
1174 let is_async = signature.asyncness.is_some();
1175 let is_unsafe = signature.unsafety.is_some();
1176 let mismatch = match spec.mismatch_for(&method) {
1177 Some(expr) => Some(expr.clone()),
1178 None => method_mismatch(function)?,
1179 };
1180 let mut arms = TokenStream2::new();
1181 for variant in variants {
1182 let VariantInfo {
1183 ident: variant_ident,
1184 payload,
1185 cfgs,
1186 deref,
1187 } = variant;
1188 if *deref && matches!(return_shape, ReturnShape::RewrapSelf | ReturnShape::MapSelf) {
1189 return Err(Error::new_spanned(
1190 payload,
1191 format!("`{method}` returns `Self`, which cannot be wrapped back into the dereferenced payload of variant `{variant_ident}`; drop `deref` on it or implement the method by hand"),
1192 ));
1193 }
1194 let target = target_type(variant);
1195 let receiver_expression = payload_expression(*deref, receiver_access, &inner);
1196 let call_arguments = argument_idents.iter().enumerate().map(|(index, ident)| {
1197 match same_variant
1198 .iter()
1199 .find(|(same_index, _)| *same_index == index)
1200 {
1201 Some((_, access)) => {
1202 payload_expression(*deref, *access, &hygienic(&format!("__same{index}")))
1203 }
1204 None => quote!(#ident),
1205 }
1206 });
1207 let mut call = quote! {
1208 <#target as #trait_path>::#method(#receiver_expression #(, #call_arguments)*)
1209 };
1210 if is_unsafe {
1211 call = quote! { unsafe { #call } };
1212 }
1213 if is_async {
1214 call = quote! { #call.await };
1215 }
1216 call = match return_shape {
1217 ReturnShape::Plain => call,
1218 ReturnShape::RewrapSelf => quote! { #enum_ident::#variant_ident(#call) },
1219 ReturnShape::MapSelf => quote! { #call.map(#enum_ident::#variant_ident) },
1220 ReturnShape::IntoOverridden => quote! { ::core::convert::Into::into(#call) },
1221 ReturnShape::MapErrOverridden => {
1222 quote! { #call.map_err(::core::convert::Into::into) }
1223 }
1224 };
1225 let pattern = if same_variant.is_empty() {
1226 quote! { #enum_ident::#variant_ident(#inner) }
1227 } else {
1228 let same_patterns = same_variant.iter().map(|(index, _)| {
1229 let same = hygienic(&format!("__same{index}"));
1230 quote! { #enum_ident::#variant_ident(#same) }
1231 });
1232 quote! { (#enum_ident::#variant_ident(#inner), #(#same_patterns),*) }
1233 };
1234 arms.extend(quote! { #(#cfgs)* #pattern => #call, });
1235 }
1236 let scrutinee = if same_variant.is_empty() {
1237 quote!(#self_token)
1238 } else {
1239 let arguments = same_variant
1240 .iter()
1241 .map(|(index, _)| &argument_idents[*index]);
1242 quote! { (#self_token, #(#arguments),*) }
1243 };
1244 let (fallback, allow_unreachable) = if same_variant.is_empty() {
1245 (TokenStream2::new(), TokenStream2::new())
1246 } else {
1247 let Some(expr) = mismatch else {
1248 return Err(Error::new_spanned(
1249 &function.sig,
1250 format!("`{method}` takes `Self` arguments; give the result for two different variants with `#[dispatch(mismatch = <expr>)]` on the trait method or `mismatch({method} = <expr>)` on the enum"),
1251 ));
1252 };
1253 (
1254 quote! { _ => #expr, },
1255 quote! { #[allow(unreachable_patterns)] },
1256 )
1257 };
1258 let inline = if function
1259 .attrs
1260 .iter()
1261 .any(|attribute| attribute.path().is_ident("inline"))
1262 {
1263 TokenStream2::new()
1264 } else {
1265 quote! { #[inline] }
1266 };
1267 let cfgs = cfg_attributes(&function.attrs);
1268 Ok(quote! {
1269 #(#cfgs)*
1270 #inline
1271 #allow_unreachable
1272 #signature {
1273 match #scrutinee {
1274 #arms
1275 #fallback
1276 }
1277 }
1278 })
1279}
1280
1281struct InherentInput {
1286 item_enum: ItemEnum,
1287 item_impl: ItemImpl,
1288}
1289
1290impl Parse for InherentInput {
1291 fn parse(input: ParseStream) -> Result<Self> {
1292 input.parse::<Token![@]>()?;
1293 let version: Ident = input.parse()?;
1294 if version != "v1" {
1295 return Err(Error::new_spanned(
1296 version,
1297 "unsupported `whichever` protocol version",
1298 ));
1299 }
1300 input.parse::<Token![enum]>()?;
1301 let enum_tokens;
1302 braced!(enum_tokens in input);
1303 let item_enum: ItemEnum = enum_tokens.parse()?;
1304 input.parse::<Token![impl]>()?;
1305 let impl_tokens;
1306 braced!(impl_tokens in input);
1307 let item_impl: ItemImpl = impl_tokens.parse()?;
1308 Ok(Self {
1309 item_enum,
1310 item_impl,
1311 })
1312 }
1313}
1314
1315#[doc(hidden)]
1318#[proc_macro]
1319pub fn implement_inherent(input: TokenStream) -> TokenStream {
1320 let input = parse_macro_input!(input as InherentInput);
1321 expand_inherent(input)
1322 .unwrap_or_else(|error| error.to_compile_error())
1323 .into()
1324}
1325
1326struct Forwarding {
1328 all: bool,
1329 variants: Vec<Ident>,
1330 call: Option<Ident>,
1331 with: Option<Path>,
1332}
1333
1334impl Parse for Forwarding {
1335 fn parse(input: ParseStream) -> Result<Self> {
1336 let mut forwarding = Forwarding {
1337 all: false,
1338 variants: vec![],
1339 call: None,
1340 with: None,
1341 };
1342 while !input.is_empty() {
1343 if input.peek(Token![*]) {
1344 input.parse::<Token![*]>()?;
1345 forwarding.all = true;
1346 } else {
1347 let ident: Ident = input.parse()?;
1348 if input.peek(Token![=]) {
1349 input.parse::<Token![=]>()?;
1350 match ident.to_string().as_str() {
1351 "call" => forwarding.call = Some(input.parse()?),
1352 "with" => forwarding.with = Some(input.parse()?),
1353 other => {
1354 return Err(Error::new(
1355 ident.span(),
1356 format!("unknown option `{other}`; the options are `call = method` and `with = path`"),
1357 ))
1358 }
1359 }
1360 } else {
1361 forwarding.variants.push(ident);
1362 }
1363 }
1364 if input.is_empty() {
1365 break;
1366 }
1367 input.parse::<Token![,]>()?;
1368 }
1369 if forwarding.all && !forwarding.variants.is_empty() {
1370 return Err(Error::new(
1371 forwarding.variants[0].span(),
1372 "`*` already names every variant",
1373 ));
1374 }
1375 if !forwarding.all && forwarding.variants.is_empty() {
1376 return Err(input.error("name the variants to forward to, or `*` for all of them"));
1377 }
1378 Ok(forwarding)
1379 }
1380}
1381
1382struct BodylessMethod {
1384 attrs: Vec<Attribute>,
1385 vis: Visibility,
1386 sig: Signature,
1387}
1388
1389impl Parse for BodylessMethod {
1390 fn parse(input: ParseStream) -> Result<Self> {
1391 let attrs = input.call(Attribute::parse_outer)?;
1392 let vis = input.parse()?;
1393 let sig = input.parse()?;
1394 input.parse::<Token![;]>()?;
1395 Ok(Self { attrs, vis, sig })
1396 }
1397}
1398
1399fn expand_inherent(input: InherentInput) -> Result<TokenStream2> {
1400 let InherentInput {
1401 item_enum,
1402 mut item_impl,
1403 } = input;
1404 let variants = item_enum
1405 .variants
1406 .iter()
1407 .map(variant_info)
1408 .collect::<Result<Vec<_>>>()?;
1409 let mut items = Vec::with_capacity(item_impl.items.len());
1410 for item in std::mem::take(&mut item_impl.items) {
1411 match item {
1412 ImplItem::Fn(mut function) => {
1413 let Some(position) = function.attrs.iter().position(is_dispatch_attribute) else {
1414 items.push(ImplItem::Fn(function));
1415 continue;
1416 };
1417 let forwarding: Forwarding = function.attrs.remove(position).parse_args()?;
1418 let generated = forward_method(
1419 function.attrs,
1420 function.vis,
1421 function.sig,
1422 Some(function.block),
1423 forwarding,
1424 &variants,
1425 &item_enum.ident,
1426 )?;
1427 items.push(ImplItem::Verbatim(generated));
1428 }
1429 ImplItem::Verbatim(tokens) => {
1430 let method: BodylessMethod = syn::parse2(tokens.clone()).map_err(|_| {
1431 Error::new_spanned(
1432 &tokens,
1433 "only methods, with or without a body, are forwarded in a `#[dispatch]` impl block",
1434 )
1435 })?;
1436 let BodylessMethod {
1437 mut attrs,
1438 vis,
1439 sig,
1440 } = method;
1441 let Some(position) = attrs.iter().position(is_dispatch_attribute) else {
1442 return Err(Error::new_spanned(
1443 &sig,
1444 "a method without a body needs `#[dispatch(*)]` to be forwarded to every variant",
1445 ));
1446 };
1447 let forwarding: Forwarding = attrs.remove(position).parse_args()?;
1448 let generated = forward_method(
1449 attrs,
1450 vis,
1451 sig,
1452 None,
1453 forwarding,
1454 &variants,
1455 &item_enum.ident,
1456 )?;
1457 items.push(ImplItem::Verbatim(generated));
1458 }
1459 other => items.push(other),
1460 }
1461 }
1462 item_impl.items = items;
1463 Ok(item_impl.to_token_stream())
1464}
1465
1466fn forward_method(
1470 attrs: Vec<Attribute>,
1471 vis: Visibility,
1472 sig: Signature,
1473 body: Option<Block>,
1474 forwarding: Forwarding,
1475 variants: &[VariantInfo],
1476 enum_ident: &Ident,
1477) -> Result<TokenStream2> {
1478 let method = &sig.ident;
1479 let Some(FnArg::Receiver(receiver)) = sig.inputs.first() else {
1480 return Err(Error::new_spanned(
1481 &sig,
1482 format!("`{method}` has no receiver, so there is no variant to forward to"),
1483 ));
1484 };
1485 if receiver.colon_token.is_some() {
1486 return Err(Error::new_spanned(
1487 receiver,
1488 "explicit `self: Type` receivers are not supported",
1489 ));
1490 }
1491 let self_token = receiver.self_token;
1492 let mut arguments = Vec::new();
1493 for input in sig.inputs.iter().skip(1) {
1494 let FnArg::Typed(pat_type) = input else {
1495 unreachable!("receiver handled above")
1496 };
1497 let Pat::Ident(pattern) = &*pat_type.pat else {
1498 return Err(Error::new_spanned(
1499 &pat_type.pat,
1500 "a forwarded parameter must be a plain identifier",
1501 ));
1502 };
1503 arguments.push(pattern.ident.clone());
1504 }
1505 let selected: Vec<&VariantInfo> = if forwarding.all {
1506 variants.iter().collect()
1507 } else {
1508 forwarding
1509 .variants
1510 .iter()
1511 .map(|name| {
1512 variants
1513 .iter()
1514 .find(|variant| &variant.ident == name)
1515 .ok_or_else(|| {
1516 Error::new(
1517 name.span(),
1518 format!("`{enum_ident}` has no variant `{name}`"),
1519 )
1520 })
1521 })
1522 .collect::<Result<_>>()?
1523 };
1524 let covers_all = forwarding.all
1525 || variants
1526 .iter()
1527 .all(|variant| forwarding.variants.contains(&variant.ident));
1528 match (covers_all, &body) {
1529 (true, Some(body)) => {
1530 return Err(Error::new_spanned(
1531 body,
1532 format!("`{method}` is forwarded to every variant, so its body would never run; end the signature with `;` instead"),
1533 ))
1534 }
1535 (false, None) => {
1536 return Err(Error::new_spanned(
1537 &sig,
1538 format!("`{method}` is forwarded to some variants only; give it a body, which is what the other variants return"),
1539 ))
1540 }
1541 _ => {}
1542 }
1543 let inner = Ident::new("__inner", Span::mixed_site());
1544 let call_name = forwarding.call.as_ref().unwrap_or(method);
1545 let awaited = sig.asyncness.is_some();
1546 let arms = selected.iter().map(|variant| {
1547 let VariantInfo { ident, cfgs, .. } = variant;
1548 let mut call = quote! { #inner.#call_name(#(#arguments),*) };
1549 if awaited {
1550 call = quote! { #call.await };
1551 }
1552 if let Some(with) = &forwarding.with {
1553 call = quote! { #with(#call) };
1554 }
1555 quote! { #(#cfgs)* Self::#ident(#inner) => #call, }
1556 });
1557 let fallback = body.map(|body| quote! { _ => #body, });
1558 Ok(quote! {
1559 #(#attrs)*
1560 #vis #sig {
1561 match #self_token {
1562 #(#arms)*
1563 #fallback
1564 }
1565 }
1566 })
1567}