forge_doc/

builder.rs

use crate::{
    AsDoc, BufWriter, Document, ParseItem, ParseSource, Parser, Preprocessor,
    document::DocumentContent, helpers::merge_toml_table,
};
use alloy_primitives::map::HashMap;
use eyre::{Context, Result};
use foundry_compilers::{compilers::solc::SOLC_EXTENSIONS, utils::source_files_iter};
use foundry_config::{DocConfig, FormatterConfig, filter::expand_globs};
use itertools::Itertools;
use mdbook_driver::MDBook;
use rayon::prelude::*;
use std::{
    cmp::Ordering,
    fs,
    path::{Path, PathBuf},
};
use toml::value;

/// Build Solidity documentation for a project from natspec comments.
/// The builder parses the source files using [Parser],
/// then formats and writes the elements as the output.
#[derive(Debug)]
pub struct DocBuilder {
    /// The project root
    root: PathBuf,
    /// Path to Solidity source files.
    sources: PathBuf,
    /// Paths to external libraries.
    libraries: Vec<PathBuf>,
    /// Flag whether to build mdbook.
    should_build: bool,
    /// Documentation configuration.
    config: DocConfig,
    /// The array of preprocessors to apply.
    preprocessors: Vec<Box<dyn Preprocessor>>,
    /// The formatter config.
    fmt: FormatterConfig,
    /// Whether to include libraries to the output.
    include_libraries: bool,
}

impl DocBuilder {
    pub(crate) const SRC: &'static str = "src";
    const SOL_EXT: &'static str = "sol";
    const README: &'static str = "README.md";
    const SUMMARY: &'static str = "SUMMARY.md";

    /// Create new instance of builder.
    pub fn new(
        root: PathBuf,
        sources: PathBuf,
        libraries: Vec<PathBuf>,
        include_libraries: bool,
    ) -> Self {
        Self {
            root,
            sources,
            libraries,
            include_libraries,
            should_build: false,
            config: DocConfig::default(),
            preprocessors: Default::default(),
            fmt: Default::default(),
        }
    }

    /// Set `should_build` flag on the builder
    pub const fn with_should_build(mut self, should_build: bool) -> Self {
        self.should_build = should_build;
        self
    }

    /// Set config on the builder.
    pub fn with_config(mut self, config: DocConfig) -> Self {
        self.config = config;
        self
    }

    /// Set formatter config on the builder.
    pub fn with_fmt(mut self, fmt: FormatterConfig) -> Self {
        self.fmt = fmt;
        self
    }

