1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use std::collections::hash_map::DefaultHasher;
6use std::hash::{Hash, Hasher};
7use syn::{
8 parse::{Parse, ParseStream},
9 punctuated::Punctuated,
10 Block, FnArg, Ident, Token, Type, Visibility,
11};
12
13struct PolymorphicFnInput {
14 vis: Visibility,
15 fn_token: Token![fn],
16 name: Ident,
17 inputs: Punctuated<FnArg, Token![,]>,
18 output: Type,
19 implementations: Punctuated<Block, Token![,]>,
20}
21
22impl Parse for PolymorphicFnInput {
23 fn parse(input: ParseStream) -> syn::Result<Self> {
24 let vis = input.parse()?;
25 let fn_token = input.parse()?;
26 let name = input.parse()?;
27 let content;
28 let _paren_token = syn::parenthesized!(content in input);
29 let inputs = Punctuated::parse_terminated(&content)?;
30 let _arrow_token = input.parse::<Token![->]>()?;
31 let output = input.parse()?;
32 let brace_content;
33 let _brace_token = syn::braced!(brace_content in input);
34 let implementations = Punctuated::parse_terminated(&brace_content)?;
35 Ok(PolymorphicFnInput {
36 vis,
37 fn_token,
38 name,
39 inputs,
40 output,
41 implementations,
42 })
43 }
44}
45
46#[proc_macro]
85pub fn polymorphic_fn(input: TokenStream) -> TokenStream {
86 let input = syn::parse_macro_input!(input as PolymorphicFnInput);
87 let implementations: Vec<&Block> = input.implementations.iter().collect();
88 if implementations.is_empty() {
89 return syn::Error::new_spanned(
90 &input.name,
91 "At least one implementation must be provided",
92 )
93 .to_compile_error()
94 .into();
95 }
96 let selected_index = deterministic_selection(&input.name, implementations.len());
97 let selected_block = implementations[selected_index];
98 let vis = &input.vis;
99 let fn_token = &input.fn_token;
100 let name = &input.name;
101 let inputs = &input.inputs;
102 let output = &input.output;
103 let output = quote! {
104 #vis #fn_token #name(#inputs) -> #output #selected_block
105 };
106 output.into()
107}
108
109fn deterministic_selection(name: &Ident, count: usize) -> usize {
111 let mut hasher = DefaultHasher::new();
113 name.to_string().hash(&mut hasher);
115 if let Ok(timestamp) = std::env::var("BUILD_TIMESTAMP") {
117 timestamp.hash(&mut hasher);
118 }
119 if let Ok(package) = std::env::var("CARGO_PKG_NAME") {
120 package.hash(&mut hasher);
121 }
122 if let Ok(version) = std::env::var("CARGO_PKG_VERSION") {
123 version.hash(&mut hasher);
124 }
125 if let Ok(target) = std::env::var("TARGET") {
126 target.hash(&mut hasher);
127 }
128 if let Ok(profile) = std::env::var("PROFILE") {
129 profile.hash(&mut hasher);
130 }
131 if let Ok(counter) = std::env::var("BUILD_COUNTER") {
134 counter.hash(&mut hasher);
135 }
136 let hash = hasher.finish();
138 (hash as usize) % count
139}