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::{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.
                            // Bump the generation to cancel any previously running jobs.
                            let token = scheduler.cancel_and_start();

                            tokio::task::spawn_blocking(move || {
                                let local = SyntaxHighlighter::new().with_theme(&theme).with_syntax(&syntax);
                                // Use the cancellable highlighter variant which checks the
                                // cancellation token between lines.
                                let tokens = local.highlight_tokens_cancellable(&lines, version, start_line, count, || token.is_cancelled());
                                // Only send results if this job was not cancelled.
                                if !token.is_cancelled()
                                    && reply.send(tokens).is_err() {
                                        log::debug!("Highlight request receiver dropped");
                                    }
                            });
                        }
                }
            }
        }
    }

    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;
    }
}

/// Lock-free cancellation token based on a generation counter.
///
/// Each token captures the generation at creation time. When the scheduler
/// advances the generation (via [`HighlightScheduler::cancel_and_start`]),
/// all previously issued tokens become cancelled. No `Mutex`, `Vec`, or
/// per-job registration is needed.
#[derive(Clone)]
pub struct CancellationToken {
    generation: u64,
    current: Arc<AtomicU64>,
}

impl CancellationToken {
    pub fn is_cancelled(&self) -> bool {
        // Relaxed is sufficient: the generation counter is the sole piece of
        // shared state and no other data needs to be ordered relative to it.
        self.generation != self.current.load(Ordering::Relaxed)
    }
}

/// Generation-based highlight job scheduler.
///
/// Instead of maintaining a `Mutex<Vec<Arc<AtomicBool>>>` of active jobs,
/// cancellation is achieved by atomically incrementing a generation counter.
/// Each [`CancellationToken`] compares its captured generation against the
/// current value — if they differ, the job is stale and should abort.
///
/// This eliminates lock contention, `retain`-based removal, and all per-job
/// registration overhead.
pub struct HighlightScheduler {
    generation: Arc<AtomicU64>,
}

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

impl HighlightScheduler {
    pub fn new() -> Self {
        Self {
            generation: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Cancels all previously issued tokens and returns a fresh
    /// [`CancellationToken`] for the new job.
    pub fn cancel_and_start(&self) -> CancellationToken {
        let new_gen = self.generation.fetch_add(1, Ordering::Relaxed) + 1;
        CancellationToken {
            generation: new_gen,
            current: self.generation.clone(),
        }
    }

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

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

    #[test]
    fn cancellation_token_starts_valid() {
        let scheduler = HighlightScheduler::new();
        let token = scheduler.cancel_and_start();
        assert!(!token.is_cancelled());
    }

    #[test]
    fn cancellation_token_invalidated_by_new_generation() {
        let scheduler = HighlightScheduler::new();
        let old_token = scheduler.cancel_and_start();
        let new_token = scheduler.cancel_and_start();

        assert!(old_token.is_cancelled(), "old token must be cancelled");
        assert!(!new_token.is_cancelled(), "new token must still be valid");
    }

    #[test]
    fn multiple_generations_cancel_all_prior_tokens() {
        let scheduler = HighlightScheduler::new();
        let t1 = scheduler.cancel_and_start();
        let t2 = scheduler.cancel_and_start();
        let t3 = scheduler.cancel_and_start();

        assert!(t1.is_cancelled());
        assert!(t2.is_cancelled());
        assert!(!t3.is_cancelled());
    }

    #[test]
    fn scheduler_generation_advances() {
        let scheduler = HighlightScheduler::new();
        let gen0 = scheduler.current_generation();
        let _ = scheduler.cancel_and_start();
        let gen1 = scheduler.current_generation();
        assert!(gen1 > gen0);
    }

    #[test]
    fn cloned_token_shares_cancellation_state() {
        let scheduler = HighlightScheduler::new();
        let token = scheduler.cancel_and_start();
        let cloned = token.clone();

        assert!(!cloned.is_cancelled());
        let _ = scheduler.cancel_and_start();
        assert!(cloned.is_cancelled());
        assert!(token.is_cancelled());
    }
}