diff --git a/benches/worker.rs b/benches/worker.rs index a87e3f3..64c7803 100644 --- a/benches/worker.rs +++ b/benches/worker.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use criterion::{criterion_group, criterion_main, Criterion}; use tree_sitter_highlight::{HighlightConfiguration, Highlighter}; -use kak_highlight::daemon::worker::highlight_content; +use kak_highlight::daemon::worker::{highlight_content, make_highlighter_config}; fn make_tokens() -> HashMap { [ @@ -33,15 +33,7 @@ fn run_benchmark(c: &mut Criterion, name: &str, content: &str) { let names: Vec<_> = tokens.keys().collect(); let mut highlighter = Highlighter::new(); - - let mut highlight_config = HighlightConfiguration::new( - tree_sitter_rust::language(), - tree_sitter_rust::HIGHLIGHT_QUERY, - "", - tree_sitter_rust::LOCALS_QUERY, - ) - .unwrap(); - highlight_config.configure(&names); + let highlight_config = make_highlighter_config(&names).unwrap(); c.bench_function(name, |b| { b.iter(|| { diff --git a/src/daemon/worker.rs b/src/daemon/worker.rs index 530d543..62bcfae 100644 --- a/src/daemon/worker.rs +++ b/src/daemon/worker.rs @@ -34,6 +34,22 @@ pub struct TaskContext { tokens: Arc>, } +pub fn make_highlighter_config>( + names: impl AsRef<[S]>, +) -> Result { + let mut config = HighlightConfiguration::new( + tree_sitter_rust::language(), + tree_sitter_rust::HIGHLIGHT_QUERY, + tree_sitter_rust::INJECTIONS_QUERY, + tree_sitter_rust::LOCALS_QUERY, + ) + .wrap_err("Invalid highlighter config")?; + + config.configure(names.as_ref()); + + Ok(config) +} + impl TaskScheduler { pub fn new(workers: u8, tokens: HashMap) -> Result { let terminate = Arc::new(AtomicBool::new(false)); @@ -43,18 +59,10 @@ impl TaskScheduler { let stealers: Vec<_> = workers.iter().map(|w| w.stealer()).collect(); let stealers = Arc::new(stealers); - let mut highlight_config = HighlightConfiguration::new( - tree_sitter_rust::language(), - tree_sitter_rust::HIGHLIGHT_QUERY, - "", - tree_sitter_rust::LOCALS_QUERY, - ) - .wrap_err("Invalid highlighter config")?; - let names: Vec<_> = tokens.keys().collect(); tracing::debug!("Highlighter tokens: {:?}", names); - highlight_config.configure(&names); + let highlight_config = make_highlighter_config(names)?; let highlight_config = Arc::new(highlight_config); let tokens = Arc::new(tokens); @@ -290,7 +298,9 @@ pub fn highlight_content( mod test { use std::collections::HashMap; - use tree_sitter_highlight::{HighlightConfiguration, Highlighter}; + use tree_sitter_highlight::Highlighter; + + use crate::daemon::worker::make_highlighter_config; use super::highlight_content; @@ -322,15 +332,7 @@ mod test { let names: Vec<_> = tokens.keys().collect(); let mut highlighter = Highlighter::new(); - - let mut highlight_config = HighlightConfiguration::new( - tree_sitter_rust::language(), - tree_sitter_rust::HIGHLIGHT_QUERY, - "", - tree_sitter_rust::LOCALS_QUERY, - ) - .unwrap(); - highlight_config.configure(&names); + let highlight_config = make_highlighter_config(&names).unwrap(); let response = highlight_content( &mut highlighter,