ezffi-macros 0.1.1

use std::{
    collections::{BTreeMap, HashMap, HashSet},
    env, fs,
    path::{Path, PathBuf},
    sync::LazyLock,
};

use quote::{format_ident, quote};
use syn::{Item, Type};

use crate::{Namer, is_c_compatible_enum, typemap};

// TODO: see what happends with nam collisions beetwen dependencies and how to solve it

const PRIMITIVE_TYPES: &[&str] = &[
    "bool", "char", "i8", "u8", "i16", "u16", "i32", "u32", "i64", "u64", "i128", "u128", "isize",
    "usize", "f32", "f64",
];

struct TypeInfo {
    c_type: String,
    c_compatible: bool,
    krate: Option<String>,
}

struct Registry {
    map: HashMap<String, TypeInfo>,

    // Only filled when an init error occurs such as codegen mismatch
    init_error: Option<String>,
}

static REGISTRY: LazyLock<Registry> = LazyLock::new(build_registry);

pub struct FFITypeResolver;

impl FFITypeResolver {
    pub fn is_primitive(ty: &syn::Type) -> bool {
        if let syn::Type::Path(p) = ty {
            if let Some(ident) = p.path.get_ident() {
                PRIMITIVE_TYPES.contains(&ident.to_string().as_str())
            } else {
                false
            }
        } else {
            false
        }
    }

    pub fn is_c_callback(ty: &syn::Type) -> bool {
        let syn::Type::BareFn(bare) = ty else {
            return false;
        };
        match &bare.abi {
            // `extern "C" fn`, or bare `extern fn` which defaults to the C ABI
            Some(abi) => abi.name.as_ref().is_none_or(|n| n.value() == "C"),
            // Plain `fn` is `extern "Rust"`, not FFI-safe
            None => false,
        }
    }

    /// Whether the exported type with this name is C-compatible, that means,
    /// all their fields are also C-compatible.
    pub fn is_rust_type_c_compatible(name: &str) -> bool {
        REGISTRY
            .map
            .get(name)
            .map(|i| i.c_compatible)
            .unwrap_or(false)
    }

    /// Surface any error that fired while building the registry (typemap
    /// codegen mismatch is the only current cause).
    ///
    /// Each entry-point macro has to call this before doing any work.
    pub fn check_registry() -> syn::Result<()> {
        match &REGISTRY.init_error {
            None => Ok(()),
            Some(msg) => Err(syn::Error::new(proc_macro2::Span::call_site(), msg)),
        }
    }

