oo_ide/editor/

highlight_worker.rs

//! Async syntax highlighting worker.
//!
//! Provides background highlighting for large files or slow syntaxes.
//! The main SyntaxHighlighter handles checkpointing and caching.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::{mpsc, oneshot};

use crate::editor::buffer::Version;
use crate::editor::highlight::{SyntaxHighlighter, Token};

pub struct HighlightWorker {
    tx: mpsc::Sender<HighlightRequest>,
    handle: tokio::task::JoinHandle<()>,
}

#[derive(Debug)]
pub struct HighlightRequest {
    pub version: Version,
    pub start_line: usize,
    pub count: usize,
    /// Shared buffer lines to avoid cloning the Vec for every request.
    pub lines: Arc<Vec<String>>,
    pub reply: oneshot::Sender<Vec<Arc<Vec<Token>>>>,
}

impl HighlightWorker {
    pub fn new(highlighter: SyntaxHighlighter) -> Self {
        let (tx, rx) = mpsc::channel::<HighlightRequest>(32);

        let handle = tokio::spawn(async move {
            Self::run_worker(rx, highlighter).await;
        });

        Self { tx, handle }
    }

    async fn run_worker(
        mut rx: mpsc::Receiver<HighlightRequest>,
        highlighter: SyntaxHighlighter,
    ) {
        const DEBOUNCE_MS: u64 = 50;

        let mut pending: Option<HighlightRequest> = None;
        let mut last_request_time = std::time::Instant::now();

        // Scheduler used to track and cancel active highlight jobs.
        let scheduler = Arc::new(HighlightScheduler::new());

        loop {
            tokio::select! {
                biased;

                request = rx.recv() => {
                    match request {
                        None => break,
                        Some(req) => {
                            pending = Some(req);
                            last_request_time = std::time::Instant::now();
                        }
                    }
                }

                _ = tokio::time::sleep(Duration::from_millis(DEBOUNCE_MS)), if pending.is_some() => {
                    if last_request_time.elapsed() >= Duration::from_millis(DEBOUNCE_MS)
                        && let Some(req) = pending.take() {
                                    // Offload CPU-heavy highlighting to the blocking thread pool
                            // to avoid blocking the Tokio runtime.
                            let theme = highlighter.theme_name.clone();
                            let syntax = highlighter.syntax_name.clone();
                            let lines = req.lines;
                            let version = req.version;
                            let start_line = req.start_line;
                            let count = req.count;
                            let reply = req.reply;

                            // Spawn a blocking task that constructs a local
                            // SyntaxHighlighter and performs the highlighting. The
                            // blocking task sends the result back on the oneshot
                            // channel when complete.
                            // Cancel any previously running jobs before starting this one.
                            let job = HighlightJob::new(version, start_line, count);
                            let cancelled = job.cancelled.clone();
                            scheduler.cancel_all_jobs();
                            scheduler.register_job(cancelled.clone());

                            let scheduler_clone = scheduler.clone();
                            tokio::task::spawn_blocking(move || {
                                let local = SyntaxHighlighter::new().with_theme(&theme).with_syntax(&syntax);
                                // Use the cancellable highlighter variant which checks the
                                // cancellation flag between lines.
                                let tokens = local.highlight_tokens_cancellable(&lines, version, start_line, count, &cancelled);
                                // Only send results if this job was not cancelled.
                                if !cancelled.load(std::sync::atomic::Ordering::SeqCst)
                                    && reply.send(tokens).is_err() {
                                        log::debug!("Highlight request receiver dropped");
                                    }
                                // Unregister job so scheduler stops tracking it.
                                scheduler_clone.unregister_job(&cancelled);
                            });
                        }
                }
            }
        }
    }

    pub async fn highlight(
        &self,
        version: Version,
        start_line: usize,
        count: usize,
        lines: Arc<Vec<String>>,
    ) -> Vec<Arc<Vec<Token>>> {
        let (reply, received) = oneshot::channel();

        let request = HighlightRequest {
            version,
            start_line,
            count,
            lines,
            reply,
        };

        if self.tx.send(request).await.is_err() {
            log::warn!("Highlight worker channel closed");
            return Vec::new();
        }

        received.await.unwrap_or_default()
    }

    pub async fn shutdown(self) {
        drop(self.tx);
        let _ = self.handle.await;
    }
}

#[derive(Debug, Clone)]
pub struct HighlightJob {
    pub version: Version,
    pub start_line: usize,
    pub count: usize,
    pub cancelled: Arc<AtomicBool>,
}

impl HighlightJob {
    pub fn new(version: Version, start_line: usize, count: usize) -> Self {
        Self {
            version,
            start_line,
            count,
            cancelled: Arc::new(AtomicBool::new(false)),
        }
    }

    pub fn cancel(&self) {
        self.cancelled.store(true, Ordering::SeqCst);
    }

    pub fn is_cancelled(&self) -> bool {
        self.cancelled.load(Ordering::SeqCst)
    }
}

pub struct HighlightScheduler {
    generation: Arc<AtomicU64>,
    active_jobs: std::sync::Mutex<Vec<Arc<AtomicBool>>>,
}

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

impl HighlightScheduler {
    pub fn new() -> Self {
        Self {
            generation: Arc::new(AtomicU64::new(0)),
            active_jobs: std::sync::Mutex::new(Vec::new()),
        }
    }

    pub fn next_generation(&self) -> u64 {
        self.generation.fetch_add(1, Ordering::SeqCst) + 1
    }

    pub fn current_generation(&self) -> u64 {
        self.generation.load(Ordering::SeqCst)
    }

    pub fn cancel_all_jobs(&self) {
        let jobs = self.active_jobs.lock().unwrap();
        for job in jobs.iter() {
            job.store(true, Ordering::SeqCst);
        }
    }

    pub fn register_job(&self, job: Arc<AtomicBool>) {
        let mut jobs = self.active_jobs.lock().unwrap();
        jobs.push(job);
        if jobs.len() > 100 {
            jobs.remove(0);
        }
    }

    pub fn unregister_job(&self, job: &Arc<AtomicBool>) {
        let mut jobs = self.active_jobs.lock().unwrap();
        jobs.retain(|j| !j.as_ptr().eq(&job.as_ptr()));
    }
}

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

    #[test]
    fn job_cancellation() {
        let job = HighlightJob::new(Version::new(), 0, 10);
        assert!(!job.is_cancelled());
        
        job.cancel();
        assert!(job.is_cancelled());
    }

    #[test]
    fn scheduler_generation() {
        let scheduler = HighlightScheduler::new();
        let gen1 = scheduler.current_generation();
        let gen2 = scheduler.next_generation();
        assert!(gen2 > gen1);
    }
}