hcie-rust-v3.05/hcie-engine-api/src/stroke_cache.rs

//! Stroke-level cache management for the HCIE engine.
//!
//! ## Purpose
//! Holds the functions that set up and tear down a brush stroke, together with
//! the `below_cache` snapshot of all layers below the active drawing layer.
//! Keeping these in one module makes the stroke-side performance logic
//! (mask pooling, below-layer cache reuse, sub-rect undo bounds) easier to
//! protect from accidental regression.
//!
//! ## Logic & Workflow
//! - `begin_stroke()` pools `active_stroke_mask`, snapshots the layer for undo,
//!   and either reuses or rebuilds the `below_cache` composite.
//! - `end_stroke()` commits a sub-rect undo snapshot and retains the below-layer
//!   cache so the next stroke on the same active layer can reuse it.
//! - `expand_stroke_bounds()` tracks the union of all brush stamps for the
//!   sub-rect snapshot.
//!
//! ## Side Effects / Dependencies
//! All functions mutate `Engine` state (`document`, `tile_layers`, pooled
//! buffers, dirty flags). They depend on `hcie_tile::TiledLayer` and the
//! dynamic `tiled::composite_tiled_into` compositor.

use hcie_protocol::LayerData;
use hcie_tile::TiledLayer;
use crate::Engine;
use crate::dynamic_loader::tiled;

/// Wrapper to send a raw pixel pointer to a background thread as a `usize`.
///
/// # Safety Invariant
///
/// The pointer is valid for the entire lifetime of the `Engine`. The background
/// thread only reads from it, and no mutations happen to `layer.pixels` between
/// `end_stroke()` and the next `begin_stroke()`.
struct SendPtr(usize);
impl SendPtr {
    #[inline]
    fn new(p: *const u8) -> Self { Self(p as usize) }
    #[inline]
    fn as_ptr(&self) -> *const u8 { self.0 as *const u8 }
}
// SAFETY: The usize encodes a pointer to layer.pixels which is stable for the engine lifetime.
// The background thread only reads, never writes.
unsafe impl Send for SendPtr {}

pub struct PendingHistoryItem {
    pub layer_idx: usize,
    pub before_pixels: Vec<u8>,
    pub after_pixels: Vec<u8>,
    pub bounds: Option<[u32; 4]>,
    pub before_shapes: Option<Vec<hcie_protocol::VectorShape>>,
    pub after_shapes: Option<Vec<hcie_protocol::VectorShape>>,
    pub description: String,
    /// Full-layer before buffer returned from background thread for pool reuse.
    pub return_before: Option<Vec<u8>>,
    /// Full-layer after buffer returned from background thread for pool reuse.
    pub return_after: Option<Vec<u8>>,
}

impl Engine {
    /// Commit any pending history items computed in the background thread.
    /// This should be polled on the UI thread.
    ///
    /// Also recycles returned buffers back into the engine's pool to avoid
    /// allocations on subsequent strokes.
    pub fn commit_pending_history(&mut self) {
        let items = {
            let mut pending = self.pending_history.lock().unwrap();
            if pending.is_empty() {
                return;
            }
            std::mem::take(&mut *pending)
        };
        for item in items {
            // Recycle returned before buffer back into the pool
            if let Some(buf) = item.return_before {
                if self.stroke_before_buf.as_ref().map_or(true, |b| b.len() != buf.len()) {
                    self.stroke_before_buf = Some(buf);
                }
            }

            if let Some(bounds) = item.bounds {
                self.document.push_draw_snapshot_subrect(
                    item.layer_idx,
                    item.before_pixels,
                    item.after_pixels,
                    (bounds[0], bounds[1], bounds[2], bounds[3]),
                    item.description,
                );
            } else {
                // If bounds are None, perform full-layer snapshot using the cached after pixels.
                // We use document's internal layers state to push the action directly.
                self.document.push_draw_snapshot(
                    item.layer_idx,
                    item.before_pixels,
                    item.before_shapes,
                    item.description,
                );
            }
        }
    }