    /// Set preprocessors on the builder.
    pub fn with_preprocessor<P: Preprocessor + 'static>(mut self, preprocessor: P) -> Self {
        self.preprocessors.push(Box::new(preprocessor) as Box<dyn Preprocessor>);
        self
    }

    /// Get the output directory
    pub fn out_dir(&self) -> Result<PathBuf> {
        Ok(self.root.join(&self.config.out).canonicalize()?)
    }

    /// Parse the sources and build the documentation.
    pub fn build(self, compiler: &mut solar::sema::Compiler) -> eyre::Result<()> {
        fs::create_dir_all(self.root.join(&self.config.out))
            .wrap_err("failed to create output directory")?;

        // Expand ignore globs
        let ignored = expand_globs(&self.root, self.config.ignore.iter())?;

        // Collect and parse source files
        let sources = source_files_iter(&self.sources, SOLC_EXTENSIONS)
            .filter(|file| !ignored.contains(file))
            .collect::<Vec<_>>();

        if sources.is_empty() {
            sh_println!("No sources detected at {}", self.sources.display())?;
            return Ok(());
        }

        let library_sources = self
            .libraries
            .iter()
            .flat_map(|lib| source_files_iter(lib, SOLC_EXTENSIONS))
            .collect::<Vec<_>>();

        let combined_sources = sources
            .iter()
            .map(|path| (path, false))
            .chain(library_sources.iter().map(|path| (path, true)))
            .collect::<Vec<_>>();

        let out_dir = self.out_dir()?;
        let out_target_dir = out_dir.clone();
        let documents = compiler.enter_mut(|compiler| -> eyre::Result<Vec<Vec<Document>>> {
            let gcx = compiler.gcx();
            let documents = combined_sources
                .par_iter()
                .map(|(path, from_library)| {
                    let path = *path;
                    let from_library = *from_library;
                    let mut files = vec![];

                    // Read and parse source file
                    if let Some((_, ast_source)) = gcx.get_ast_source(path)
                        && let Some(source_unit) = ast_source.ast.as_ref()
                    {
                        // Solar uses a global SourceMap: span BytePos values are global
                        // offsets, not per-file offsets. Subtract file.start_pos so that
                        // span-based indexing into the per-file source string is correct.
                        let source = ast_source.file.src.to_string();
                        let file_start = ast_source.file.start_pos.to_usize();

                        // Walk the solar AST directly
                        let doc = Parser::new(source, file_start, self.fmt.tab_width);
                        let all_items = doc.parse(source_unit);

                        // Split the parsed items on top-level constants and rest.
                        let (items, consts): (Vec<ParseItem>, Vec<ParseItem>) = all_items
                            .into_iter()
                            .partition(|item| !matches!(item.source, ParseSource::Variable(_)));

                        // Attempt to group overloaded top-level functions
                        let mut remaining = Vec::with_capacity(items.len());
                        let mut funcs: HashMap<String, Vec<ParseItem>> = HashMap::default();
                        for item in items {
                            if matches!(item.source, ParseSource::Function(_)) {
                                funcs.entry(item.source.ident()).or_default().push(item);
                            } else {
                                // Put the item back
                                remaining.push(item);
                            }
                        }
                        let (items, overloaded): (
                            HashMap<String, Vec<ParseItem>>,
                            HashMap<String, Vec<ParseItem>>,
                        ) = funcs.into_iter().partition(|(_, v)| v.len() == 1);
                        remaining.extend(items.into_values().flatten());

                        // Each regular item will be written into its own file.
                        files = remaining
                            .into_iter()
                            .map(|item| {
                                let relative_path =
                                    path.strip_prefix(&self.root)?.join(item.filename());

                                let target_path = out_dir.join(Self::SRC).join(relative_path);
                                let ident = item.source.ident();
                                Ok(Document::new(
                                    path.clone(),
                                    target_path,
                                    from_library,
                                    out_target_dir.clone(),
                                )
                                .with_content(DocumentContent::Single(item), ident))
                            })
                            .collect::<eyre::Result<Vec<_>>>()?;

                        // If top-level constants exist, they will be written to the same file.
                        if !consts.is_empty() {
                            let filestem = path.file_stem().and_then(|stem| stem.to_str());

                            let filename = {
                                let mut name = "constants".to_owned();
                                if let Some(stem) = filestem {
                                    name.push_str(&format!(".{stem}"));
                                }
                                name.push_str(".md");
                                name
                            };
                            let relative_path = path.strip_prefix(&self.root)?.join(filename);
                            let target_path = out_dir.join(Self::SRC).join(relative_path);

                            let identity = match filestem {
                                Some(stem) if stem.to_lowercase().contains("constants") => {
                                    stem.to_owned()
                                }
                                Some(stem) => format!("{stem} constants"),
                                None => "constants".to_owned(),
                            };

                            files.push(
                                Document::new(
                                    path.clone(),
                                    target_path,
                                    from_library,
                                    out_target_dir.clone(),
                                )
                                .with_content(DocumentContent::Constants(consts), identity),
                            )
                        }

                        // If overloaded functions exist, they will be written to the same file
                        if !overloaded.is_empty() {
                            for (ident, funcs) in overloaded {
                                let filename =
                                    funcs.first().expect("no overloaded functions").filename();
                                let relative_path = path.strip_prefix(&self.root)?.join(filename);

                                let target_path = out_dir.join(Self::SRC).join(relative_path);
                                files.push(
                                    Document::new(
                                        path.clone(),
                                        target_path,
                                        from_library,
                                        out_target_dir.clone(),
                                    )
                                    .with_content(
                                        DocumentContent::OverloadedFunctions(funcs),
                                        ident,
                                    ),
                                );
                            }
                        }
                    };

                    Ok(files)
                })
                .collect::<eyre::Result<Vec<_>>>()?;

            Ok(documents)
        })?;

        // Flatten results and apply preprocessors to files
        let documents = self
            .preprocessors
            .iter()
            .try_fold(documents.into_iter().flatten().collect_vec(), |docs, p| {
                p.preprocess(docs)
            })?;

        // Sort the results and filter libraries.
        let documents = documents
            .into_iter()
            .sorted_by(|doc1, doc2| {
                doc1.item_path.display().to_string().cmp(&doc2.item_path.display().to_string())
            })
            .filter(|d| !d.from_library || self.include_libraries)
            .collect_vec();

        // Write mdbook related files
        self.write_mdbook(documents)?;

        // Build the book if requested
        if self.should_build {
            MDBook::load(self.out_dir().wrap_err("failed to construct output directory")?)
                .and_then(|book| book.build())
                .map_err(|err| eyre::eyre!("failed to build book: {err:?}"))?;
        }

        Ok(())
    }

    fn write_mdbook(&self, documents: Vec<Document>) -> eyre::Result<()> {
        let out_dir = self.out_dir().wrap_err("failed to construct output directory")?;
        let out_dir_src = out_dir.join(Self::SRC);
        fs::create_dir_all(&out_dir_src)?;

        // Write readme content if any
        let homepage_content = {
            // Default to the homepage README if it's available.
            // If not, use the src README as a fallback.
            let homepage_or_src_readme = self
                .config
                .homepage
                .as_ref()
                .map(|homepage| self.root.join(homepage))
                .unwrap_or_else(|| self.sources.join(Self::README));
            // Grab the root readme.
            let root_readme = self.root.join(Self::README);

            // Check to see if there is a 'homepage' option specified in config.
            // If not, fall back to src and root readme files, in that order.
            if homepage_or_src_readme.exists() {
                fs::read_to_string(homepage_or_src_readme)?
            } else if root_readme.exists() {
                fs::read_to_string(root_readme)?
            } else {
                String::new()
            }
        };

        let readme_path = out_dir_src.join(Self::README);
        fs::write(readme_path, homepage_content)?;

        // Write summary and section readmes
        let mut summary = BufWriter::default();
        summary.write_title("Summary")?;
        summary.write_link_list_item("Home", Self::README, 0)?;
        self.write_summary_section(&mut summary, &documents.iter().collect::<Vec<_>>(), None, 0)?;
        fs::write(out_dir_src.join(Self::SUMMARY), summary.finish())?;

        // Write solidity syntax highlighting
        fs::write(out_dir.join("solidity.min.js"), include_str!("../static/solidity.min.js"))?;

        // Write css files
        fs::write(out_dir.join("book.css"), include_str!("../static/book.css"))?;

        // Write book config
        fs::write(out_dir.join("book.toml"), self.book_config()?)?;

        // Write .gitignore
        let gitignore = "book/";
        fs::write(out_dir.join(".gitignore"), gitignore)?;

        // Write doc files
        for document in documents {
            fs::create_dir_all(
                document
                    .target_path
                    .parent()
                    .ok_or_else(|| eyre::format_err!("empty target path; noop"))?,
            )?;
            fs::write(&document.target_path, document.as_doc()?)?;
        }

        Ok(())
    }

    fn book_config(&self) -> eyre::Result<String> {
        // Read the default book first
        let mut book: value::Table = toml::from_str(include_str!("../static/book.toml"))?;
        book["book"]
            .as_table_mut()
            .unwrap()
            .insert(String::from("title"), self.config.title.clone().into());
        if let Some(ref repo) = self.config.repository {
            // Create the full repository URL.
            let git_repo_url = if let Some(path) = &self.config.path {
                // If path is specified, append it to the repository URL.
                format!("{}/{}", repo.trim_end_matches('/'), path.trim_start_matches('/'))
            } else {
                // If no path specified, use repository URL as-is.
                repo.clone()
            };

            book["output"].as_table_mut().unwrap()["html"]
                .as_table_mut()
                .unwrap()
                .insert(String::from("git-repository-url"), git_repo_url.into());
        }

        // Attempt to find the user provided book path
        let book_path = {
            if self.config.book.is_file() {
                Some(self.config.book.clone())
            } else {
                let book_path = self.config.book.join("book.toml");
                book_path.is_file().then_some(book_path)
            }
        };

        // Merge two book configs
        if let Some(book_path) = book_path {
            merge_toml_table(&mut book, toml::from_str(&fs::read_to_string(book_path)?)?);
        }

        Ok(toml::to_string_pretty(&book)?)
    }

    fn write_summary_section(
        &self,
        summary: &mut BufWriter,
        files: &[&Document],
        base_path: Option<&Path>,
        depth: usize,
    ) -> eyre::Result<()> {
        if files.is_empty() {
            return Ok(());
        }

        if let Some(path) = base_path {
            let title = path.iter().next_back().unwrap().to_string_lossy();
            if depth == 1 {
                summary.write_title(&title)?;
            } else {
                let summary_path = path.join(Self::README);
                summary.write_link_list_item(
                    &format!("❱ {title}"),
                    &summary_path.display().to_string(),
                    depth - 1,
                )?;
            }
        }

        // Group entries by path depth
        let mut grouped = HashMap::new();
        for file in files {
            let path = file.item_path.strip_prefix(&self.root)?;
            let key = path.iter().take(depth + 1).collect::<PathBuf>();
            grouped.entry(key).or_insert_with(Vec::new).push(*file);
        }
        // Sort entries by path depth
        let grouped = grouped.into_iter().sorted_by(|(lhs, _), (rhs, _)| {
            let lhs_at_end = lhs.extension().map(|ext| ext == Self::SOL_EXT).unwrap_or_default();
            let rhs_at_end = rhs.extension().map(|ext| ext == Self::SOL_EXT).unwrap_or_default();
            if lhs_at_end == rhs_at_end {
                lhs.cmp(rhs)
            } else if lhs_at_end {
                Ordering::Greater
            } else {
                Ordering::Less
            }
        });

        let out_dir = self.out_dir().wrap_err("failed to construct output directory")?;
        let mut readme = BufWriter::new("\n\n# Contents\n");
        for (path, files) in grouped {
            if path.extension().map(|ext| ext == Self::SOL_EXT).unwrap_or_default() {
                for file in files {
                    let ident = &file.identity;

                    let summary_path = &file.target_path.strip_prefix(out_dir.join(Self::SRC))?;
                    summary.write_link_list_item(
                        ident,
                        &summary_path.display().to_string(),
                        depth,
                    )?;

                    let readme_path = base_path
                        .map(|path| summary_path.strip_prefix(path))
                        .transpose()?
                        .unwrap_or(summary_path);
                    readme.write_link_list_item(ident, &readme_path.display().to_string(), 0)?;
                }
            } else {
                let name = path.iter().next_back().unwrap().to_string_lossy();
                let readme_path = Path::new("/").join(&path).display().to_string();
                readme.write_link_list_item(&name, &readme_path, 0)?;
                self.write_summary_section(summary, &files, Some(&path), depth + 1)?;
            }
        }
        if !readme.is_empty()
            && let Some(path) = base_path
        {
            let path = out_dir.join(Self::SRC).join(path);
            fs::create_dir_all(&path)?;
            fs::write(path.join(Self::README), readme.finish())?;
        }
        Ok(())
    }
}