1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

use proc_macro::TokenStream;
use quote::format_ident;
use quote::quote;
use quote::ToTokens;
use std::cell::RefCell;
use std::collections::HashMap;
use syn::{parse_macro_input, AttributeArgs, ItemFn, Lit, LitInt, NestedMeta};

thread_local! {
    static AOC_SOLVERS: RefCell<HashMap<(String, String), String>> = RefCell::new(HashMap::new());
}

#[proc_macro_attribute]
pub fn aoc(input: TokenStream, annotated_item: TokenStream) -> TokenStream {
    let attributes = parse_macro_input!(input as AttributeArgs);
    let attributes: Vec<String> = attributes
        .iter()
        .map(|a| match a {
            NestedMeta::Lit(l) => match l {
                Lit::Str(s) => s.value(),
                _ => panic!("Attribute is not a string"),
            },
            _ => panic!("Attribute is not a string"),
        })
        .collect();
    if attributes.len() != 2 {
        panic!("Number of attributes must be two");
    }

    let func = parse_macro_input!(annotated_item as ItemFn);
    AOC_SOLVERS.with(|solvers| {
        solvers.borrow_mut().insert(
            (attributes[0].clone(), attributes[1].clone()),
            func.sig.ident.to_string(),
        );
    });
    func.into_token_stream().into()
}

#[proc_macro]
pub fn aoc_main(input: TokenStream) -> TokenStream {
    let year = parse_macro_input!(input as LitInt);

    let solvers = AOC_SOLVERS.with(|s| s.take()).into_iter();
    let quote_solvers = solvers.map(|(key, value)| {
        let (day, part) = key;
        let id = format_ident!("{}", value);
        quote! { aoc.add_solver(concat!(#day, #part), #id); }
    });
    quote! {
        fn main() {
            let mut aoc = Aoc::new(#year);
            #( #quote_solvers )*

            aoc.run();
        }
    }
    .into()
}