compose_macro/
lib.rs

1// Function composition:
2//    - Function composition is the process of combining two or more functions to produce a new function.
3//    - The result of each function is passed as an argument to the next function.
4//
5// The macro invocation should look like so:
6//    - compose!(h -> g -> f)(x) == f(g(h(x)))
7//
8// This module defines a procedural macro for function composition.
9
10use proc_macro::TokenStream;
11use syn::{parse::Parse, Ident};
12
13// The `Composed` struct represents the parsed input for the `compose!` macro.
14// It contains the first function and a list of additional functions to compose.
15struct Composed {
16    function: Ident,    // The first function in the composition chain.
17    others: Vec<Ident>, // A vector of additional functions to apply in sequence.
18}
19
20// Implement the `Parse` trait for `Composed` to enable parsing the macro input.
21impl Parse for Composed {
22    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
23        // Parse the first function identifier.
24        let function = input.parse()?;
25        let mut others = Vec::new();
26
27        // Parse additional functions separated by the `->` token.
28        while input.parse::<syn::Token![->]>().is_ok() {
29            let other: Ident = input.parse()?;
30            others.push(other);
31        }
32
33        // Return the parsed `Composed` struct.
34        Ok(Composed { function, others })
35    }
36}
37
38// Implement the `ToTokens` trait for `Composed` to generate the composed function.
39impl quote::ToTokens for Composed {
40    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
41        let func = &self.function;
42
43        if self.others.is_empty() {
44            // If there are no additional functions, output the first function directly.
45            func.to_tokens(tokens);
46            return;
47        }
48
49        // Start with the first function applied to the input `x`.
50        let mut composed = quote::quote! {
51            #func(x)
52        };
53
54        // Chain the remaining functions in reverse order (right-to-left composition).
55        for other in &self.others {
56            composed = quote::quote! {
57                #other(#composed)
58            };
59        }
60
61        // Output the generated code.
62        composed.to_tokens(tokens);
63    }
64}
65
66// The `compose` procedural macro entry point.
67#[proc_macro]
68pub fn compose(input: TokenStream) -> TokenStream {
69    // Parse the input into a `Composed` struct.
70    let composed = syn::parse_macro_input!(input as Composed);
71
72    let out = if composed.others.is_empty() {
73        // If there are no additional functions, return the first function directly.
74        quote::quote! {
75            #composed
76        }
77    } else {
78        // Generate a closure that takes an input `x` and applies the composed functions.
79        quote::quote! {
80          |x| #composed
81        }
82    };
83
84    // Return the generated code as a `TokenStream`.
85    out.into()
86}