Skip to main content

iris_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, parse::{Parse, ParseStream}, Token, Type, Ident};
4use std::path::Path;
5use std::fs::OpenOptions;
6use std::io::{Read, Write, Seek, SeekFrom};
7use fs2::FileExt;
8
9const MACRO_CRATE_DIR: &str = env!("CARGO_MANIFEST_DIR");
10
11fn next_slot(key: &str) -> u64 {
12    let state_path = Path::new(MACRO_CRATE_DIR).join(".iris_slots.json");
13
14    let mut file = OpenOptions::new()
15        .read(true)
16        .write(true)
17        .create(true)
18        .open(&state_path)
19        .unwrap();
20
21    file.lock_exclusive().unwrap();
22
23    let mut content = String::new();
24    file.read_to_string(&mut content).unwrap();
25
26    let mut map: std::collections::HashMap<String, u64> = if content.is_empty() {
27        std::collections::HashMap::new()
28    } else {
29        serde_json::from_str(&content).unwrap_or_default()
30    };
31
32    if let Some(&slot) = map.get(key) {
33        file.unlock().unwrap();
34        return slot;
35    }
36
37    let next = map.get("__next").copied().unwrap_or(0);
38    map.insert(key.to_string(), next);
39    map.insert("__next".to_string(), next + 1);
40
41    let new_content = serde_json::to_string_pretty(&map).unwrap();
42    file.seek(SeekFrom::Start(0)).unwrap();
43    file.set_len(0).unwrap();
44    file.write_all(new_content.as_bytes()).unwrap();
45    file.unlock().unwrap();
46
47    next
48}
49
50struct RegisterCapability {
51    cap_name: Ident,
52    trait_bound: Type,
53    returns: Type,
54}
55
56impl Parse for RegisterCapability {
57    fn parse(input: ParseStream) -> syn::Result<Self> {
58        let _: Ident = input.parse()?;
59        let _: Token![=] = input.parse()?;
60        let cap_name: Ident = input.parse()?;
61        let _: Token![,] = input.parse()?;
62
63        let _: Ident = input.parse()?;
64        let _: Token![=] = input.parse()?;
65        let trait_bound: Type = input.parse()?;
66        let _: Token![,] = input.parse()?;
67
68        let _: Ident = input.parse()?;
69        let _: Token![=] = input.parse()?;
70        let returns: Type = input.parse()?;
71
72        let _ = input.parse::<Token![,]>();
73
74        Ok(Self { cap_name, trait_bound, returns })
75    }
76}
77
78#[proc_macro]
79pub fn register_capability(input: TokenStream) -> TokenStream {
80    let RegisterCapability { cap_name, trait_bound, returns } =
81        parse_macro_input!(input as RegisterCapability);
82
83    let pkg_name = std::env::var("CARGO_PKG_NAME").unwrap_or_else(|_| "unknown".to_string());
84    let key = format!("{}::{}", pkg_name, cap_name);
85    let slot = next_slot(&key);
86
87    let output = quote! {
88        pub struct #cap_name;
89        impl iris::Capability for #cap_name {
90            type Handle = #returns;
91            const SLOT: u64 = #slot;
92        }
93
94        impl<'a, T: #trait_bound + 'static> iris::HasCap<#cap_name, #slot>
95            for iris::Probe<'a, T, #slot>
96        {
97            fn probe(&self, caps: &mut iris::CapabilityMap) {
98                caps.insert_cast::<#cap_name>(
99                    |any| {
100                        let val = any.downcast_ref::<T>().unwrap();
101                        val as &#returns
102                    },
103                    |any| {
104                        let val = any.downcast_mut::<T>().unwrap();
105                        val as &mut #returns
106                    },
107                );
108            }
109        }
110    };
111
112    output.into()
113}