    /// Begin a new brush stroke on the given layer.
    ///
    /// # Allocation Behaviour
    ///
    /// This function performs two pooled buffer operations per call:
    ///
    /// 1. `active_stroke_mask` — **Pooled**: if the existing buffer matches
    ///    `layer_size`, it is zeroed in-place with `fill(0)`. Only if the
    ///    canvas dimensions changed does a fresh `vec![0u8; layer_size]`
    ///    allocation occur (~8MB on 4K). The buffer is retained across strokes.
    ///
    /// 2. `stroke_before_buf` — **Pooled**: `layer.pixels` is copied into the
    ///    existing buffer via `copy_from_slice` (no allocation after first use).
    ///    If the buffer doesn't exist or has wrong size, a fresh allocation
    ///    occurs (~33MB on 4K). The buffer is returned to the pool after the
    ///    background thread extracts the sub-rect snapshot.
    ///
    /// # Risk
    ///
    /// If `active_stroke_mask` is accidentally set to `None` at any point
    /// after `end_stroke()` (e.g. a partial merge revert), the pooling
    /// optimization is silently negated and ~8MB will be allocated per stroke.
    /// The merge artifact cleanup of 2026-05-28 removed exactly this regression.
    pub fn begin_stroke(&mut self, layer_id: u64, x: f32, y: f32) {
        self.last_stroke_pos = Some((x, y, 1.0));
        self.sketch_history.clear();
        if let Some(layer) = self.document.get_layer_by_id_mut(layer_id) {
            let layer_pixels = (layer.width * layer.height * 4) as usize;  // RGBA byte size
            let layer_size = (layer.width * layer.height) as usize;  // pixel count (for mask)

            // Pool active_stroke_mask: reuse buffer if size matches, otherwise allocate.
            // This avoids a ~16MB allocation per stroke on a 4K canvas.
            match &mut self.active_stroke_mask {
                Some(mask) if mask.len() == layer_size => {
                    // Reuse existing buffer — just zero it
                    mask.fill(0);
                }
                _ => {
                    // Size changed (different document/layer) — allocate fresh
                    self.active_stroke_mask = Some(vec![0u8; layer_size]);
                }
            }

            // Backup and clear layer.effects to disable expensive effect compositing during the stroke
            if !layer.effects.is_empty() {
                self.stroke_effects_backup = Some(layer.effects.clone());
                layer.effects.clear();
                // Clear the cache so it stops rendering the old effects immediately
                *layer.effects_cache.lock().unwrap() = None;
            }
            if !layer.styles.is_empty() {
                self.stroke_styles_backup = Some(layer.styles.clone());
                layer.styles.clear();
                *layer.effects_cache.lock().unwrap() = None;
            }

            // Initialize last_stroke_bounds with brush radius around first point
            // Used for sub-rect snapshot in end_stroke (avoid 33MB clone on 4K)
            let tip = &self.current_tip;
            let dr = ((tip.size.max(2.0) * 1.3) as i32).max(1);
            let ix = x as i32;
            let iy = y as i32;
            let x0 = (ix - dr).max(0) as u32;
            let y0 = (iy - dr).max(0) as u32;
            let x1 = ((ix + dr + 1).min(layer.width as i32 - 1)).max(0) as u32;
            let y1 = ((iy + dr + 1).min(layer.height as i32 - 1)).max(0) as u32;
            self.last_stroke_bounds = Some([x0, y0, x1, y1]);

            // Pool the "before" pixel buffer: reuse if size matches, otherwise allocate.
            // This avoids a ~33MB allocation per stroke on 4K canvases.
            match &mut self.stroke_before_buf {
                Some(buf) if buf.len() == layer_pixels => {
                    buf.copy_from_slice(&layer.pixels);
                }
                _ => {
                    self.stroke_before_buf = Some(layer.pixels.clone());
                }
            }
            let before_shapes = if let LayerData::Vector { shapes } = &layer.data {
                Some(shapes.clone())
            } else {
                None
            };
            self.stroke_before = Some((layer_id, before_shapes));

            // Cache selection mask once at stroke start to avoid ~8MB clone per stroke_to().
            self.cached_selection_mask = self.document.selection_mask.clone();

            // ── Below-layer composite cache ──────────────────────────────
            self.rebuild_below_cache_if_needed();
        }
    }

