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}