lutra-compiler 0.6.0

use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock};

use itertools::Itertools;

use crate::diagnostic::{Diagnostic, WithErrorInfo};
use crate::error::Error;
use crate::pr;
use crate::project::SourceOverlay;
use crate::project::SourceProvider;
use crate::project::{self, Dependency};
use crate::resolver::NS_STD;
use crate::{SourceTree, diagnostic};
use crate::{Span, error};

#[cfg_attr(feature = "clap", derive(clap::Parser))]
#[derive(Default, Clone)]
pub struct CheckParams {
    #[cfg_attr(feature = "clap", arg(skip))]
    pub dependencies: Vec<Dependency>,
}

pub fn check(
    mut source: project::SourceTree,
    params: CheckParams,
) -> Result<project::Project, error::Error> {
    // parse
    if source.is_empty() {
        source.insert(PathBuf::from(""), "".into());
    }
    let root_mod = parse(&source)?.map_err(|e| Error::from_diagnostics(e, &source))?;

    // resolve
    let metadata = root_mod.get_anno_at(&pr::Path::empty(), pr::Anno::as_std_metadata);
    let is_std = metadata.is_some_and(|m| m == NS_STD);
    let mut dependencies = params.dependencies;
    if !is_std {
        dependencies.push(Dependency {
            name: NS_STD.into(),
            inner: std_project()?,
        });
    }
    let mut project = crate::resolver::resolve(root_mod, dependencies, is_std)
        .map_err(|e| Error::from_diagnostics(e, &source))?;
    project.source = source;
    Ok(project)
}

pub fn check_overlay(
    project: &project::Project,
    overlay: &str,
    overlay_name: Option<&str>,
) -> Result<pr::Expr, error::Error> {
    let source = crate::project::SourceOverlay::new(&project.source, overlay, overlay_name);

    parse_overlay(&source)
        .and_then(|expr| crate::resolver::resolve_overlay_expr(&project.root_module, expr))
        .map_err(|e| Error::from_diagnostics(e, &source))
}

fn parse(tree: &SourceTree) -> Result<Result<pr::ModuleDef, Vec<Diagnostic>>, error::Error> {
    // init the root module def
    let mut root = pr::ModuleDef::default();

    // parse and insert into the root
    let mut diags = Vec::new();
    for source_id in tree.get_ids() {
        let (path, content) = tree.get_by_id(source_id).unwrap();

        let module_path = os_path_to_mod_path(path)?;

        let (parsed, errs, _) = crate::parser::parse_source(content, source_id);
        diags.extend(errs);
        if let Some(parsed) = parsed {
            // TODO: improve these error messages

            if module_path.is_empty() && parsed.is_submodule {
                diags.push(
                    Diagnostic::new_custom("cannot load the project root")
                        .with_span(Some(Span {
                            start: 0,
                            len: 1,
                            source_id,
                        }))
                        .push_hint(format!("file {} is a submodule", path.display())),
                );
            }
            let included = module_path.is_empty() || parsed.is_submodule;
            if included {
                diags.extend(insert_module_at_path(&mut root, module_path, parsed.root));
            }
        }
    }
    Ok(if diags.is_empty() {
        Ok(root)
    } else {
        Err(diags)
    })
}

fn parse_overlay(overlay: &SourceOverlay) -> Result<pr::Expr, Vec<Diagnostic>> {
    let id = SourceOverlay::overlay_id();
    let (_path, content) = overlay.get_by_id(id).unwrap();
    let (ast, diagnostics) = crate::parser::parse_expr(content, id);
    if diagnostics.is_empty() {
        Ok(ast.unwrap())
    } else {
        Err(diagnostics)
    }
}

pub fn insert_module_at_path(
    module: &mut pr::ModuleDef,
    mut path: Vec<String>,
    to_insert: pr::ModuleDef,
) -> Vec<diagnostic::Diagnostic> {
    if path.is_empty() {
        let mut d = Vec::new();

        module.annotations.extend(to_insert.annotations);
        module.span_content = to_insert.span_content;
        module.imports.extend(to_insert.imports);
        for (name, def) in to_insert.defs {
            let conflict = module.defs.insert(name, def);
            if let Some(conflict) = conflict {
                d.push(
                    diagnostic::Diagnostic::new_custom("duplicate name").with_span(conflict.span),
                );
            }
        }
        return d;
    }

    let step = path.remove(0);

    // find submodule def
    let submodule = module
        .defs
        .entry(step)
        .or_insert_with(|| pr::Def::new(pr::ModuleDef::default()));
    let pr::DefKind::Module(submodule) = &mut submodule.kind else {
        return vec![
            diagnostic::Diagnostic::new_custom("duplicate name").with_span(submodule.span),
        ];
    };
    insert_module_at_path(submodule, path, to_insert)
}

fn os_path_to_mod_path(path: &Path) -> Result<Vec<String>, error::Error> {
    let path = if path.ends_with("module.lt") {
        // remove module.lt suffix
        path.parent().unwrap().to_path_buf()
    } else {
        // remove file format extension
        path.with_extension("")
    };

    // split by /
    path.components()
        .map(|x| {
            x.as_os_str()
                .to_str()
                .map(str::to_string)
                .ok_or_else(|| error::Error::InvalidPath { path: path.clone() })
        })
        .try_collect()
}

pub fn std_source() -> SourceTree {
    SourceTree::single(
        std::path::PathBuf::new(),
        include_str!("std.lt").to_string(),
    )
}

/// Compile `std.lt` once and cache the result for the lifetime of the process.
static STD_PROJECT: OnceLock<Arc<crate::Project>> = OnceLock::new();

fn std_project() -> Result<Arc<crate::Project>, error::Error> {
    if let Some(cached) = STD_PROJECT.get() {
        return Ok(Arc::clone(cached));
    }

    let project = Arc::new(check(std_source(), CheckParams::new())?);
    let _ = STD_PROJECT.set(project);
    Ok(Arc::clone(STD_PROJECT.get().unwrap()))
}

impl CheckParams {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_dep(mut self, name: impl Into<String>, project: Arc<crate::Project>) -> Self {
        self.dependencies.push(crate::project::Dependency {
            name: name.into(),
            inner: project,
        });
        self
    }
}