    /// End the current brush stroke and commit a sub-rect undo snapshot.
    ///
    /// Offloads the expensive sub-rect copy and vector comparison to a background
    /// thread using `std::thread::spawn` so the UI remains completely responsive and free of pauses.
    pub fn end_stroke(&mut self, layer_id: u64) {
        log::trace!(
            "[end_stroke_async] layer_id={}, stroke_before={}, below_cache={}, below_cache_active_idx={:?}, last_stroke_bounds={:?}",
            layer_id,
            self.stroke_before.is_some(),
            self.below_cache.is_some(),
            self.below_cache_active_idx,
            self.last_stroke_bounds
        );
        if let Some((id, before_shapes)) = self.stroke_before.take() {
            if id == layer_id {
                if let Some(idx) = self.document.layer_index_by_id(layer_id) {
                    let layer = &mut self.document.layers[idx];
                    let lw = layer.width;
                    let layer_pixels = (layer.width * layer.height * 4) as usize;  // RGBA byte size
                    // Restore layer effects from backup so they are re-composited correctly
                    if let Some(backup) = self.stroke_effects_backup.take() {
                        layer.effects = backup;
                    }
                    if let Some(backup) = self.stroke_styles_backup.take() {
                        layer.styles = backup;
                    }

                    if !layer.effects.is_empty() || !layer.styles.is_empty() {
                        layer.effects_dirty.store(true, std::sync::atomic::Ordering::Release);
                        self.document.composite_dirty = true;
                    }

                    // Take the pooled before buffer (no allocation)
                    log::debug!("[end_stroke] layer_pixels={}, stroke_before_buf={:?}", layer_pixels, self.stroke_before_buf.as_ref().map(|b| b.len()));
                    let before = match self.stroke_before_buf.take() {
                        Some(buf) if buf.len() == layer_pixels => buf,
                        _ => vec![0u8; layer_pixels],
                    };
                    log::debug!("[end_stroke] before.len()={}, lw={}", before.len(), lw);

                    // Layer 3: Zero-copy after-snapshot via raw pointer.
                    // SAFETY: After end_stroke(), no mutations happen to layer.pixels
                    // until the next begin_stroke(). The background thread only reads.
                    let after_ptr = SendPtr::new(layer.pixels.as_ptr());
                    let after_len = layer.pixels.len();

                    let after_shapes = if let LayerData::Vector { shapes } = &layer.data {
                        Some(shapes.clone())
                    } else {
                        None
                    };

                    let bounds = self.last_stroke_bounds;
                    let pending_history = self.pending_history.clone();

                    std::thread::spawn(move || {
                        let t_start = std::time::Instant::now();
                        // SAFETY: after_ptr points to layer.pixels which is valid for the
                        // engine lifetime. No mutations happen between end_stroke() and next begin_stroke().
                        let after_slice = unsafe { std::slice::from_raw_parts(after_ptr.as_ptr(), after_len) };
                        let mut item = PendingHistoryItem {
                            layer_idx: idx,
                            before_pixels: Vec::new(),
                            after_pixels: Vec::new(),
                            bounds: None,
                            before_shapes,
                            after_shapes,
                            description: "Brush Stroke".to_string(),
                            return_before: None,
                            return_after: None,
                        };

                        if let Some([sx0, sy0, sx1, sy1]) = bounds {
                            if sx0 < sx1 && sy0 < sy1 {
                                let rw = sx1 - sx0;
                                let rh = sy1 - sy0;
                                let rect_size = (rw * rh * 4) as usize;
                                log::debug!("[end_stroke_bg] bounds=[{},{},{},{}], rw={}, rh={}, rect_size={}, before.len()={}, after.len()={}, lw={}", sx0, sy0, sx1, sy1, rw, rh, rect_size, before.len(), after_slice.len(), lw);
                                let mut before_rect = vec![0u8; rect_size];
                                let mut after_rect = vec![0u8; rect_size];
                                for row in 0..rh {
                                    let src_start = (((sy0 + row) * lw + sx0) * 4) as usize;
                                    let dst_start = (row * rw * 4) as usize;
                                    let len = (rw * 4) as usize;
                                    before_rect[dst_start..dst_start + len]
                                        .copy_from_slice(&before[src_start..src_start + len]);
                                    after_rect[dst_start..dst_start + len]
                                        .copy_from_slice(&after_slice[src_start..src_start + len]);
                                }

                                if before_rect != after_rect {
                                    item.before_pixels = before_rect;
                                    item.after_pixels = after_rect;
                                    item.bounds = Some([sx0, sy0, sx1, sy1]);
                                    // Return before buffer to pool (after is a borrowed pointer)
                                    item.return_before = Some(before);
                                    pending_history.lock().unwrap().push(item);
                                }
                                log::trace!("[end_stroke_async] Background snapshot comparison took {}ms", t_start.elapsed().as_millis());
                                return;
                            }
                        }

                        // Fallback: use full buffers as pixels (no sub-rect optimization)
                        item.before_pixels = before;
                        // SAFETY: after_slice is valid for the engine lifetime.
                        // Copy into owned memory for PendingHistoryItem.
                        item.after_pixels = after_slice.to_vec();
                        pending_history.lock().unwrap().push(item);
                        log::trace!("[end_stroke_async] Background fallback snapshot comparison took {}ms", t_start.elapsed().as_millis());
                    });

                    // Stroke changed layer.pixels — invalidate effects backup so
                    // the next style edit captures the post-stroke raw pixels.
                    self.raw_pixel_backup.remove(&layer_id);
                    self.last_stroke_bounds = None;
                    self.last_stroke_pos = None;
                    self.below_cache_dirty = false;
                    self.document.composite_dirty = true;
                    if let Some(mask) = &mut self.active_stroke_mask {
                        mask.fill(0);
                    }
                    self.cached_selection_mask = None;
                    return;
                }
            }
        }
        self.last_stroke_pos = None;
        self.last_stroke_bounds = None;
        self.below_cache_dirty = false;
        self.document.composite_dirty = true;
        log::trace!("[end_stroke_async] fallback cleanup done, below_cache retained for reuse");
        if let Some(mask) = &mut self.active_stroke_mask {
            mask.fill(0);
        }
        self.cached_selection_mask = None;
    }

