iris-macros 0.0.1

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, parse::{Parse, ParseStream}, Token, Type, Ident};
use std::path::Path;
use std::fs::OpenOptions;
use std::io::{Read, Write, Seek, SeekFrom};
use fs2::FileExt;

const MACRO_CRATE_DIR: &str = env!("CARGO_MANIFEST_DIR");

fn next_slot(key: &str) -> u64 {
    let state_path = Path::new(MACRO_CRATE_DIR).join(".iris_slots.json");

    let mut file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(&state_path)
        .unwrap();

    file.lock_exclusive().unwrap();

    let mut content = String::new();
    file.read_to_string(&mut content).unwrap();

    let mut map: std::collections::HashMap<String, u64> = if content.is_empty() {
        std::collections::HashMap::new()
    } else {
        serde_json::from_str(&content).unwrap_or_default()
    };

    if let Some(&slot) = map.get(key) {
        file.unlock().unwrap();
        return slot;
    }

    let next = map.get("__next").copied().unwrap_or(0);
    map.insert(key.to_string(), next);
    map.insert("__next".to_string(), next + 1);

    let new_content = serde_json::to_string_pretty(&map).unwrap();
    file.seek(SeekFrom::Start(0)).unwrap();
    file.set_len(0).unwrap();
    file.write_all(new_content.as_bytes()).unwrap();
    file.unlock().unwrap();

    next
}

struct RegisterCapability {
    cap_name: Ident,
    trait_bound: Type,
    returns: Type,
}

impl Parse for RegisterCapability {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let _: Ident = input.parse()?;
        let _: Token![=] = input.parse()?;
        let cap_name: Ident = input.parse()?;
        let _: Token![,] = input.parse()?;

        let _: Ident = input.parse()?;
        let _: Token![=] = input.parse()?;
        let trait_bound: Type = input.parse()?;
        let _: Token![,] = input.parse()?;

        let _: Ident = input.parse()?;
        let _: Token![=] = input.parse()?;
        let returns: Type = input.parse()?;

        let _ = input.parse::<Token![,]>();

        Ok(Self { cap_name, trait_bound, returns })
    }
}

#[proc_macro]
pub fn register_capability(input: TokenStream) -> TokenStream {
    let RegisterCapability { cap_name, trait_bound, returns } =
        parse_macro_input!(input as RegisterCapability);

    let pkg_name = std::env::var("CARGO_PKG_NAME").unwrap_or_else(|_| "unknown".to_string());
    let key = format!("{}::{}", pkg_name, cap_name);
    let slot = next_slot(&key);

    let output = quote! {
        pub struct #cap_name;
        impl iris::Capability for #cap_name {
            type Handle = #returns;
            const SLOT: u64 = #slot;
        }

        impl<'a, T: #trait_bound + 'static> iris::HasCap<#cap_name, #slot>
            for iris::Probe<'a, T, #slot>
        {
            fn probe(&self, caps: &mut iris::CapabilityMap) {
                caps.insert_cast::<#cap_name>(
                    |any| {
                        let val = any.downcast_ref::<T>().unwrap();
                        val as &#returns
                    },
                    |any| {
                        let val = any.downcast_mut::<T>().unwrap();
                        val as &mut #returns
                    },
                );
            }
        }
    };

    output.into()
}