♻️ Refactor wasm shapes state management

render-wasm/src/state.rs

@@ -1,9 +1,9 @@
-use std::collections::{hash_map::Entry, HashMap};
-use std::{iter, vec};
-
 use skia_safe::{self as skia, Path, Point};
+use std::collections::HashMap;
+
+mod shapes_pool;
+pub use shapes_pool::*;
 
-use crate::performance;
 use crate::render::RenderState;
 use crate::shapes::Shape;
 use crate::shapes::StructureEntry;
@@ -12,93 +12,29 @@ use crate::uuid::Uuid;
 
 use crate::shapes::modifiers::grid_layout::grid_cell_data;
 
-const SHAPES_POOL_ALLOC_MULTIPLIER: f32 = 1.3;
-
-/// A pool allocator for `Shape` objects that attempts to minimize memory reallocations.
-///
-/// `ShapesPool` pre-allocates a contiguous vector of boxed `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<Box<Shape>>`, which keeps the `Box` pointers
-/// in a contiguous memory block. The actual `Shape` instances are heap-allocated,
-/// and this approach ensures that pushing new shapes does not invalidate
-/// previously returned mutable references.
-///
-/// This is especially important because references to `Shape` are also held in the
-/// state shapes attribute
-pub(crate) struct ShapesPool {
-    // We need a box so that pushing here doesn't invalidate state.shapes references
-    // FIXME: See if we can avoid this
-    #[allow(clippy::vec_box)]
-    shapes: Vec<Box<Shape>>,
-    counter: usize,
-}
-
-impl ShapesPool {
-    pub fn new() -> Self {
-        ShapesPool {
-            shapes: vec![],
-            counter: 0,
-        }
-    }
-
-    pub fn initialize(&mut self, capacity: usize) {
-        performance::begin_measure!("shapes_pool_initialize");
-        self.counter = 0;
-        let additional = capacity as i32 - self.shapes.len() as i32;
-        if additional <= 0 {
-            return;
-        }
-
-        self.shapes.extend(
-            iter::repeat_with(|| Box::new(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() {
-            let additional = (self.shapes.len() as f32 * SHAPES_POOL_ALLOC_MULTIPLIER) as usize;
-            self.shapes
-                .extend(iter::repeat_with(|| Box::new(Shape::new(Uuid::nil()))).take(additional));
-        }
-        let new_shape = &mut self.shapes[self.counter];
-        new_shape.id = id;
-        self.counter += 1;
-        new_shape
-    }
-}
-
 /// This struct holds the state of the Rust application between JS calls.
 ///
 /// It is created by [init] and passed to the other exported functions.
 /// Note that rust-skia data structures are not thread safe, so a state
 /// must not be shared between different Web Workers.
-pub(crate) struct State<'a> {
+pub(crate) struct State {
     pub render_state: RenderState,
     pub current_id: Option<Uuid>,
-    pub current_shape: Option<&'a mut Shape>,
-    pub shapes: HashMap<Uuid, &'a mut Shape>,
+    pub shapes: ShapesPool,
     pub modifiers: HashMap<Uuid, skia::Matrix>,
     pub scale_content: HashMap<Uuid, f32>,
     pub structure: HashMap<Uuid, Vec<StructureEntry>>,
-    pub shapes_pool: ShapesPool,
 }
 
-impl<'a> State<'a> {
-    pub fn new(width: i32, height: i32, capacity: usize) -> Self {
+impl State {
+    pub fn new(width: i32, height: i32) -> Self {
         State {
             render_state: RenderState::new(width, height),
             current_id: None,
-            current_shape: None,
-            shapes: HashMap::with_capacity(capacity),
+            shapes: ShapesPool::new(),
             modifiers: HashMap::new(),
             scale_content: HashMap::new(),
             structure: HashMap::new(),
-            shapes_pool: ShapesPool::new(),
         }
     }
 
@@ -106,11 +42,11 @@ impl<'a> State<'a> {
         self.render_state.resize(width, height);
     }
 
-    pub fn render_state_mut(&'a mut self) -> &'a mut RenderState {
+    pub fn render_state_mut(&mut self) -> &mut RenderState {
         &mut self.render_state
     }
 
-    pub fn render_state(&'a self) -> &'a RenderState {
+    pub fn render_state(&self) -> &RenderState {
         &self.render_state
     }
 
@@ -150,16 +86,14 @@ impl<'a> State<'a> {
     }
 
     pub fn init_shapes_pool(&mut self, capacity: usize) {
-        self.shapes_pool.initialize(capacity);
+        self.shapes.initialize(capacity);
     }
 
-    pub fn use_shape(&'a mut self, id: Uuid) {
-        if let Entry::Vacant(e) = self.shapes.entry(id) {
-            let new_shape = self.shapes_pool.add_shape(id);
-            e.insert(new_shape);
+    pub fn use_shape(&mut self, id: Uuid) {
+        if !self.shapes.has(&id) {
+            self.shapes.add_shape(id);
         }
         self.current_id = Some(id);
-        self.current_shape = self.shapes.get_mut(&id).map(|r| &mut **r);
     }
 
     pub fn delete_shape(&mut self, id: Uuid) {
@@ -177,11 +111,11 @@ impl<'a> State<'a> {
     }
 
     pub fn current_shape_mut(&mut self) -> Option<&mut Shape> {
-        self.current_shape.as_deref_mut()
+        self.shapes.get_mut(&self.current_id?)
     }
 
     pub fn current_shape(&self) -> Option<&Shape> {
-        self.current_shape.as_deref()
+        self.shapes.get(&self.current_id?)
     }
 
     pub fn set_background_color(&mut self, color: skia::Color) {
@@ -189,14 +123,17 @@ impl<'a> State<'a> {
     }
 
     pub fn set_selrect_for_current_shape(&mut self, left: f32, top: f32, right: f32, bottom: f32) {
-        let Some(shape) = self.current_shape.as_deref_mut() else {
-            panic!("Invalid current shape")
+        let shape = {
+            let Some(shape) = self.current_shape_mut() else {
+                panic!("Invalid current shape")
+            };
+            shape.set_selrect(left, top, right, bottom);
+            shape.clone()
         };
 
-        shape.set_selrect(left, top, right, bottom);
         // We don't need to update the tile for the root shape.
         if !shape.id.is_nil() {
-            self.render_state.update_tile_for(shape);
+            self.render_state.update_tile_for(&shape);
         }
     }
 
@@ -207,34 +144,31 @@ impl<'a> State<'a> {
     }
 
     pub fn update_tile_for_current_shape(&mut self) {
-        let Some(shape) = self.current_shape.as_deref() else {
+        let Some(shape) = self.current_shape() else {
             panic!("Invalid current shape")
         };
-        if !shape.id.is_nil() && self.shapes.contains_key(&shape.id) {
-            self.render_state.update_tile_for(shape);
+        if !shape.id.is_nil() {
+            self.render_state.update_tile_for(&shape.clone());
         }
     }
 
     pub fn rebuild_tiles_shallow(&mut self) {
         self.render_state
-            .rebuild_tiles_shallow(&mut self.shapes, &self.modifiers, &self.structure);
+            .rebuild_tiles_shallow(&self.shapes, &self.modifiers, &self.structure);
     }
 
     pub fn rebuild_tiles(&mut self) {
         self.render_state
-            .rebuild_tiles(&mut self.shapes, &self.modifiers, &self.structure);
+            .rebuild_tiles(&self.shapes, &self.modifiers, &self.structure);
     }
 
     pub fn rebuild_modifier_tiles(&mut self) {
         self.render_state
-            .rebuild_modifier_tiles(&mut self.shapes, &self.modifiers);
+            .rebuild_modifier_tiles(&self.shapes, &self.modifiers);
     }
 
-    pub fn get_grid_coords(&self, pos_x: f32, pos_y: f32) -> (i32, i32) {
-        let Some(shape) = self.current_shape() else {
-            return (-1, -1);
-        };
-
+    pub fn get_grid_coords(&self, pos_x: f32, pos_y: f32) -> Option<(i32, i32)> {
+        let shape = self.current_shape()?;
         let bounds = shape.bounds();
         let position = Point::new(pos_x, pos_y);
 
@@ -251,10 +185,10 @@ impl<'a> State<'a> {
             let polygon = Path::polygon(points, true, None, None);
 
             if polygon.contains(position) {
-                return (cell.row as i32 + 1, cell.column as i32 + 1);
+                return Some((cell.row as i32 + 1, cell.column as i32 + 1));
             }
         }
 
-        (-1, -1)
+        None
     }
 }