dodo-bindgen-codegen 2.0.0

use std::collections::{BTreeSet, HashMap, HashSet};

use crate::ffi::FfiSymbol;
use crate::util::*;
use crate::{ClassDefs, FfiFunction, FfiLib, FfiType, FunctionDefs, RustTypeDefs, UserDefinedDefs};
use dodo_bindgen_rust_ir::{Item, ItemClass, ItemFn, ItemHelper, ItemMethod, ItemTs, Type};
use syn::Ident;

use super::fn_::ItemFnExt as _;

/* -------------------------------------------------------------------------- */

#[derive(Clone, Debug)]
pub struct CodegenOpts {
    pub file_name: String,
    pub lib_name: String,
    pub lazy: bool,
}

#[derive(Clone, Debug, Default)]
pub struct TsModule {
    pub ffi_lib: FfiLib, // symbol definitions
    pub type_defs: RustTypeDefs, /* import statements or type definitions if `inline = true`.
                          * links to standard types */
    pub user_defs: UserDefinedDefs,
    pub functions: FunctionDefs,
    pub class_defs: ClassDefs,
    pub class_types: HashMap<Type, Ident>,
    pub used_helpers: BTreeSet<String>,
    pub uses_pin: bool,
    pub runtime: String,
    pub helpers: String,
}

const PIN_HELPERS: &str = r#"
const __dodo_pinned = new Set<unknown[]>();
"#;

// TODO: use buffer where possible to handle large projects

fn type_contains(haystack: &Type, needle: &Type) -> bool {
    if haystack == needle {
        return true;
    }
    match haystack {
        Type::Ptr(elem)
        | Type::PtrMut(elem)
        | Type::Ref(elem)
        | Type::RefMut(elem)
        | Type::Box(elem)
        | Type::Slice(elem)
        | Type::Vec(elem)
        | Type::Array { elem, .. } => type_contains(elem, needle),
        Type::Tuple(elems) => elems.iter().any(|elem| type_contains(elem, needle)),
        Type::FnPtr { inputs, output } => {
            inputs.iter().any(|input| type_contains(input, needle)) || type_contains(output, needle)
        },
        _ => false,
    }
}

fn fn_types(func: &ItemFn) -> impl Iterator<Item = &Type> {
    func.inputs
        .iter()
        .map(|arg| &arg.ty)
        .chain(std::iter::once(&func.output))
}

fn is_foreign(class: &ItemClass) -> bool {
    !matches!(class.ty, Type::Object { .. })
}

fn kept_classes(classes: &[ItemClass], methods: &[ItemMethod], fns: &[ItemFn]) -> HashSet<String> {
    let mut kept: HashSet<String> = classes
        .iter()
        .filter(|class| !is_foreign(class))
        .map(|class| class.name.clone())
        .collect();
    let mut roots: Vec<&Type> = fns.iter().flat_map(fn_types).collect();
    roots.extend(
        methods
            .iter()
            .filter(|method| kept.contains(&method.class))
            .flat_map(|method| fn_types(&method.func)),
    );

    loop {
        let newly_kept: Vec<&ItemClass> = classes
            .iter()
            .filter(|class| is_foreign(class) && !kept.contains(&class.name))
            .filter(|class| roots.iter().any(|root| type_contains(root, &class.ty)))
            .collect();
        if newly_kept.is_empty() {
            return kept;
        }
        for class in newly_kept {
            kept.insert(class.name.clone());
            roots.extend(
                methods
                    .iter()
                    .filter(|method| method.class == class.name)
                    .flat_map(|method| fn_types(&method.func)),
            );
        }
    }
}

impl TsModule {
    pub fn new(items: Vec<Item>, opts: &CodegenOpts) -> Self {
        let mut module = TsModule::default();
        module.ffi_lib.lib_name = opts.lib_name.clone();
        module.ffi_lib.lazy = opts.lazy;

        module.unwrap(items);
        module.user_defs.dedup(&module.class_defs);
        module
    }

