karnaugh/src/syntax.rs

use std::{
    borrow::Cow,
    collections::{HashMap, HashSet},
    error::Error,
    fs,
    path::Path,
    sync::OnceLock,
    time::Duration,
};

use ropey::Rope;
use tracing::{info, warn};
use tree_house::{
    InjectionLanguageMarker, Language, LanguageConfig, LanguageLoader, Syntax,
    highlighter::{Highlight, HighlightEvent, Highlighter},
    read_query,
};
use tree_house_bindings::Grammar;
use typst::{
    foundations::{Content, Func, NativeElement, NativeFunc, func},
    text::TextElem,
};
use typst_html::{HtmlAttr, HtmlElem, attr, tag};

use crate::AppConfig;

static REGISTRY: OnceLock<LanguageRegistry> = OnceLock::new();

/// Class names used in highlighted code.
const HIGHLIGHT_NAMES: &[&str] = &[
    "attribute",
    "comment",
    "comment.block",
    "comment.line",
    "constant",
    "constant.builtin",
    "constructor",
    "function",
    "function.builtin",
    "function.macro",
    "function.method",
    "keyword",
    "keyword.control",
    "keyword.operator",
    "label",
    "namespace",
    "number",
    "operator",
    "property",
    "punctuation",
    "punctuation.bracket",
    "punctuation.delimiter",
    "special",
    "string",
    "string.special",
    "tag",
    "type",
    "type.builtin",
    "variable",
    "variable.builtin",
    "variable.parameter",
];

/// Get the longest matching highlight name for a given scope.
fn highlight_for_scope(scope: &str) -> Option<Highlight> {
    let mut candidate = scope;
    loop {
        if let Some(i) = HIGHLIGHT_NAMES.iter().position(|&n| n == candidate) {
            return Some(Highlight::new(i as u32));
        }
        match candidate.rfind('.') {
            Some(dot) => candidate = &candidate[..dot],
            None => return None,
        }
    }
}

struct LanguageRegistry {
    by_name: HashMap<String, Language>,
    configs: Vec<LanguageConfig>,
}

impl LanguageRegistry {
    /// Loads grammars from `syntax_root`. When `languages` is non-empty, only
    /// grammars with those names are loaded; otherwise every grammar found is.
    fn new(syntax_root: &Path, languages: &[String]) -> Self {
        let filter: Option<HashSet<&str>> =
            (!languages.is_empty()).then(|| languages.iter().map(String::as_str).collect());

        let mut by_name = HashMap::new();
        let mut configs = Vec::new();

        for entry in fs::read_dir(syntax_root).into_iter().flatten().flatten() {
            let name = entry.file_name().to_string_lossy().into_owned();
            if filter.as_ref().is_some_and(|f| !f.contains(name.as_str())) {
                continue;
            }
            let so_path = entry.path().join(format!("{name}.so"));
            if !so_path.exists() {
                continue;
            }

            match load_language(&name, &entry.path(), syntax_root) {
                Ok(config) => {
                    by_name.insert(name, Language::new(configs.len() as u32));
                    configs.push(config);
                }
                Err(e) => warn!("Failed to load grammar {name}: {e}"),
            }
        }

        // Warn about explicitly requested grammars that were not found.
        for lang in languages {
            if !by_name.contains_key(lang) {
                warn!("Requested grammar {lang} was not found in the syntax root");
            }
        }

        info!("Loaded {} syntax highlighting grammars", configs.len());
        Self { by_name, configs }
    }
}

fn load_language(
    name: &str,
    lang_dir: &Path,
    syntax_root: &Path,
) -> Result<LanguageConfig, Box<dyn Error>> {
    let grammar = unsafe { Grammar::new(name, &lang_dir.join(format!("{name}.so")))? };

    let read = |kind: &str| {
        read_query(name, |lang| {
            fs::read_to_string(syntax_root.join(lang).join(format!("{kind}.scm")))
                .unwrap_or_default()
        })
    };

    let config = LanguageConfig::new(
        grammar,
        &read("highlights"),
        &read("injections"),
        &read("locals"),
    )?;
    config.configure(highlight_for_scope);

    Ok(config)
}

impl LanguageLoader for LanguageRegistry {
    fn language_for_marker(&self, marker: InjectionLanguageMarker<'_>) -> Option<Language> {
        match marker {
            InjectionLanguageMarker::Name(name) => self.by_name.get(name).copied(),
            InjectionLanguageMarker::Match(text) => {
                let name: Cow<str> = text.into();
                self.by_name.get(name.as_ref()).copied()
            }
            _ => None,
        }
    }

    fn get_config(&self, lang: Language) -> Option<&LanguageConfig> {
        self.configs.get(lang.idx())
    }
}

pub fn init(config: &AppConfig) {
    match &config.syntax_root {
        Some(syntax_root) => {
            info!("Loading syntax files from {}", syntax_root.display());
            REGISTRY.get_or_init(|| LanguageRegistry::new(syntax_root, &config.languages));
        }
        None => {
            warn!(
                "No syntax root has been configured, improved tree-sitter highlighting will not be possible."
            )
        }
    }
}

pub fn highlight_func() -> Func {
    highlight::func()
}

#[func]
fn highlight(
    text: String,
    #[named]
    #[default(false)]
    block: bool,
    #[named]
    #[default(None)]
    lang: Option<String>,
) -> Content {
    let highlighted = highlight_code(&text, lang.as_deref()).unwrap_or(TextElem::packed(text));

    if block {
        HtmlElem::new(tag::pre)
            .with_body(Some(highlighted))
            .with_attr(HtmlAttr::constant("class"), "code-block")
            .pack()
    } else {
        HtmlElem::new(tag::code)
            .with_body(Some(highlighted))
            .with_attr(HtmlAttr::constant("class"), "code-inline")
            .pack()
    }
}

/// Highlights `source` as `lang`.
fn highlight_code(source: &str, lang: Option<&str>) -> Option<Content> {
    let registry = REGISTRY.get()?;
    let language = *registry.by_name.get(lang?)?;

    let rope = Rope::from_str(source);
    let slice = rope.slice(..);
    let syntax = Syntax::new(slice, language, Duration::from_millis(500), registry).ok()?;
    let mut hl = Highlighter::new(&syntax, slice, registry, 0..source.len() as u32);

    let src = source.as_bytes();
    let src_len = source.len() as u32;
    let mut pos: u32 = 0;
    let mut classes = String::new();
    let mut nodes = Vec::<Content>::new();

    loop {
        let next = hl.next_event_offset().min(src_len);

        if next > pos {
            let text = std::str::from_utf8(&src[pos as usize..next as usize]).ok()?;
            let node = TextElem::packed(text);
            nodes.push(if classes.is_empty() {
                node
            } else {
                HtmlElem::new(tag::span)
                    .with_attr(attr::class, classes.clone())
                    .with_body(Some(node))
                    .pack()
            });
            pos = next;
        }

        if pos >= src_len {
            break;
        }

        let (event, highlights) = hl.advance();
        let base = match event {
            // Replace the active highlights
            HighlightEvent::Refresh => String::new(),
            // Extend the active highlights
            HighlightEvent::Push => classes.clone(),
        };
        classes = highlights.fold(base, |mut acc, h| {
            let name = HIGHLIGHT_NAMES[h.idx()].replace('.', "-");
            if !acc.is_empty() {
                acc.push(' ');
            }
            acc.push_str("hl-");
            acc.push_str(&name);
            acc
        });
    }

    Some(Content::sequence(nodes))
}