    pub fn c_type_of(
        ty: &syn::Type,
        self_repl: Option<&Type>,
    ) -> syn::Result<proc_macro2::TokenStream> {
        if FFITypeResolver::is_primitive(ty) || FFITypeResolver::is_c_callback(ty) {
            return Ok(quote! { #ty });
        }

        match ty {
            syn::Type::Reference(r) => Self::c_type_of(&r.elem, self_repl),
            syn::Type::Slice(_) => Ok(quote! { ::ezffi::EzffiSlice }),
            syn::Type::Path(p) if p.path.segments[0].ident == "str" => {
                Ok(quote! { ::ezffi::EzffiStr })
            }
            syn::Type::Path(p) if p.path.segments[0].ident == "Self" => {
                let target = self_repl.ok_or_else(|| {
                    syn::Error::new_spanned(p, "`Self` used outside an `impl` block")
                })?;

                Self::c_type_of(target, None)
            }
            syn::Type::Path(path) => {
                let ident = path.path.segments[0].ident.to_string();

                if let Some(info) = REGISTRY.map.get(&ident) {
                    let c_type = format_ident!("{}", &info.c_type);
                    match &info.krate {
                        None => Ok(quote! { #c_type }),
                        Some(krate) => {
                            let krate = format_ident!("{}", krate);
                            Ok(quote! { #krate::#c_type })
                        }
                    }
                } else {
                    // TODO: Is this reallt needed??
                    // We should ensure all Ezffi types are in the map
                    let ty = format_ident!("Ezffi{}", ident);
                    Ok(quote! { ezffi::#ty })
                }
            }
            _ => Err(syn::Error::new_spanned(
                ty,
                format!("unsupported type in #[ezffi::export]: `{}`", quote!(#ty)),
            )),
        }
    }
}

fn type_name(ty: &syn::Type) -> Option<String> {
    match ty {
        syn::Type::Path(p) => p.path.segments.last().map(|s| s.ident.to_string()),
        _ => None,
    }
}

fn build_registry() -> Registry {
    let Ok(manifest) = env::var("CARGO_MANIFEST_DIR") else {
        return Registry {
            map: HashMap::new(),
            init_error: None,
        };
    };

    // --- Scan this crate's source for every exported type ---
    let mut files = Vec::new();
    collect_rs_files(&Path::new(&manifest).join("src"), &mut files);

    // Non-generic exported structs: name -> field types
    let mut struct_fields: HashMap<String, Vec<Type>> = HashMap::new();
    // Non-generic exported enums: name -> all-variants-are-unit
    let mut enum_c_compat: HashSet<String> = HashSet::new();
    // Types forced opaque up-front: generics + `export_extern_type!` + complex enum types
    let mut forced_opaque: HashSet<String> = HashSet::new();
    // Every exported name, to tell "exported but undecided" from "external"
    let mut all_exported: HashSet<String> = HashSet::new();

    for file in &files {
        let Ok(src) = fs::read_to_string(file) else {
            continue;
        };
        let Ok(ast) = syn::parse_file(&src) else {
            continue;
        };
        scan_items(
            &ast.items,
            &mut struct_fields,
            &mut enum_c_compat,
            &mut forced_opaque,
            &mut all_exported,
        );
    }

    // --- Load the types exported by dependency crates. ---
    let this_crate = crate::PKG_NAME.replace('-', "_");
    let type_map = typemap::TypeMap::load();

    let mut registry: HashMap<String, TypeInfo> = HashMap::new();
    let mut foreign_c_compat: HashMap<String, bool> = HashMap::new();

    let mut init_error: Option<String> = None;

    if let Some(type_map) = &type_map {
        for (krate, section) in type_map.sections() {
            if *krate == this_crate {
                continue; // our own section is recomputed below
            }

            let v = section.codegen_version;

            if v < crate::MIN_COMPATIBLE_CODEGEN {
                let min_crate = crate::crate_version_for_codegen(crate::MIN_COMPATIBLE_CODEGEN);

                init_error.get_or_insert(format!(
                    "ezffi codegen mismatch: dependency `{krate}` was generated with \
                     codegen v{v}, but this crate only reads v{}..=v{}. \
                     Rebuild `{krate}` against ezffi >= {min_crate} if possible.",
                    crate::MIN_COMPATIBLE_CODEGEN,
                    crate::EZFFI_CODEGEN_VERSION,
                ));

                continue;
            }

            if v > crate::EZFFI_CODEGEN_VERSION {
                let needed_crate = crate::crate_version_for_codegen(v);

                init_error.get_or_insert(format!(
                    "ezffi codegen mismatch: dependency `{krate}` was generated with \
                     codegen v{v}, but this crate only reads v{}..=v{}. \
                     Upgrade this crate's ezffi to >= {needed_crate}.",
                    crate::MIN_COMPATIBLE_CODEGEN,
                    crate::EZFFI_CODEGEN_VERSION,
                ));

                continue;
            }

            for (rust_ty, entry) in &section.types {
                foreign_c_compat.insert(rust_ty.clone(), entry.c_compatible);
                registry.insert(
                    rust_ty.clone(),
                    TypeInfo {
                        c_type: entry.c_type.clone(),
                        c_compatible: entry.c_compatible,
                        krate: Some(krate.clone()),
                    },
                );
            }
        }
    }

    // --- Decide C-compatibility for this crate's structs. ---
    // Fixpoint: a struct is C-compatible iff it is non-empty and every field
    // is a primitive, a callback, or another C-compatible exported type.
    let mut c_compat: HashMap<String, bool> = HashMap::new();

    for name in &enum_c_compat {
        c_compat.insert(name.clone(), true);
    }

    for name in &forced_opaque {
        c_compat.insert(name.clone(), false);
    }

    loop {
        let mut changed = false;

        for (name, fields) in &struct_fields {
            if c_compat.contains_key(name) {
                continue;
            }
            if let Some(decided) =
                decide_struct(fields, &c_compat, &foreign_c_compat, &all_exported)
            {
                c_compat.insert(name.clone(), decided);
                changed = true;
            }
        }

        if !changed {
            break;
        }
    }

    // Anything still undecided depends on an unresolvable type → opaque.
    for name in struct_fields.keys() {
        c_compat.entry(name.clone()).or_insert(false);
    }

    // --- Register this crate's types and write its section back. ---
    let mut own_types: BTreeMap<String, typemap::TypeEntry> = BTreeMap::new();

    for name in &all_exported {
        let c_compatible = c_compat.get(name).copied().unwrap_or(false);
        let c_type_name = Namer::name_struct(&format_ident!("{}", name)).to_string();

        own_types.insert(
            name.clone(),
            typemap::TypeEntry {
                c_type: c_type_name.clone(),
                c_compatible,
            },
        );

        // Own types override any dependency type that shares their name.
        registry.insert(
            name.clone(),
            TypeInfo {
                c_type: c_type_name,
                c_compatible,
                krate: None,
            },
        );
    }

    // Write this crate's section back; consuming the handle releases the lock.
    if let Some(type_map) = type_map {
        type_map.store(
            this_crate,
            typemap::Section {
                codegen_version: crate::EZFFI_CODEGEN_VERSION,
                types: own_types,
            },
        );
    }

    Registry {
        map: registry,
        init_error,
    }
}

fn decide_struct(
    fields: &[Type],
    c_compat: &HashMap<String, bool>,
    foreign_c_compat: &HashMap<String, bool>,
    all_exported: &HashSet<String>,
) -> Option<bool> {
    if fields.is_empty() {
        return Some(false);
    }
    for ty in fields {
        if FFITypeResolver::is_primitive(ty) || FFITypeResolver::is_c_callback(ty) {
            continue;
        }
        let Some(name) = type_name(ty) else {
            // References, tuples, etc. — not a plain C value.
            return Some(false);
        };
        // A type exported by this crate, still being decided.
        match c_compat.get(&name) {
            Some(true) => continue,
            Some(false) => return Some(false),
            None => {}
        }
        // A type exported by a dependency crate — already decided.
        match foreign_c_compat.get(&name) {
            Some(true) => continue,
            Some(false) => return Some(false),
            None => {}
        }
        if all_exported.contains(&name) {
            return None; // exported here but not decided yet — defer
        }
        return Some(false); // external type (String, Rc, ...)
    }
    Some(true)
}

fn scan_items(
    items: &[Item],
    struct_fields: &mut HashMap<String, Vec<Type>>,
    enum_c_compat: &mut HashSet<String>,
    forced_opaque: &mut HashSet<String>,
    all_exported: &mut HashSet<String>,
) {
    for item in items {
        match item {
            Item::Mod(m) => {
                if let Some((_, inner)) = &m.content {
                    scan_items(
                        inner,
                        struct_fields,
                        enum_c_compat,
                        forced_opaque,
                        all_exported,
                    );
                }
            }
            Item::Struct(s) if has_export_attr(&s.attrs) => {
                let name = s.ident.to_string();
                all_exported.insert(name.clone());
                if s.generics.gt_token.is_some() {
                    forced_opaque.insert(name);
                } else {
                    struct_fields.insert(name, s.fields.iter().map(|f| f.ty.clone()).collect());
                }
            }
            Item::Enum(e) if has_export_attr(&e.attrs) => {
                let name = e.ident.to_string();
                all_exported.insert(name.clone());
                if is_c_compatible_enum(e) {
                    enum_c_compat.insert(name);
                } else {
                    forced_opaque.insert(name);
                }
            }
            Item::Macro(mac) if mac.mac.path.is_ident("export_extern_type") => {
                if let Ok(path) = mac.mac.parse_body::<syn::Path>()
                    && let Some(seg) = path.segments.last()
                {
                    let name = seg.ident.to_string();
                    all_exported.insert(name.clone());
                    forced_opaque.insert(name);
                }
            }
            _ => {}
        }
    }
}

fn has_export_attr(attrs: &[syn::Attribute]) -> bool {
    attrs.iter().any(|a| {
        // Direct form: `#[ezffi::export]`
        if a.path()
            .segments
            .last()
            .map(|s| s.ident == "export")
            .unwrap_or(false)
        {
            return true;
        }

        // `#[cfg_attr(..., ezffi::export)]`
        if a.path().is_ident("cfg_attr")
            && let Ok(args) = a.parse_args_with(
                syn::punctuated::Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated,
            )
        {
            return args.iter().any(|m| {
                m.path()
                    .segments
                    .last()
                    .map(|s| s.ident == "export")
                    .unwrap_or(false)
            });
        }
        false
    })
}

fn collect_rs_files(dir: &Path, out: &mut Vec<PathBuf>) {
    let Ok(entries) = fs::read_dir(dir) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            collect_rs_files(&path, out);
        } else if path.extension().and_then(|e| e.to_str()) == Some("rs") {
            out.push(path);
        }
    }
}