    fn unwrap(&mut self, items: Vec<Item>) {
        let mut classes = Vec::new();
        let mut methods = Vec::new();
        let mut fns = Vec::new();
        let mut snippets = Vec::new();
        let mut helpers = Vec::new();
        for item in items {
            match item {
                Item::Fn(item_fn) => fns.push(item_fn),
                Item::Method(item_method) => methods.push(item_method),
                Item::Class(item_class) => classes.push(item_class),
                Item::Ts(item_ts) => snippets.push(item_ts),
                Item::Helper(item_helper) => helpers.push(item_helper),
            }
        }

        let kept = kept_classes(&classes, &methods, &fns);
        let foreign: HashSet<String> = classes
            .iter()
            .filter(|class| is_foreign(class))
            .map(|class| class.name.clone())
            .collect();
        let is_kept = |class: &str| !foreign.contains(class) || kept.contains(class);

        for ItemClass {
            name,
            drop_symbol,
            ty,
            ..
        } in classes.into_iter().filter(|class| is_kept(&class.name))
        {
            let class = format_ident!("{}", name);
            let drop_symbol = format_ident!("{}", drop_symbol);
            self.ffi_lib.interface.push_fn(FfiFunction {
                ident: drop_symbol.clone(),
                inputs: vec![FfiType::Pointer],
                output: FfiType::Void,
                non_blocking: false,
            });
            self.class_types.insert(ty, class.clone());
            self.class_defs.store.entry(class).or_default().drop_symbol = Some(drop_symbol);
        }

        for item_fn in fns {
            let method = item_fn.unwrap(self, None);
            self.functions.push(method);
        }

        for ItemMethod {
            class,
            receiver,
            func,
        } in methods.into_iter().filter(|method| is_kept(&method.class))
        {
            let class = format_ident!("{}", class);
            let method = func.unwrap(self, Some((class.clone(), receiver)));
            self.class_defs
                .store
                .entry(class)
                .or_default()
                .methods
                .push(method);
        }

        for ItemTs { class, code } in snippets
            .into_iter()
            .filter(|snippet| is_kept(&snippet.class))
        {
            self.class_defs
                .store
                .entry(format_ident!("{}", class))
                .or_default()
                .snippets
                .push(code);
        }

        self.add_helpers(helpers);
    }

    fn add_helpers(&mut self, helpers: Vec<ItemHelper>) {
        let mut seen = HashSet::new();
        for ItemHelper {
            id,
            runtime,
            code,
            symbols,
        } in helpers
        {
            let used = runtime || id.is_empty() || self.used_helpers.contains(&id);
            if !used || (!id.is_empty() && !seen.insert(id)) {
                continue;
            }
            for symbol in symbols {
                let ident = format_ident!("{}", symbol.name);
                let exists = self.ffi_lib.interface.symbols.iter().any(|existing| {
                    let FfiSymbol::Function(existing) = existing;
                    existing.ident == ident
                });
                if !exists {
                    self.ffi_lib.interface.push_fn(FfiFunction {
                        ident,
                        inputs: symbol.inputs.into_iter().map(FfiType::from).collect(),
                        output: FfiType::from(symbol.output),
                        non_blocking: false,
                    });
                }
            }
            let target = if runtime {
                &mut self.runtime
            } else {
                &mut self.helpers
            };
            target.push('\n');
            target.push_str(&code);
        }

        for id in &self.used_helpers {
            if !seen.contains(id) {
                eprintln!("warning: TypeScript helper `{id}` is used but no crate defines it");
            }
        }
    }

    fn helpers(&self) -> String {
        let mut helpers = self.helpers.clone();
        if self.uses_pin {
            helpers.push_str(PIN_HELPERS);
        }
        helpers
    }

    /// Generate a single typescript module with rust type definitions inlined
    pub fn generate_single(self, _opts: &CodegenOpts) -> String {
        let helpers = self.helpers();
        let TsModule {
            ffi_lib,
            user_defs,
            functions,
            class_defs,
            runtime: type_defs,
            ..
        } = self;

        let ffi_lib = ffi_lib.to_token_stream().to_string();
        let user_defs = user_defs.to_token_stream().to_string();
        let functions = functions.to_token_stream().to_string();
        let class_defs = class_defs.print();

        expand_raw(format!(
            "// deno-lint-ignore-file\n
            {ffi_lib}
            {type_defs}
            {helpers}
            {user_defs}
            {class_defs}
            {functions}
            "
        ))
    }

    /// Generate multiple typescript modules with rust type definitions on a
    /// separate file
    pub fn generate_multiple(self, _opts: &CodegenOpts, type_defs_name: &str) -> (String, String) {
        let helpers = self.helpers();
        let TsModule {
            mut ffi_lib,
            type_defs,
            user_defs,
            functions,
            class_defs,
            runtime,
            ..
        } = self;
        ffi_lib.export = true;

        let mut type_defs_builder = String::from("// deno-lint-ignore-file\n");
        type_defs_builder.push_str(&runtime);
        let (type_defs, imports) = type_defs.print_separate(type_defs_builder, type_defs_name);

        let ffi_lib = ffi_lib.to_token_stream().to_string();
        let user_defs = user_defs.to_token_stream().to_string();
        let functions = functions.to_token_stream().to_string();
        let class_defs = class_defs.print();

        let module = expand_raw(format!(
            "
            {imports}
            {ffi_lib}
            {helpers}
            {user_defs}
            {class_defs}
            {functions}
            "
        ));

        (module, type_defs)
    }
}