lutra-codegen 0.6.1

mod client;
mod encode;
mod functions;
mod program;
mod types;

use std::collections::{HashMap, HashSet, VecDeque};
use std::path;

use lutra_bin::ir;
use lutra_compiler::Project;

use crate::{GenerateOptions, infer_names};

fn canonical_std_type_reexport(module_path: &[String], name: &str) -> Option<&'static str> {
    match (module_path, name) {
        ([std], "Timestamp") if std == "std" => Some("std::Timestamp"),
        ([std], "Date") if std == "std" => Some("std::Date"),
        ([std], "Time") if std == "std" => Some("std::Time"),
        ([std], "Decimal") if std == "std" => Some("std::Decimal"),
        ([std, ops], "Ordering") if std == "std" && ops == "ops" => Some("std::ops::Ordering"),
        _ => None,
    }
}

#[derive(Debug)]
pub struct Context<'a> {
    current_rust_mod: Vec<String>,

    /// Buffer for types that don't have their own Lutra decl, but need their own Rust decl.
    /// When such type ref is encountered, it is pushed into here and generated later.
    def_buffer: VecDeque<ir::Ty>,

    /// Track already-emitted type names to avoid duplicate definitions.
    emitted_types: HashSet<String>,

    // static env
    options: &'a GenerateOptions,
    ty_defs: &'a HashMap<ir::Path, &'a ir::Ty>,
    project: &'a Project,
    out_dir: path::PathBuf,
}

impl<'a> Context<'a> {
    fn is_done(&self) -> bool {
        self.def_buffer.is_empty()
    }

    #[allow(dead_code)]
    fn get_ty_mat<'t: 'a>(&'t self, ty: &'t ir::Ty) -> &'t ir::Ty {
        if let ir::TyKind::Ident(path) = &ty.kind {
            self.ty_defs.get(path).unwrap()
        } else {
            ty
        }
    }
}

pub(crate) fn run(
    project: &Project,
    options: &GenerateOptions,
    out_dir: path::PathBuf,
) -> Result<String, std::fmt::Error> {
    use std::fmt::Write;

    let module = lutra_compiler::project_to_types(project);

    let ty_defs = module.iter_types_re().collect();

    let mut w = String::new();
    writeln!(w, "//# Generated by lutra-codegen\n")?;
    writeln!(w, "pub use {}::std;\n", options.lutra_bin_path)?;

    let mut ctx = Context {
        current_rust_mod: vec![],
        def_buffer: VecDeque::new(),
        emitted_types: HashSet::new(),

        options,
        ty_defs: &ty_defs,
        project,
        out_dir,
    };

    let module_path = vec![];
    codegen_module(&mut w, &module, module_path, &mut ctx)?;
    Ok(w)
}

fn codegen_module(
    w: &mut impl std::fmt::Write,
    module: &ir::Module,
    module_path: Vec<String>,
    ctx: &mut Context,
) -> Result<(), std::fmt::Error> {
    // collect defs
    let mut tys = Vec::new();
    let mut functions = Vec::new();
    let mut sub_modules = Vec::new();

    // iterate pr defs (which keep the order in the source)
    let root_mod = &ctx.project.root_module;
    let pr_mod = root_mod.get_module(&module_path).unwrap();
    for (name, pr_def) in &pr_mod.defs {
        let Some(decl) = module.decls.iter().find(|d| &d.name == name) else {
            continue;
        };

        match &decl.decl {
            ir::Decl::Mod(module) => {
                let is_dep = module_path.is_empty()
                    && ctx.project.dependencies.iter().any(|d| &d.name == name);
                if is_dep {
                    continue;
                }

                sub_modules.push((name, module));
            }
            ir::Decl::Ty(ty) => {
                let mut ty = ty.clone();
                infer_names(name, &mut ty);

                tys.push((ty, pr_def.annotations.as_slice()));
            }
            ir::Decl::Var(ty) => {
                let mut ty = ty.clone();
                infer_names(name, &mut ty);

                if let ir::TyKind::Function(func) = ty.kind {
                    functions.push((name, *func));
                }
            }
        }
    }

    ctx.current_rust_mod = module_path.clone();

    let mut all_tys = Vec::new();

    // write types
    if ctx.options.generates_types() {
        // partition
        let mut tys_to_import = Vec::new();
        let mut tys_to_generate = Vec::new();
        for (ty, annotations) in tys {
            let name = ty.name.as_deref().unwrap();
            if let Some(path) = canonical_std_type_reexport(&module_path, name) {
                tys_to_import.push((name.to_string(), path));
            } else {
                tys_to_generate.push((ty, annotations));
            }
        }

        // write imports
        let lutra_bin = &ctx.options.lutra_bin_path;
        for (name, path) in &tys_to_import {
            writeln!(w, "pub use {lutra_bin}::{path} as {name};")?;
        }
        if !tys_to_import.is_empty() {
            writeln!(w)?;
        }

        // write ty defs
        all_tys.extend(types::write_tys(w, tys_to_generate, ctx)?);
    };

    // write traits for functions
    if ctx.options.generates_function_traits() {
        functions::write_functions(w, &functions, ctx)?;

        all_tys.extend(types::write_tys_in_buffer(w, ctx)?);
    }

    // write programs
    let prog_repr = get_programs_repr_of(ctx, &module_path);
    if let Some(format) = prog_repr {
        program::write_rr_programs(w, &functions, format, ctx)?;

        all_tys.extend(types::write_tys_in_buffer(w, ctx)?);
    }

    let client_sub_modules = sub_modules
        .iter()
        .map(|(name, _)| *name)
        .filter(|name| {
            let mut path = module_path.clone();
            path.push(name.to_string());
            has_programs_in_subtree(ctx, &path)
        })
        .collect::<Vec<_>>();
    if ctx.options.generates_client() && (prog_repr.is_some() || !client_sub_modules.is_empty()) {
        let functions_here = if prog_repr.is_some() {
            functions.as_slice()
        } else {
            &[]
        };
        client::write_client(w, functions_here, &client_sub_modules, prog_repr, ctx)?;

        all_tys.extend(types::write_tys_in_buffer(w, ctx)?);
    }

    // recurse into sub modules
    for (name, sub_mod) in sub_modules {
        writeln!(w, "pub mod {name} {{")?;

        let mut path = module_path.clone();
        path.push(name.clone());
        codegen_module(w, sub_mod, path, ctx)?;
        writeln!(w, "}}\n")?;
    }

    // write encode/decode impls
    ctx.current_rust_mod = module_path.clone();
    if ctx.options.generates_encode_decode() {
        encode::write_encode_impls(w, &all_tys, ctx)?;
    }

    assert!(ctx.is_done(), "{ctx:?}");

    Ok(())
}

fn module_path_string(module_path: &[String]) -> String {
    module_path.join("::")
}

fn get_programs_repr_of(
    ctx: &Context,
    module_path: &[String],
) -> Option<lutra_compiler::ProgramRepr> {
    let module_path_str = module_path_string(module_path);
    ctx.options.included_program_repr(&module_path_str)
}

fn has_programs_in_subtree(ctx: &Context, module_path: &[String]) -> bool {
    let module_path_str = module_path_string(module_path);
    ctx.options.has_programs_in_subtree(&module_path_str)
}