slovo/compiler/src/docgen.rs

use std::{fs, path::Path};

use crate::{
    ast::Program,
    diag::Diagnostic,
    lexer, lower,
    project::{self, SourceFile, ToolFailure},
    sexpr::{Atom, SExpr, SExprKind},
};

pub fn generate(input: &str, output_dir: &str) -> Result<(), ToolFailure> {
    let docs = if project::is_project_input(input) {
        let loaded = project::load_project_sources_for_tools(input)?;
        render_project(&loaded.artifact.project_name, &loaded.sources)?
    } else {
        let source = fs::read_to_string(input).map_err(|err| ToolFailure {
            diagnostics: vec![Diagnostic::new(
                input,
                "InputReadFailed",
                format!("cannot read `{}`: {}", input, err),
            )],
            sources: Vec::new(),
            artifact: None,
        })?;
        let sources = vec![SourceFile {
            path: input.to_string(),
            source,
        }];
        render_file(input, &sources)?
    };

    fs::create_dir_all(output_dir).map_err(|err| ToolFailure {
        diagnostics: vec![Diagnostic::new(
            output_dir,
            "OutputWriteFailed",
            format!(
                "cannot create documentation directory `{}`: {}",
                output_dir, err
            ),
        )],
        sources: Vec::new(),
        artifact: None,
    })?;
    fs::write(Path::new(output_dir).join("index.md"), docs).map_err(|err| ToolFailure {
        diagnostics: vec![Diagnostic::new(
            output_dir,
            "OutputWriteFailed",
            format!("cannot write documentation into `{}`: {}", output_dir, err),
        )],
        sources: Vec::new(),
        artifact: None,
    })
}

fn render_project(title: &str, sources: &[SourceFile]) -> Result<String, ToolFailure> {
    let mut out = String::new();
    out.push_str("# Project ");
    out.push_str(title);
    out.push_str("\n\n");
    for source in sources {
        let module = document_source(source)?;
        render_module(&mut out, &module);
    }
    Ok(out)
}

fn render_file(title: &str, sources: &[SourceFile]) -> Result<String, ToolFailure> {
    let mut out = String::new();
    out.push_str("# File ");
    out.push_str(title);
    out.push_str("\n\n");
    let module = document_source(&sources[0])?;
    render_module(&mut out, &module);
    Ok(out)
}

fn document_source(source: &SourceFile) -> Result<DocModule, ToolFailure> {
    let tokens = lexer::lex(&source.path, &source.source).map_err(|diagnostics| ToolFailure {
        diagnostics,
        sources: vec![source.clone()],
        artifact: None,
    })?;
    let forms = crate::sexpr::parse(&source.path, &tokens).map_err(|diagnostics| ToolFailure {
        diagnostics,
        sources: vec![source.clone()],
        artifact: None,
    })?;
    let program = lower::lower_program(&source.path, &forms).ok();
    Ok(module_from_forms(&source.path, &forms, program.as_ref()))
}

struct DocModule {
    title: String,
    imports: Vec<String>,
    exports: Vec<String>,
    structs: Vec<String>,
    functions: Vec<String>,
    tests: Vec<String>,
}

fn module_from_forms(file: &str, forms: &[SExpr], program: Option<&Program>) -> DocModule {
    let mut module = DocModule {
        title: file.to_string(),
        imports: Vec::new(),
        exports: Vec::new(),
        structs: Vec::new(),
        functions: Vec::new(),
        tests: Vec::new(),
    };

    for form in forms {
        let Some(items) = list(form) else {
            continue;
        };
        match items.first().and_then(ident) {
            Some("module") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.title = name.to_string();
                }
                if let Some(exports) = items.get(2).and_then(list) {
                    if matches!(exports.first().and_then(ident), Some("export")) {
                        module
                            .exports
                            .extend(exports[1..].iter().filter_map(ident).map(str::to_string));
                    }
                }
            }
            Some("import") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.imports.push(name.to_string());
                }
            }
            Some("struct") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.structs.push(name.to_string());
                }
            }
            Some("fn") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.functions.push(name.to_string());
                }
            }
            Some("test") => {
                if let Some(name) = items.get(1).and_then(string_atom) {
                    module.tests.push(name.to_string());
                }
            }
            _ => {}
        }
    }

    if let Some(program) = program {
        module.title = program.module.clone();
        module.structs = program
            .structs
            .iter()
            .map(|decl| decl.name.clone())
            .collect();
        module.functions = program
            .functions
            .iter()
            .map(|function| {
                let params = function
                    .params
                    .iter()
                    .map(|param| format!("{} {}", param.name, param.ty))
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("{}({}) -> {}", function.name, params, function.return_type)
            })
            .collect();
        module.tests = program.tests.iter().map(|test| test.name.clone()).collect();
    }

    module.imports.sort();
    module.exports.sort();
    module.structs.sort();
    module.functions.sort();
    module.tests.sort();
    module
}

fn render_module(out: &mut String, module: &DocModule) {
    out.push_str("## Module ");
    out.push_str(&module.title);
    out.push_str("\n\n");
    render_list(out, "Imports", &module.imports);
    render_list(out, "Exports", &module.exports);
    render_list(out, "Structs", &module.structs);
    render_list(out, "Functions", &module.functions);
    render_list(out, "Tests", &module.tests);
}

fn render_list(out: &mut String, title: &str, values: &[String]) {
    out.push_str("### ");
    out.push_str(title);
    out.push('\n');
    if values.is_empty() {
        out.push_str("\nNone.\n\n");
        return;
    }
    for value in values {
        out.push_str("- `");
        out.push_str(&value.replace('`', "\\`"));
        out.push_str("`\n");
    }
    out.push('\n');
}

fn list(expr: &SExpr) -> Option<&[SExpr]> {
    match &expr.kind {
        SExprKind::List(items) => Some(items),
        _ => None,
    }
}

fn ident(expr: &SExpr) -> Option<&str> {
    match &expr.kind {
        SExprKind::Atom(Atom::Ident(value)) => Some(value),
        _ => None,
    }
}

fn string_atom(expr: &SExpr) -> Option<&str> {
    match &expr.kind {
        SExprKind::Atom(Atom::String(value)) => Some(value),
        _ => None,
    }
}