rustmorphism/
lib.rs

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/// Defines a function with multiple implementations, where one is deterministically selected at compile-time.
47///
48/// At compile-time, the macro selects one of the provided implementation blocks based on a hash of
49/// source location and other build identifiers. Other implementations are excluded from the final binary.
50///
51/// # Syntax
52///
53/// ```text
54/// polymorphic_fn! {
55///     [vis] fn name([args]) -> return_type {
56///         { implementation1 },
57///         { implementation2 },
58///         ...
59///     }
60/// }
61/// ```
62///
63/// - `[vis]`: Optional visibility (e.g., `pub`), defaults to private if omitted.
64/// - `name`: The function name.
65/// - `[args]`: Comma-separated argument list (e.g., `x: i32`).
66/// - `return_type`: The return type of the function (e.g., `i32`).
67/// - `{ implementation }`: A block containing an implementation; at least one must be provided.
68///
69/// # Example
70///
71/// ```rust
72/// use rustmorphism::polymorphic_fn;
73///
74/// polymorphic_fn! {
75///     pub fn compute(x: i32) -> i32 {
76///         { x + 1 },
77///         { x * 2 },
78///         { x - 3 }
79///     }
80/// }
81///
82/// println!("Result: {}", compute(5));
83/// ```
84#[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
109/// Deterministically selects an implementation based on build identifiers
110fn deterministic_selection(name: &Ident, count: usize) -> usize {
111    // Create a hasher for deterministic selection
112    let mut hasher = DefaultHasher::new();
113    // Hash the function name
114    name.to_string().hash(&mut hasher);
115    // Use various environment variables for entropy
116    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 we need more entropy, we can use a counter that increments on every build
132    // through the build.rs mechanism
133    if let Ok(counter) = std::env::var("BUILD_COUNTER") {
134        counter.hash(&mut hasher);
135    }
136    // Get the hash value and select an implementation
137    let hash = hasher.finish();
138    (hash as usize) % count
139}