chroma-syntaxis/src/lib.rs

use std::collections::HashMap;

use html_escape::encode_text;
use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};

pub mod languages;

fn get_hash_for_attrs(attrs: &[&str]) -> u64 {
    use std::{
        collections::hash_map::DefaultHasher,
        hash::{Hash, Hasher},
    };

    let mut hasher = DefaultHasher::new();
    attrs.hash(&mut hasher);
    hasher.finish()
}

fn write_opening_tag(out: &mut String, attrs: &[&str]) {
    out.push_str("<span class=\"");
    for (i, attr) in attrs.iter().enumerate() {
        if i != 0 {
            out.push(' ');
        }

        out.push_str(&attr.replace('.', " "));
    }
    out.push_str("\">");
}

fn build_highlighted_html<'a>(
    source: &'a str,
    events: impl Iterator<Item = HighlightEvent> + 'a,
    highlight_names: &[&str],
) -> String {
    let mut highlight_attrs: Vec<&str> = Vec::new();
    let mut out = String::new();

    // Collapse adjacent identical attribute sets
    let mut last_attrs: u64 = 0;
    let mut tag_is_open = false;

    for event in events {
        match event {
            HighlightEvent::Source { start, end } => {
                let source_section = &source[start..end];

                let attr_hash = get_hash_for_attrs(&highlight_attrs);
                if last_attrs != attr_hash && tag_is_open {
                    out.push_str("</span>");
                    tag_is_open = false;
                }

                if !highlight_attrs.is_empty() && (!tag_is_open || last_attrs != attr_hash) {
                    write_opening_tag(&mut out, &highlight_attrs);
                    tag_is_open = true;
                    last_attrs = attr_hash;
                }

                out.push_str(&encode_text(source_section));
            }

            HighlightEvent::HighlightStart(highlight) => {
                let capture_name = highlight_names[highlight.0];
                highlight_attrs.push(capture_name);
            }

            HighlightEvent::HighlightEnd => {
                highlight_attrs.pop();
            }
        }
    }

    if tag_is_open {
        out.push_str("</span>");
    }

    out
}

pub struct SyntaxHighlighter<'a> {
    languages: HashMap<String, HighlightConfiguration>,
    highlight_names: &'a [&'a str],
}

impl SyntaxHighlighter<'_> {
    pub fn new() -> Self {
        let mut highlighter = Self {
            languages: HashMap::new(),
            highlight_names: languages::COMMON_HIGHLIGHT_NAMES,
        };

        languages::register_builtin_languages(&mut highlighter);

        highlighter
    }
}

impl Default for SyntaxHighlighter<'_> {
    fn default() -> Self {
        Self::new()
    }
}

impl<'a> SyntaxHighlighter<'a> {
    pub fn new_empty(highlight_names: &'a [&'a str]) -> Self {
        Self {
            languages: HashMap::new(),
            highlight_names,
        }
    }

    pub fn register(&mut self, lang: String, config: HighlightConfiguration) {
        self.languages.insert(lang, config);
    }

    pub fn highlight(&self, lang: &str, source: &str) -> String {
        let highlight_config = match self.languages.get(lang) {
            Some(config) => config,
            None => return source.to_string(),
        };

        let source_bytes = source.as_bytes();
        let mut highlighter = Highlighter::new();

        let highlight_result =
            highlighter.highlight(highlight_config, source_bytes, None, |injected_lang| {
                self.languages.get(injected_lang)
            });

        let events = match highlight_result {
            Ok(events) => events,
            Err(_) => return source.to_string(),
        }
        .filter_map(|e| e.ok());

        build_highlighted_html(source, events, self.highlight_names)
    }
}

#[cfg(test)]
mod tests {
    use crate::*;

    #[test]
    fn highlight_rust() {
        let highlighter = SyntaxHighlighter::new();

        let source = include_str!("../samples/fizzbuzz.rs");
        let expected_result = include_str!("../samples/fizzbuzz.rs.html");

        assert_eq!(highlighter.highlight("rust", source), expected_result);
    }

    #[test]
    fn highlight_js() {
        let highlighter = SyntaxHighlighter::new();

        let source = include_str!("../samples/fizzbuzz.js");
        let expected_result = include_str!("../samples/fizzbuzz.js.html");

        assert_eq!(highlighter.highlight("js", source), expected_result);
    }

    #[test]
    fn highlight_python() {
        let highlighter = SyntaxHighlighter::new();

        let source = include_str!("../samples/fizzbuzz.py");
        let expected_result = include_str!("../samples/fizzbuzz.py.html");

        assert_eq!(highlighter.highlight("python", source), expected_result);
    }
}