    /// Expand the stroke bounds to include the given point with radius.
    pub(crate) fn expand_stroke_bounds(&mut self, x: u32, y: u32, radius: f32) {
        let r = (radius as i32).max(1);
        let x0 = (x as i32 - r).max(0) as u32;
        let y0 = (y as i32 - r).max(0) as u32;
        let x1 = ((x as i32 + r + 1).min(self.document.canvas_width as i32 - 1)).max(0) as u32;
        let y1 = ((y as i32 + r + 1).min(self.document.canvas_height as i32 - 1)).max(0) as u32;
        match &mut self.last_stroke_bounds {
            Some(b) => {
                b[0] = b[0].min(x0);
                b[1] = b[1].min(y0);
                b[2] = b[2].max(x1);
                b[3] = b[3].max(y1);
            }
            None => {
                self.last_stroke_bounds = Some([x0, y0, x1, y1]);
            }
        }
    }

    /// Rebuild the `below_cache` snapshot of layers below the active layer
    /// only when necessary. Reuses the existing buffer when the active layer
    /// index matches and no layer property that affects the below-layer
    /// composite has changed since the last stroke.
    fn rebuild_below_cache_if_needed(&mut self) {
        let active_idx = self.document.active_layer;
        let cw = self.document.canvas_width;
        let ch = self.document.canvas_height;
        let buf_size = (cw * ch * 4) as usize;

        let can_reuse = !self.below_cache_dirty
            && self.below_cache_active_idx == Some(active_idx)
            && active_idx > 0
            && self.below_cache.as_ref().map_or(false, |b| b.len() == buf_size);

        log::trace!(
            "[begin_stroke] below_cache state: active_idx={}, reuse={}, dirty={}, existing_cache={}, existing_active_idx={:?}",
            active_idx, can_reuse, self.below_cache_dirty, self.below_cache.is_some(), self.below_cache_active_idx
        );

        if can_reuse {
            log::trace!("[begin_stroke] REUSING below_cache for active_idx={}", active_idx);
            self.below_cache_dirty = false;
            return;
        }

        if active_idx > 0 {
            let lcount = self.document.layers.len();
            if self.tile_layers.len() < lcount {
                self.tile_layers.resize_with(lcount, || None);
            }
            let mut tiles_built = 0usize;
            for (ti, tl) in self.document.layers.iter().enumerate() {
                if ti >= active_idx { break; }
                if !tl.pixels.is_empty() && self.tile_layers[ti].is_none() {
                    self.tile_layers[ti] = Some(TiledLayer::from_dense(
                        &tl.pixels, tl.width, tl.height,
                    ));
                    tiles_built += 1;
                    log::trace!("[begin_stroke] built tile for below layer[{}] id={}", ti, tl.id);
                }
            }
            let mut cache = match self.below_cache.take() {
                Some(mut b) if b.len() == buf_size => {
                    b.fill(0);
                    b
                }
                _ => vec![0u8; buf_size],
            };
            let tile_slice_len = active_idx.min(self.tile_layers.len());
            let visible_below: Vec<(usize, bool)> = self.document.layers[..active_idx]
                .iter().enumerate().map(|(i, l)| (i, l.visible)).collect();
            log::trace!(
                "[begin_stroke] compositing below_cache: {} below_layers, tile_slice_len={}, visible_below={:?}, tiles_built={}",
                active_idx, tile_slice_len, visible_below, tiles_built
            );
            tiled::composite_tiled_into(
                &self.document.layers[..active_idx],
                &self.tile_layers[..tile_slice_len],
                cw, ch, 0, 0, cw, ch,
                &mut cache,
            );
            let non_zero = cache.iter().filter(|&&b| b != 0).count();
            log::trace!(
                "[begin_stroke] below_cache built: {} bytes, non-zero bytes={}, active_idx={}",
                cache.len(), non_zero, active_idx
            );
            self.below_cache = Some(cache);
            self.below_cache_active_idx = Some(active_idx);
            self.below_cache_dirty = false;
        } else {
            log::trace!("[begin_stroke] active_idx=0 (bottom layer), no below_cache");
            self.below_cache = None;
            self.below_cache_active_idx = None;
            self.below_cache_dirty = false;
        }
    }
}