penpot/render-wasm/src/state/shapes_pool.rs

use std::collections::HashMap;
use std::collections::VecDeque;
use std::iter;

use crate::performance;
use crate::shapes;
use crate::shapes::Shape;
use crate::uuid::Uuid;

use crate::shapes::StructureEntry;
use crate::skia;

use std::cell::OnceCell;

use crate::math;
use crate::math::bools as math_bools;
use crate::math::Matrix;

const SHAPES_POOL_ALLOC_MULTIPLIER: f32 = 1.3;

/// A pool allocator for `Shape` objects that attempts to minimize memory reallocations.
///
/// `ShapesPoolImpl` pre-allocates a contiguous vector of `Shape` instances,
/// which can be reused and indexed efficiently. This design helps avoid
/// memory reallocation overhead by reserving enough space in advance.
///
/// # Memory Layout
///
/// Shapes are stored in a `Vec<Shape>`, which keeps the `Shape` instances
/// in a contiguous memory block.
///
/// # Index-based Design
///
/// All auxiliary HashMaps (modifiers, structure, scale_content, modified_shape_cache)
/// use `usize` indices instead of `&'a Uuid` references. This eliminates:
/// - Unsafe lifetime extensions
/// - The need for `rebuild_references()` after Vec reallocation
/// - Complex lifetime annotations
///
/// The `uuid_to_idx` HashMap maps `Uuid` (owned) to indices, avoiding lifetime issues.
///
pub struct ShapesPoolImpl {
    shapes: Vec<Shape>,
    counter: usize,

    /// Maps UUID to index in the shapes Vec. Uses owned Uuid, no lifetime needed.
    uuid_to_idx: HashMap<Uuid, usize>,

    /// Cache for modified shapes, keyed by index
    modified_shape_cache: HashMap<usize, OnceCell<Shape>>,
    /// Transform modifiers, keyed by index
    modifiers: HashMap<usize, skia::Matrix>,
    /// UUIDs of shapes that have an active transform modifier, kept in sync
    /// with `modifiers`. Stored explicitly so that `modifier_ids()` is O(K)
    /// (K = number of modified shapes) instead of O(N_shapes) — avoids
    /// building a full reverse-index HashMap on every call.
    modifier_uuids: Vec<Uuid>,
    /// Structure entries, keyed by index
    structure: HashMap<usize, Vec<StructureEntry>>,
    /// Scale content values, keyed by index
    scale_content: HashMap<usize, f32>,
}

// Type aliases - no longer need lifetimes!
pub type ShapesPool = ShapesPoolImpl;
pub type ShapesPoolRef<'a> = &'a ShapesPoolImpl;
pub type ShapesPoolMutRef<'a> = &'a mut ShapesPoolImpl;

impl ShapesPoolImpl {
    pub fn new() -> Self {
        ShapesPoolImpl {
            shapes: vec![],
            counter: 0,
            uuid_to_idx: HashMap::default(),

            modified_shape_cache: HashMap::default(),
            modifiers: HashMap::default(),
            modifier_uuids: Vec::new(),
            structure: HashMap::default(),
            scale_content: HashMap::default(),
        }
    }

    pub fn initialize(&mut self, capacity: usize) {
        performance::begin_measure!("shapes_pool_initialize");
        self.counter = 0;
        self.uuid_to_idx = HashMap::with_capacity(capacity);

        let additional = capacity as i32 - self.shapes.len() as i32;
        if additional <= 0 {
            return;
        }

        // Reserve extra capacity to avoid future reallocations
        let target_capacity = (capacity as f32 * SHAPES_POOL_ALLOC_MULTIPLIER) as usize;
        self.shapes
            .reserve_exact(target_capacity.saturating_sub(self.shapes.len()));

        self.shapes
            .extend(iter::repeat_with(|| Shape::new(Uuid::nil())).take(additional as usize));
        performance::end_measure!("shapes_pool_initialize");
    }

