treehouse/src/state.rs

use std::{collections::HashMap, ops::Range};

use anyhow::Context;
use codespan_reporting::{
    diagnostic::{Diagnostic, Label, LabelStyle, Severity},
    term::termcolor::{ColorChoice, StandardStream},
};
use tracing::instrument;
use ulid::Ulid;

use crate::{
    doc::Doc,
    tree::{SemaBranchId, SemaRoots, SemaTree},
    vfs::{VPath, VPathBuf},
};

#[derive(Debug, Clone)]
pub enum Source {
    Tree { input: String, tree_path: VPathBuf },
    Other(String),
}

impl Source {
    pub fn input(&self) -> &str {
        match &self {
            Source::Tree { input, .. } => input,
            Source::Other(source) => source,
        }
    }
}

impl AsRef<str> for Source {
    fn as_ref(&self) -> &str {
        self.input()
    }
}

#[derive(Debug, Clone)]
pub struct File {
    pub path: VPathBuf,
    pub source: Source,
    pub line_starts: Vec<usize>,
}

impl File {
    fn line_start(&self, line_index: usize) -> Result<usize, codespan_reporting::files::Error> {
        use std::cmp::Ordering;

        match line_index.cmp(&self.line_starts.len()) {
            Ordering::Less => Ok(self
                .line_starts
                .get(line_index)
                .cloned()
                .expect("failed despite previous check")),
            Ordering::Equal => Ok(self.source.as_ref().len()),
            Ordering::Greater => Err(codespan_reporting::files::Error::LineTooLarge {
                given: line_index,
                max: self.line_starts.len() - 1,
            }),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct FileId(usize);

#[derive(Debug, Clone, Default)]
pub struct Tag {
    pub files: Vec<FileId>,
}

/// Treehouse compilation context.
pub struct Treehouse {
    pub files: Vec<File>,
    pub files_by_tree_path: HashMap<VPathBuf, FileId>, // trees only
    pub files_by_doc_path: HashMap<VPathBuf, FileId>,  // docs only
    pub tags: HashMap<String, Tag>,

    pub tree: SemaTree,
    pub branches_by_named_id: HashMap<String, SemaBranchId>,
    pub roots: HashMap<FileId, SemaRoots>,

    pub docs: HashMap<FileId, Doc>,

    pub branch_redirects: HashMap<String, SemaBranchId>,

    pub missingno_generator: ulid::Generator,
}

impl Treehouse {
    pub fn new() -> Self {
        Self {
            files: vec![],
            files_by_tree_path: HashMap::new(),
            files_by_doc_path: HashMap::new(),
            tags: HashMap::new(),

            tree: SemaTree::default(),
            branches_by_named_id: HashMap::new(),
            roots: HashMap::new(),

            docs: HashMap::new(),

            branch_redirects: HashMap::new(),

            missingno_generator: ulid::Generator::new(),
        }
    }

    pub fn add_file(&mut self, path: VPathBuf, source: Source) -> FileId {
        let id = FileId(self.files.len());
        self.files.push(File {
            line_starts: codespan_reporting::files::line_starts(source.input()).collect(),

            path,
            source,
        });
        id
    }

    /// Get the name of a file, assuming it was previously registered.
    pub fn path(&self, file_id: FileId) -> &VPath {
        &self.files[file_id.0].path
    }

    /// Get the source code of a file, assuming it was previously registered.
    pub fn source(&self, file_id: FileId) -> &Source {
        &self.files[file_id.0].source
    }

    pub fn set_source(&mut self, file_id: FileId, source: Source) {
        self.files[file_id.0].line_starts =
            codespan_reporting::files::line_starts(source.input()).collect();
        self.files[file_id.0].source = source;
    }

    pub fn tree_path(&self, file_id: FileId) -> Option<&VPath> {
        match self.source(file_id) {
            Source::Tree { tree_path, .. } => Some(tree_path),
            Source::Other(_) => None,
        }
    }

    pub fn next_missingno(&mut self) -> Ulid {
        self.missingno_generator
            .generate()
            .expect("just how much disk space do you have?")
    }
}

impl Default for Treehouse {
    fn default() -> Self {
        Self::new()
    }
}

impl<'a> codespan_reporting::files::Files<'a> for Treehouse {
    type FileId = FileId;

    type Name = &'a VPath;

    type Source = &'a str;

    fn name(&'a self, id: Self::FileId) -> Result<Self::Name, codespan_reporting::files::Error> {
        Ok(self.path(id))
    }

    fn source(
        &'a self,
        id: Self::FileId,
    ) -> Result<Self::Source, codespan_reporting::files::Error> {
        Ok(self.source(id).input())
    }

    fn line_index(
        &'a self,
        id: Self::FileId,
        byte_index: usize,
    ) -> Result<usize, codespan_reporting::files::Error> {
        let file = &self.files[id.0];
        Ok(file
            .line_starts
            .binary_search(&byte_index)
            .unwrap_or_else(|next_line| next_line - 1))
    }

    fn line_range(
        &'a self,
        id: Self::FileId,
        line_index: usize,
    ) -> Result<Range<usize>, codespan_reporting::files::Error> {
        let file = &self.files[id.0];
        let line_start = file.line_start(line_index)?;
        let next_line_start = file.line_start(line_index + 1)?;

        Ok(line_start..next_line_start)
    }
}

pub struct TomlError {
    pub message: String,
    pub span: Option<Range<usize>>,
    pub file_id: FileId,
    pub input_range: Range<usize>,
}

pub fn toml_error_to_diagnostic(error: TomlError) -> Diagnostic<FileId> {
    Diagnostic {
        severity: Severity::Error,
        code: Some("toml".into()),
        message: error.message,
        labels: error
            .span
            .map(|span| Label {
                style: LabelStyle::Primary,
                file_id: error.file_id,
                range: error.input_range.start + span.start..error.input_range.start + span.end,
                message: String::new(),
            })
            .into_iter()
            .collect(),
        notes: vec![],
    }
}

#[instrument(skip(files, diagnostics))]
pub fn report_diagnostics(
    files: &Treehouse,
    diagnostics: &[Diagnostic<FileId>],
) -> anyhow::Result<()> {
    let writer = StandardStream::stderr(ColorChoice::Auto);
    let config = codespan_reporting::term::Config::default();
    for diagnostic in diagnostics {
        codespan_reporting::term::emit(&mut writer.lock(), &config, files, diagnostic)
            .context("could not emit diagnostic")?;
    }

    Ok(())
}

pub fn has_errors(diagnostics: &[Diagnostic<FileId>]) -> bool {
    diagnostics.iter().any(|d| d.severity == Severity::Error)
}