♻️ Refactor wasm shapes state management

render-wasm/src/main.rs

@@ -82,7 +82,7 @@ macro_rules! with_current_shape {
 /// This is called from JS after the WebGL context has been created.
 #[no_mangle]
 pub extern "C" fn init(width: i32, height: i32) {
-    let state_box = Box::new(State::new(width, height, 2048));
+    let state_box = Box::new(State::new(width, height));
     unsafe {
         STATE = Some(state_box);
     }
@@ -778,7 +778,13 @@ pub extern "C" fn get_grid_coords(pos_x: f32, pos_y: f32) -> *mut u8 {
     let row: i32;
     let col: i32;
     with_state!(state, {
-        (row, col) = state.get_grid_coords(pos_x, pos_y);
+        if let Some((r, c)) = state.get_grid_coords(pos_x, pos_y) {
+            row = r;
+            col = c;
+        } else {
+            row = -1;
+            col = -1;
+        };
     });
     let mut bytes = vec![0; 8];
     bytes[0..4].clone_from_slice(&row.to_le_bytes());

render-wasm/src/render.rs

@@ -21,7 +21,8 @@ use options::RenderOptions;
 use surfaces::{SurfaceId, Surfaces};
 
 use crate::performance;
-use crate::shapes::{Corners, Fill, Shape, SolidColor, StructureEntry, Type};
+use crate::shapes::{Corners, Fill, Shape, StructureEntry, Type};
+use crate::state::ShapesPool;
 use crate::tiles::{self, PendingTiles, TileRect};
 use crate::uuid::Uuid;
 use crate::view::Viewbox;
@@ -558,7 +559,7 @@ impl RenderState {
 
     pub fn render_from_cache(
         &mut self,
-        shapes: &HashMap<Uuid, &mut Shape>,
+        shapes: &ShapesPool,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
     ) {
@@ -599,7 +600,7 @@ impl RenderState {
 
     pub fn start_render_loop(
         &mut self,
-        tree: &HashMap<Uuid, &mut Shape>,
+        tree: &ShapesPool,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
         scale_content: &HashMap<Uuid, f32>,
@@ -654,7 +655,7 @@ impl RenderState {
 
     pub fn process_animation_frame(
         &mut self,
-        tree: &HashMap<Uuid, &mut Shape>,
+        tree: &ShapesPool,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
         scale_content: &HashMap<Uuid, f32>,
@@ -850,7 +851,7 @@ impl RenderState {
 
     pub fn render_shape_tree_partial_uncached(
         &mut self,
-        tree: &HashMap<Uuid, &mut Shape>,
+        tree: &ShapesPool,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
         scale_content: &HashMap<Uuid, f32>,
@@ -967,7 +968,7 @@ impl RenderState {
 
     pub fn render_shape_tree_partial(
         &mut self,
-        tree: &HashMap<Uuid, &mut Shape>,
+        tree: &ShapesPool,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
         scale_content: &HashMap<Uuid, f32>,
@@ -1109,7 +1110,7 @@ impl RenderState {
 
     pub fn rebuild_tiles_shallow(
         &mut self,
-        tree: &mut HashMap<Uuid, &mut Shape>,
+        tree: &ShapesPool,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
     ) {
@@ -1118,7 +1119,7 @@ impl RenderState {
         self.surfaces.remove_cached_tiles();
         let mut nodes = vec![Uuid::nil()];
         while let Some(shape_id) = nodes.pop() {
-            if let Some(shape) = tree.get_mut(&shape_id) {
+            if let Some(shape) = tree.get(&shape_id) {
                 let mut shape: Cow<Shape> = Cow::Borrowed(shape);
                 if shape_id != Uuid::nil() {
                     if let Some(modifier) = modifiers.get(&shape_id) {
@@ -1139,7 +1140,7 @@ impl RenderState {
 
     pub fn rebuild_tiles(
         &mut self,
-        tree: &mut HashMap<Uuid, &mut Shape>,
+        tree: &ShapesPool,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
     ) {
@@ -1148,7 +1149,7 @@ impl RenderState {
         self.surfaces.remove_cached_tiles();
         let mut nodes = vec![Uuid::nil()];
         while let Some(shape_id) = nodes.pop() {
-            if let Some(shape) = tree.get_mut(&shape_id) {
+            if let Some(shape) = tree.get(&shape_id) {
                 let mut shape: Cow<Shape> = Cow::Borrowed(shape);
                 if shape_id != Uuid::nil() {
                     if let Some(modifier) = modifiers.get(&shape_id) {
@@ -1166,13 +1167,9 @@ impl RenderState {
         performance::end_measure!("rebuild_tiles");
     }
 
-    pub fn rebuild_modifier_tiles(
+    pub fn rebuild_modifier_tiles(&mut self, tree: &ShapesPool, modifiers: &HashMap<Uuid, Matrix>) {
-        &mut self,
-        tree: &mut HashMap<Uuid, &mut Shape>,
-        modifiers: &HashMap<Uuid, Matrix>,
-    ) {
         for (uuid, matrix) in modifiers {
-            if let Some(shape) = tree.get_mut(uuid) {
+            if let Some(shape) = tree.get(uuid) {
                 let mut shape: Cow<Shape> = Cow::Borrowed(shape);
                 shape.to_mut().apply_transform(matrix);
                 self.update_tile_for(&shape);

render-wasm/src/render/grid_layout.rs

@@ -4,13 +4,14 @@ use std::collections::HashMap;
 use crate::math::{Matrix, Rect};
 use crate::shapes::modifiers::grid_layout::grid_cell_data;
 use crate::shapes::{Shape, StructureEntry};
+use crate::state::ShapesPool;
 use crate::uuid::Uuid;
 
 pub fn render_overlay(
     zoom: f32,
     canvas: &skia::Canvas,
     shape: &Shape,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     modifiers: &HashMap<Uuid, Matrix>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) {

render-wasm/src/render/ui.rs

@@ -1,16 +1,15 @@
 use skia_safe::{self as skia, Color4f};
 use std::collections::HashMap;
 
+use super::{RenderState, ShapesPool, SurfaceId};
 use crate::math::Matrix;
 use crate::render::grid_layout;
-use crate::shapes::{Shape, StructureEntry};
+use crate::shapes::StructureEntry;
 use crate::uuid::Uuid;
 
-use super::{RenderState, SurfaceId};
-
 pub fn render(
     render_state: &mut RenderState,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     modifiers: &HashMap<Uuid, Matrix>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) {

render-wasm/src/shapes.rs

@@ -44,6 +44,8 @@ use crate::math;
 use crate::math::{Bounds, Matrix, Point};
 use indexmap::IndexSet;
 
+use crate::state::ShapesPool;
+
 const MIN_VISIBLE_SIZE: f32 = 2.0;
 const ANTIALIAS_THRESHOLD: f32 = 15.0;
 
@@ -781,7 +783,7 @@ impl Shape {
 
     pub fn all_children_with_self(
         &self,
-        shapes: &HashMap<Uuid, &mut Shape>,
+        shapes: &ShapesPool,
         include_hidden: bool,
     ) -> IndexSet<Uuid> {
         once(self.id)

render-wasm/src/shapes/modifiers.rs

@@ -11,13 +11,14 @@ use crate::shapes::{
     auto_height, set_paragraphs_width, ConstraintH, ConstraintV, Frame, Group, GrowType, Layout,
     Modifier, Shape, StructureEntry, TransformEntry, Type,
 };
+use crate::state::ShapesPool;
 use crate::state::State;
 use crate::uuid::Uuid;
 
 #[allow(clippy::too_many_arguments)]
 fn propagate_children(
     shape: &Shape,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     parent_bounds_before: &Bounds,
     parent_bounds_after: &Bounds,
     transform: Matrix,
@@ -88,7 +89,7 @@ fn propagate_children(
 
 fn calculate_group_bounds(
     shape: &Shape,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> Option<Bounds> {
@@ -420,21 +421,25 @@ mod tests {
 
     #[test]
     fn test_propagate_shape() {
-        let mut shapes = HashMap::<Uuid, &mut Shape>::new();
-
-        let child_id = Uuid::new_v4();
-        let mut child = Shape::new(child_id);
-        child.set_selrect(3.0, 3.0, 2.0, 2.0);
-        shapes.insert(child_id, &mut child);
-
         let parent_id = Uuid::new_v4();
-        let mut parent = Shape::new(parent_id);
-        parent.set_shape_type(Type::Group(Group::default()));
-        parent.add_child(child_id);
-        parent.set_selrect(1.0, 1.0, 5.0, 5.0);
-        let mut parent_clone = parent.clone();
-        shapes.insert(parent_id, &mut parent_clone);
 
+        let shapes = {
+            let mut shapes = ShapesPool::new();
+            shapes.initialize(10);
+
+            let child_id = Uuid::new_v4();
+            let child = shapes.add_shape(child_id);
+            child.set_selrect(3.0, 3.0, 2.0, 2.0);
+
+            let parent = shapes.add_shape(parent_id);
+            parent.set_shape_type(Type::Group(Group::default()));
+            parent.add_child(child_id);
+            parent.set_selrect(1.0, 1.0, 5.0, 5.0);
+
+            shapes
+        };
+
+        let parent = shapes.get(&parent_id).unwrap();
         let mut transform = Matrix::scale((2.0, 1.5));
         let x = parent.selrect.x();
         let y = parent.selrect.y();
@@ -445,7 +450,7 @@ mod tests {
         let bounds_after = bounds_before.transform(&transform);
 
         let result = propagate_children(
-            &parent,
+            parent,
             &shapes,
             &bounds_before,
             &bounds_after,
@@ -460,29 +465,31 @@ mod tests {
 
     #[test]
     fn test_group_bounds() {
-        let mut shapes = HashMap::<Uuid, &mut Shape>::new();
-
-        let child1_id = Uuid::new_v4();
-        let mut child1 = Shape::new(child1_id);
-        child1.set_selrect(3.0, 3.0, 2.0, 2.0);
-        shapes.insert(child1_id, &mut child1);
-
-        let child2_id = Uuid::new_v4();
-        let mut child2 = Shape::new(child2_id);
-        child2.set_selrect(0.0, 0.0, 1.0, 1.0);
-        shapes.insert(child2_id, &mut child2);
-
         let parent_id = Uuid::new_v4();
-        let mut parent = Shape::new(parent_id);
+        let shapes = {
-        parent.set_shape_type(Type::Group(Group::default()));
+            let mut shapes = ShapesPool::new();
-        parent.add_child(child1_id);
+            shapes.initialize(10);
-        parent.add_child(child2_id);
+
-        parent.set_selrect(0.0, 0.0, 3.0, 3.0);
+            let child1_id = Uuid::new_v4();
-        let mut parent_clone = parent.clone();
+            let child1 = shapes.add_shape(child1_id);
-        shapes.insert(parent_id, &mut parent_clone);
+            child1.set_selrect(3.0, 3.0, 2.0, 2.0);
+
+            let child2_id = Uuid::new_v4();
+            let child2 = shapes.add_shape(child2_id);
+            child2.set_selrect(0.0, 0.0, 1.0, 1.0);
+
+            let parent = shapes.add_shape(parent_id);
+            parent.set_shape_type(Type::Group(Group::default()));
+            parent.add_child(child1_id);
+            parent.add_child(child2_id);
+            parent.set_selrect(0.0, 0.0, 3.0, 3.0);
+            shapes
+        };
+
+        let parent = shapes.get(&parent_id).unwrap();
 
         let bounds =
-            calculate_group_bounds(&parent, &shapes, &HashMap::new(), &HashMap::new()).unwrap();
+            calculate_group_bounds(parent, &shapes, &HashMap::new(), &HashMap::new()).unwrap();
 
         assert_eq!(bounds.width(), 3.0);
         assert_eq!(bounds.height(), 3.0);

render-wasm/src/shapes/modifiers/flex_layout.rs

@@ -4,6 +4,7 @@ use crate::shapes::{
     AlignContent, AlignItems, AlignSelf, FlexData, JustifyContent, LayoutData, LayoutItem,
     Modifier, Shape, StructureEntry,
 };
+use crate::state::ShapesPool;
 use crate::uuid::Uuid;
 
 use std::collections::{HashMap, VecDeque};
@@ -178,7 +179,7 @@ fn initialize_tracks(
     layout_bounds: &Bounds,
     layout_axis: &LayoutAxis,
     flex_data: &FlexData,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> Vec<TrackData> {
@@ -430,7 +431,7 @@ fn calculate_track_data(
     layout_data: &LayoutData,
     flex_data: &FlexData,
     layout_bounds: &Bounds,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> Vec<TrackData> {
@@ -570,7 +571,7 @@ pub fn reflow_flex_layout(
     shape: &Shape,
     layout_data: &LayoutData,
     flex_data: &FlexData,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &mut HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> VecDeque<Modifier> {

render-wasm/src/shapes/modifiers/grid_layout.rs

@@ -4,6 +4,7 @@ use crate::shapes::{
     JustifyContent, JustifyItems, JustifySelf, Layout, LayoutData, LayoutItem, Modifier, Shape,
     StructureEntry, Type,
 };
+use crate::state::ShapesPool;
 use crate::uuid::Uuid;
 use indexmap::IndexSet;
 use std::collections::{HashMap, VecDeque};
@@ -44,7 +45,7 @@ pub fn calculate_tracks(
     grid_data: &GridData,
     layout_bounds: &Bounds,
     cells: &Vec<GridCell>,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &HashMap<Uuid, Bounds>,
 ) -> Vec<TrackData> {
     let layout_size = if is_column {
@@ -121,7 +122,7 @@ fn set_auto_base_size(
     column: bool,
     tracks: &mut [TrackData],
     cells: &Vec<GridCell>,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &HashMap<Uuid, Bounds>,
 ) {
     for cell in cells {
@@ -172,7 +173,7 @@ fn set_auto_multi_span(
     column: bool,
     tracks: &mut [TrackData],
     cells: &[GridCell],
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &HashMap<Uuid, Bounds>,
 ) {
     // Remove groups with flex (will be set in flex_multi_span)
@@ -247,7 +248,7 @@ fn set_flex_multi_span(
     layout_data: &LayoutData,
     tracks: &mut [TrackData],
     cells: &[GridCell],
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &HashMap<Uuid, Bounds>,
 ) {
     // Remove groups without flex
@@ -539,7 +540,7 @@ fn cell_bounds(
 pub fn create_cell_data<'a>(
     layout_bounds: &Bounds,
     children: &IndexSet<Uuid>,
-    shapes: &'a HashMap<Uuid, &mut Shape>,
+    shapes: &'a ShapesPool,
     cells: &Vec<GridCell>,
     column_tracks: &[TrackData],
     row_tracks: &[TrackData],
@@ -552,7 +553,7 @@ pub fn create_cell_data<'a>(
             if !children.contains(&shape_id) {
                 None
             } else {
-                shapes.get(&shape_id).map(|v| &**v)
+                shapes.get(&shape_id)
             }
         } else {
             None
@@ -602,7 +603,7 @@ pub fn create_cell_data<'a>(
 
 pub fn grid_cell_data<'a>(
     shape: &Shape,
-    shapes: &'a HashMap<Uuid, &mut Shape>,
+    shapes: &'a ShapesPool,
     modifiers: &HashMap<Uuid, Matrix>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
     allow_empty: bool,
@@ -723,7 +724,7 @@ pub fn reflow_grid_layout(
     shape: &Shape,
     layout_data: &LayoutData,
     grid_data: &GridData,
-    shapes: &HashMap<Uuid, &mut Shape>,
+    shapes: &ShapesPool,
     bounds: &mut HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> VecDeque<Modifier> {

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: ShapesPool,
-    pub shapes: HashMap<Uuid, &'a mut Shape>,
     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> {
+impl State {
-    pub fn new(width: i32, height: i32, capacity: usize) -> Self {
+    pub fn new(width: i32, height: i32) -> Self {
         State {
             render_state: RenderState::new(width, height),
             current_id: None,
-            current_shape: None,
+            shapes: ShapesPool::new(),
-            shapes: HashMap::with_capacity(capacity),
             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) {
+    pub fn use_shape(&mut self, id: Uuid) {
-        if let Entry::Vacant(e) = self.shapes.entry(id) {
+        if !self.shapes.has(&id) {
-            let new_shape = self.shapes_pool.add_shape(id);
+            self.shapes.add_shape(id);
-            e.insert(new_shape);
         }
         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 {
+        let shape = {
-            panic!("Invalid current 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) {
+        if !shape.id.is_nil() {
-            self.render_state.update_tile_for(shape);
+            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) {
+    pub fn get_grid_coords(&self, pos_x: f32, pos_y: f32) -> Option<(i32, i32)> {
-        let Some(shape) = self.current_shape() else {
+        let shape = self.current_shape()?;
-            return (-1, -1);
-        };
-
         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
     }
 }

render-wasm/src/state/shapes_pool.rs

@@ -0,0 +1,81 @@
+use std::collections::HashMap;
+use std::iter;
+
+use crate::performance;
+use crate::shapes::Shape;
+use crate::uuid::Uuid;
+
+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 `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.
+///
+pub struct ShapesPool {
+    shapes: Vec<Shape>,
+    shapes_uuid_to_idx: HashMap<Uuid, usize>,
+    counter: usize,
+}
+
+impl ShapesPool {
+    pub fn new() -> Self {
+        ShapesPool {
+            shapes: vec![],
+            counter: 0,
+            shapes_uuid_to_idx: HashMap::default(),
+        }
+    }
+
+    pub fn initialize(&mut self, capacity: usize) {
+        performance::begin_measure!("shapes_pool_initialize");
+        self.counter = 0;
+        self.shapes_uuid_to_idx = HashMap::with_capacity(capacity);
+
+        let additional = capacity as i32 - self.shapes.len() as i32;
+        if additional <= 0 {
+            return;
+        }
+
+        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() {
+            let additional = (self.shapes.len() as f32 * SHAPES_POOL_ALLOC_MULTIPLIER) as usize;
+            self.shapes
+                .extend(iter::repeat_with(|| Shape::new(Uuid::nil())).take(additional));
+        }
+        let new_shape = &mut self.shapes[self.counter];
+        new_shape.id = id;
+        self.shapes_uuid_to_idx.insert(id, self.counter);
+        self.counter += 1;
+        new_shape
+    }
+
+    pub fn len(&self) -> usize {
+        self.shapes_uuid_to_idx.len()
+    }
+
+    pub fn has(&self, id: &Uuid) -> bool {
+        self.shapes_uuid_to_idx.contains_key(id)
+    }
+
+    pub fn get_mut(&mut self, id: &Uuid) -> Option<&mut Shape> {
+        let idx = *self.shapes_uuid_to_idx.get(id)?;
+        Some(&mut self.shapes[idx])
+    }
+
+    pub fn get(&self, id: &Uuid) -> Option<&Shape> {
+        let idx = *self.shapes_uuid_to_idx.get(id)?;
+        Some(&self.shapes[idx])
+    }
+}