    pub fn add_shape(&mut self, id: Uuid) -> &mut Shape {
        if self.counter >= self.shapes.len() {
            // We need more space
            let current_capacity = self.shapes.capacity();
            // Ensure we add at least 1 shape when the pool is empty
            let additional =
                ((self.shapes.len() as f32 * SHAPES_POOL_ALLOC_MULTIPLIER) as usize).max(1);
            let needed_capacity = self.shapes.len() + additional;

            if needed_capacity > current_capacity {
                // Reserve extra space to minimize future reallocations
                let extra_reserve = (needed_capacity as f32 * 0.5) as usize;
                self.shapes
                    .reserve(needed_capacity + extra_reserve - current_capacity);
            }

            self.shapes
                .extend(iter::repeat_with(|| Shape::new(Uuid::nil())).take(additional));
        }

        let idx = self.counter;
        let new_shape = &mut self.shapes[idx];
        new_shape.id = id;

        // Simply store the UUID -> index mapping. No unsafe lifetime tricks needed!
        self.uuid_to_idx.insert(id, idx);
        self.counter += 1;

        &mut self.shapes[idx]
    }
    // No longer needed! Index-based storage means no references to rebuild.
    // The old rebuild_references() function has been removed entirely.

    pub fn len(&self) -> usize {
        self.uuid_to_idx.len()
    }

    pub fn has(&self, id: &Uuid) -> bool {
        self.uuid_to_idx.contains_key(id)
    }

    pub fn get_mut(&mut self, id: &Uuid) -> Option<&mut Shape> {
        let idx = *self.uuid_to_idx.get(id)?;
        Some(&mut self.shapes[idx])
    }

    /// Returns the current transform modifier matrix for the shape, if any.
    pub fn get_modifier(&self, id: &Uuid) -> Option<&skia::Matrix> {
        let idx = *self.uuid_to_idx.get(id)?;
        self.modifiers.get(&idx)
    }

    /// Modifier applied to `id`, including one inherited from an ancestor.
    pub fn get_layout_modifier(&self, id: &Uuid) -> Option<skia::Matrix> {
        if let Some(matrix) = self.get_modifier(id) {
            return Some(*matrix);
        }
        let idx = *self.uuid_to_idx.get(id)?;
        self.find_nearest_ancestor_modifier(idx)
    }

    /// Get a shape by UUID without applying modifiers/structure/scale-content.
    pub fn get_raw(&self, id: &Uuid) -> Option<&Shape> {
        let idx = *self.uuid_to_idx.get(id)?;
        Some(&self.shapes[idx])
    }

    /// Get a shape by UUID. Returns the modified shape if modifiers/structure
    /// are applied, otherwise returns the base shape.
    pub fn get(&self, id: &Uuid) -> Option<&Shape> {
        let idx = *self.uuid_to_idx.get(id)?;

        let shape = &self.shapes[idx];

        // Check if this shape needs modification (has modifiers, structure changes, or is a bool)
        let needs_modification = shape.is_bool()
            || self.modifiers.contains_key(&idx)
            || self.structure.contains_key(&idx)
            || self.scale_content.contains_key(&idx);

        if needs_modification {
            // Check if we have a cached modified version
            if let Some(cell) = self.modified_shape_cache.get(&idx) {
                Some(cell.get_or_init(|| {
                    let mut modified_shape =
                        shape.transformed(self.modifiers.get(&idx), self.structure.get(&idx));

                    if self.to_update_bool(&modified_shape) {
                        math_bools::update_bool_to_path(&mut modified_shape, self);
                    }

                    if let Some(scale) = self.scale_content.get(&idx) {
                        modified_shape.scale_content(*scale);
                    }
                    modified_shape
                }))
            } else {
                Some(shape)
            }
        } else {
            if let Some(cell) = self.modified_shape_cache.get(&idx) {
                return Some(cell.get_or_init(|| {
                    if let Some(m) = self.find_nearest_ancestor_modifier(idx) {
                        shape.transformed(Some(&m), None)
                    } else {
                        shape.clone()
                    }
                }));
            }
            Some(shape)
        }
    }

