chroma-syntaxis/src/lib.rs

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

mod languages;

#[cfg(feature = "built-in")]
pub use languages::built_in;
pub use languages::register_language;

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("\">");
}

pub fn highlight(lang: &str, source: &str, highlight_names: &[&str]) -> String {
    let highlight_config = match languages::get_highlight_config(lang, highlight_names) {
        Some(conf) => conf,
        None => return source.to_string(),
    };

    let source_bytes = source.as_bytes();

    let mut highlighter = Highlighter::new();
    let items = highlighter
        .highlight(&highlight_config, source_bytes, None, |injected_lang| {
            languages::get_highlight_config(injected_lang, highlight_names)
                .map(|x| /* ugh */ &*Box::leak(Box::new(x)))
        })
        .unwrap();

    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 item in items {
        match item.unwrap() {
            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
}

#[cfg(test)]
#[cfg(feature = "built-in")]
mod tests {
    use crate::*;

    #[test]
    fn highlight_rust() {
        built_in::register_builtin_languages();

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

        assert_eq!(
            highlight("rust", source, built_in::COMMON_HIGHLIGHT_NAMES),
            expected_result
        );
    }

    #[test]
    fn highlight_js() {
        built_in::register_builtin_languages();

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

        assert_eq!(
            highlight("js", source, built_in::COMMON_HIGHLIGHT_NAMES),
            expected_result
        );
    }

    #[test]
    fn highlight_python() {
        built_in::register_builtin_languages();

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

        assert_eq!(
            highlight("python", source, built_in::COMMON_HIGHLIGHT_NAMES),
            expected_result
        );
    }
}