    // Given an id, returns the depth in the tree-shaped structure
    // of shapes.
    pub fn get_depth(&self, id: &Uuid) -> usize {
        if id == &Uuid::nil() {
            return 0;
        }

        let Some(idx) = self.uuid_to_idx.get(id) else {
            return 0;
        };

        let shape = &self.shapes[*idx];

        let Some(parent_id) = shape.parent_id else {
            return 0;
        };

        self.get_depth(&parent_id) + 1
    }

    #[allow(dead_code)]
    pub fn iter(&self) -> std::slice::Iter<'_, Shape> {
        self.shapes.iter()
    }

    #[allow(dead_code)]
    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, Shape> {
        self.shapes.iter_mut()
    }

    fn clean_shape_cache(&mut self) {
        self.modified_shape_cache.clear()
    }

    pub fn dependent_ancestor_ids<'a>(&'a self, id: &Uuid) -> impl Iterator<Item = Uuid> + 'a {
        let mut current = self
            .uuid_to_idx
            .get(id)
            .and_then(|idx| self.shapes[*idx].parent_id);

        std::iter::from_fn(move || {
            let parent_id = current.filter(|parent_id| !parent_id.is_nil())?;
            let parent_idx = self.uuid_to_idx.get(&parent_id).copied()?;
            let parent = &self.shapes[parent_idx];
            if !parent.extrect_depends_on_children() {
                return None;
            }
            current = parent.parent_id;
            Some(parent_id)
        })
    }

    /// Drops the extrect cache of every ancestor whose extrect is affected by this shape
    /// stopping at the first ancestor that clips.
    pub fn invalidate_ancestors_extrect(&mut self, id: &Uuid) {
        let mut current = self
            .uuid_to_idx
            .get(id)
            .and_then(|idx| self.shapes[*idx].parent_id);

        while let Some(parent_id) = current.filter(|parent_id| !parent_id.is_nil()) {
            let Some(parent_idx) = self.uuid_to_idx.get(&parent_id).copied() else {
                break;
            };
            if !self.shapes[parent_idx].extrect_depends_on_children() {
                break;
            }

            self.shapes[parent_idx].invalidate_extrect();
            // `get` returns a snapshot clone, we need to get mut
            // and replace the OnceCell to reset it.
            if let Some(cell) = self.modified_shape_cache.get_mut(&parent_idx) {
                *cell = OnceCell::new();
            }

            current = self.shapes[parent_idx].parent_id;
        }
    }

    pub fn set_modifiers(&mut self, modifiers: HashMap<Uuid, skia::Matrix>) {
        let mut ids = Vec::<Uuid>::new();
        let mut modifiers_with_idx = HashMap::with_capacity(modifiers.len());

        for (uuid, matrix) in modifiers {
            if let Some(idx) = self.uuid_to_idx.get(&uuid).copied() {
                modifiers_with_idx.insert(idx, matrix);
                ids.push(uuid);
            }
        }

        // Expand every root modifier to its full descendant subtree.
        // When CLJS sends only root shapes (translation on drag), descendants
        // need the same matrix.
        // For resize/rotate, propagate-modifiers already includes all descendants.
        // Descendants are NOT pushed into `ids` / `modifier_uuids`: rebuild_modifier_tiles
        // runs for roots, and drops the non-clipping ancestors' extrects separately.
        let root_pairs: Vec<(usize, skia::Matrix)> = ids
            .iter()
            .filter_map(|uuid| {
                let idx = self.uuid_to_idx.get(uuid).copied()?;
                let matrix = modifiers_with_idx.get(&idx).copied()?;
                Some((idx, matrix))
            })
            .collect();

        let mut descendants_idxs: Vec<usize> = Vec::new();
        for (root_idx, matrix) in root_pairs {
            for descendant_idx in self.collect_all_descendants(root_idx) {
                if let std::collections::hash_map::Entry::Vacant(e) =
                    modifiers_with_idx.entry(descendant_idx)
                {
                    e.insert(matrix);
                    descendants_idxs.push(descendant_idx);
                }
            }
        }

        self.modifiers = modifiers_with_idx;

        for descendant_idx in descendants_idxs {
            self.modified_shape_cache
                .insert(descendant_idx, OnceCell::new());
        }

        // Compute ancestors before consuming `ids` so we can move it into
        // `modifier_uuids` without a clone.
        let all_ids = shapes::all_with_ancestors(&ids, self, true);

        for uuid in all_ids {
            if let Some(idx) = self.uuid_to_idx.get(&uuid).copied() {
                self.modified_shape_cache.insert(idx, OnceCell::new());
            }
        }

        // rebuild_modifier_tiles doesn't process every descendant individually.
        self.modifier_uuids = ids;
    }

    pub fn set_structure(&mut self, structure: HashMap<Uuid, Vec<StructureEntry>>) {
        // Convert HashMap<Uuid, V> to HashMap<usize, V> using indices
        // Initialize the cache cells for affected shapes
        let mut structure_with_idx = HashMap::with_capacity(structure.len());
        let mut ids = Vec::<Uuid>::new();

        for (uuid, entries) in structure {
            if let Some(idx) = self.uuid_to_idx.get(&uuid).copied() {
                structure_with_idx.insert(idx, entries);
                ids.push(uuid);
            }
        }
        self.structure = structure_with_idx;

        let all_ids = shapes::all_with_ancestors(&ids, self, true);
        for uuid in all_ids {
            if let Some(idx) = self.uuid_to_idx.get(&uuid).copied() {
                self.modified_shape_cache.insert(idx, OnceCell::new());
            }
        }
    }

    pub fn set_scale_content(&mut self, scale_content: HashMap<Uuid, f32>) {
        // Convert HashMap<Uuid, V> to HashMap<usize, V> using indices
        // Initialize the cache cells for affected shapes
        let mut scale_content_with_idx = HashMap::with_capacity(scale_content.len());
        let mut ids = Vec::<Uuid>::new();

        for (uuid, value) in scale_content {
            if let Some(idx) = self.uuid_to_idx.get(&uuid).copied() {
                scale_content_with_idx.insert(idx, value);
                ids.push(uuid);
            }
        }
        self.scale_content = scale_content_with_idx;

        let all_ids = shapes::all_with_ancestors(&ids, self, true);
        for uuid in all_ids {
            if let Some(idx) = self.uuid_to_idx.get(&uuid).copied() {
                self.modified_shape_cache.insert(idx, OnceCell::new());
            }
        }
    }

    /// Clears transient per-frame state (modifiers, structure, scale_content)
    /// and returns the list of UUIDs that had a `modifier` applied at the
    /// moment of cleaning. The caller can use that list to re-sync the tile
    /// index / tile cache for those shapes: after cleaning their modifier is
    /// gone, but if we don't touch their tiles they keep pointing at the
    /// previous modified position and the tile texture cache may serve stale
    /// pixels.
    pub fn clean_all(&mut self) -> Vec<Uuid> {
        self.clean_shape_cache();

        // `modifier_uuids` is kept in sync with `modifiers` by `set_modifiers`,
        // so we can take it directly — no need to rebuild a reverse index.
        let modified_uuids = std::mem::take(&mut self.modifier_uuids);

        self.modifiers = HashMap::default();
        self.structure = HashMap::default();
        self.scale_content = HashMap::default();

        modified_uuids
    }

    /// UUIDs of all shapes that currently have a transform modifier.
    /// Used by the throttled drag path so per-rAF tile invalidation can
    /// be done once with the current modifier set instead of once per
    /// pointer move.
    ///
    /// Returns a reference to avoid allocation on every call — callers
    /// inside hot render loops should hold this reference rather than
    /// calling `modifier_ids()` repeatedly.
    pub fn modifier_ids(&self) -> &[Uuid] {
        &self.modifier_uuids
    }

    pub fn subtree(&self, id: &Uuid) -> ShapesPoolImpl {
        let Some(shape) = self.get(id) else {
            panic!("Subtree not found");
        };

        let mut shapes = vec![];
        let mut new_idx = 0;
        let mut uuid_to_idx = HashMap::default();

        for child_id in shape.all_children_iter(self, true, true) {
            let Some(child_shape) = self.get(&child_id) else {
                panic!("Not found");
            };
            shapes.push(child_shape.clone());
            uuid_to_idx.insert(child_id, new_idx);
            new_idx += 1;
        }

        ShapesPoolImpl {
            shapes,
            counter: new_idx,
            uuid_to_idx,
            modified_shape_cache: HashMap::default(),
            modifiers: HashMap::default(),
            modifier_uuids: Vec::new(),
            structure: HashMap::default(),
            scale_content: HashMap::default(),
        }
    }

    fn collect_all_descendants(&self, idx: usize) -> Vec<usize> {
        let mut result = Vec::new();
        let mut queue: VecDeque<&Uuid> = VecDeque::new();
        let shape = &self.shapes[idx];
        for child_id in shape.children_ids_iter(false) {
            queue.push_back(child_id);
        }
        while let Some(child_id) = queue.pop_front() {
            if let Some(&child_idx) = self.uuid_to_idx.get(child_id) {
                result.push(child_idx);
                let child_shape = &self.shapes[child_idx];
                for grandchild_id in child_shape.children_ids_iter(false) {
                    queue.push_back(grandchild_id);
                }
            }
        }
        result
    }

    fn find_nearest_ancestor_modifier(&self, idx: usize) -> Option<Matrix> {
        let mut current_idx = idx;
        loop {
            let shape = &self.shapes[current_idx];
            let parent_id = shape.parent_id?;
            if parent_id == Uuid::nil() {
                return None;
            }
            let &parent_idx = self.uuid_to_idx.get(&parent_id)?;
            if let Some(matrix) = self.modifiers.get(&parent_idx) {
                return Some(*matrix);
            }
            current_idx = parent_idx;
        }
    }

    fn to_update_bool(&self, shape: &Shape) -> bool {
        if !shape.is_bool() {
            return false;
        }

        let default = &Matrix::default();

        // Get parent modifier by index
        let parent_idx = self.uuid_to_idx.get(&shape.id);
        let parent_modifier = parent_idx
            .and_then(|idx| self.modifiers.get(idx))
            .unwrap_or(default);

        // Returns true if the transform of any child is different to the parent's
        shape.all_children_iter(self, true, false).any(|child_id| {
            let child_modifier = self
                .uuid_to_idx
                .get(&child_id)
                .and_then(|idx| self.modifiers.get(idx))
                .unwrap_or(default);
            !math::is_close_matrix(parent_modifier, child_modifier)
        })
    }
}

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

impl Clone for ShapesPoolImpl {
    fn clone(&self) -> Self {
        ShapesPoolImpl {
            shapes: self.shapes.clone(),
            counter: self.counter,
            uuid_to_idx: self.uuid_to_idx.clone(),
            // The modified_shape_cache is a derived/computed cache; reset it on clone
            // so it gets lazily rebuilt on demand rather than cloning OnceCell state.
            modified_shape_cache: HashMap::default(),
            modifiers: self.modifiers.clone(),
            modifier_uuids: self.modifier_uuids.clone(),
            structure: self.structure.clone(),
            scale_content: self.scale_content.clone(),
        }
    }
}