🎉 Add shapes buffer to improve memory allocation

frontend/resources/wasm-playground/js/lib.js

@@ -79,6 +79,19 @@ function getHeapU32() {
   return Module.HEAPU32;
 }
 
+export function clearShapeFills() {
+  Module._clear_shape_fills();
+}
+
+export function addShapeSolidFill(argb) {
+  const ptr = allocBytes(176);
+  const heap = getHeapU32();
+  const dv = new DataView(heap.buffer);
+  dv.setUint8(ptr, 0x00, true);
+  dv.setUint32(ptr + 4, argb, true);
+  Module._add_shape_fill();
+}
+
 export function setShapeChildren(shapeIds) {
   const offset = allocBytes(shapeIds.length * 16);
   const heap = getHeapU32();

frontend/resources/wasm-playground/rects.html

@@ -25,43 +25,50 @@
   <script type="module">
     import initWasmModule from '/js/render_wasm.js';
     import {
-      init, assignCanvas, hexToU32ARGB, getRandomInt, getRandomColor,
+      init, addShapeSolidFill, assignCanvas, hexToU32ARGB, getRandomInt, getRandomColor,
       getRandomFloat, useShape, setShapeChildren, setupInteraction
     } from './js/lib.js';
 
     const canvas = document.getElementById("canvas");
     canvas.width =  window.innerWidth;
     canvas.height = window.innerHeight;
+    const shapes = 1_000;
 
     initWasmModule().then(Module => {
       init(Module);
       assignCanvas(canvas);
       Module._set_canvas_background(hexToU32ARGB("#FABADA", 1));
       Module._set_view(1, 0, 0);
+      Module._init_shapes_pool(shapes + 1);
       setupInteraction(canvas);
 
       const children = [];
-      for (let i = 0; i < 1000; i++) {
+      for (let i = 0; i < shapes; i++) {
         const uuid = crypto.randomUUID();
         children.push(uuid);
 
         useShape(uuid);
         Module._set_parent(0, 0, 0, 0);
         Module._set_shape_type(3);
-        const x1 = getRandomInt(0, canvas.width);
-        const y1 = getRandomInt(0, canvas.height);
+        const x1 = getRandomInt(0, canvas.width*3);
+        const y1 = getRandomInt(0, canvas.height*3);
         const width = getRandomInt(20, 100);
         const height = getRandomInt(20, 100);
         Module._set_shape_selrect(x1, y1, x1 + width, y1 + height);
 
         const color = getRandomColor();
         const argb = hexToU32ARGB(color, getRandomFloat(0.1, 1.0));
-        Module._add_shape_solid_fill(argb);
+        addShapeSolidFill(argb)
       }
 
       useShape("00000000-0000-0000-0000-000000000000");
       setShapeChildren(children);
+
+      performance.mark('render:begin');
       Module._render(Date.now());
+      performance.mark('render:end');
+      const { duration } = performance.measure('render', 'render:begin', 'render:end');
+      // alert(`render time: ${duration.toFixed(2)}ms`);
     });
     
   </script>

frontend/src/app/render_wasm/api.cljs

@@ -842,9 +842,12 @@
 
 (defn initialize
   [base-objects zoom vbox background]
-  (let [rgba (sr-clr/hex->u32argb background 1)]
+  (let [rgba         (sr-clr/hex->u32argb background 1)
+        shapes       (into [] (vals base-objects))
+        total-shapes (count shapes)]
     (h/call wasm/internal-module "_set_canvas_background" rgba)
     (h/call wasm/internal-module "_set_view" zoom (- (:x vbox)) (- (:y vbox)))
+    (h/call wasm/internal-module "_init_shapes_pool" total-shapes)
     (set-objects base-objects)))
 
 (def ^:private canvas-options

render-wasm/src/main.rs

@@ -150,6 +150,13 @@ pub extern "C" fn set_view(zoom: f32, x: f32, y: f32) {
     });
 }
 
+#[no_mangle]
+pub extern "C" fn init_shapes_pool(capacity: usize) {
+    with_state!(state, {
+        state.init_shapes_pool(capacity);
+    });
+}
+
 #[no_mangle]
 pub extern "C" fn use_shape(a: u32, b: u32, c: u32, d: u32) {
     with_state!(state, {

render-wasm/src/render.rs

@@ -434,7 +434,7 @@ impl RenderState {
 
     pub fn start_render_loop(
         &mut self,
-        tree: &mut HashMap<Uuid, Shape>,
+        tree: &mut HashMap<Uuid, &mut Shape>,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
         timestamp: i32,
@@ -501,7 +501,7 @@ impl RenderState {
 
     pub fn process_animation_frame(
         &mut self,
-        tree: &mut HashMap<Uuid, Shape>,
+        tree: &mut HashMap<Uuid, &mut Shape>,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
         timestamp: i32,
@@ -600,7 +600,7 @@ impl RenderState {
 
     pub fn render_shape_tree(
         &mut self,
-        tree: &mut HashMap<Uuid, Shape>,
+        tree: &mut HashMap<Uuid, &mut Shape>,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
         timestamp: i32,
@@ -855,7 +855,7 @@ impl RenderState {
 
     pub fn rebuild_tiles_shallow(
         &mut self,
-        tree: &mut HashMap<Uuid, Shape>,
+        tree: &mut HashMap<Uuid, &mut Shape>,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
     ) {
@@ -864,7 +864,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(&shape_id) {
+            if let Some(shape) = tree.get_mut(&shape_id) {
                 let mut shape = shape.clone();
                 if shape_id != Uuid::nil() {
                     if let Some(modifier) = modifiers.get(&shape_id) {
@@ -885,7 +885,7 @@ impl RenderState {
 
     pub fn rebuild_tiles(
         &mut self,
-        tree: &mut HashMap<Uuid, Shape>,
+        tree: &mut HashMap<Uuid, &mut Shape>,
         modifiers: &HashMap<Uuid, Matrix>,
         structure: &HashMap<Uuid, Vec<StructureEntry>>,
     ) {
@@ -894,7 +894,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(&shape_id) {
+            if let Some(shape) = tree.get_mut(&shape_id) {
                 let mut shape = shape.clone();
                 if shape_id != Uuid::nil() {
                     if let Some(modifier) = modifiers.get(&shape_id) {
@@ -914,11 +914,11 @@ impl RenderState {
 
     pub fn rebuild_modifier_tiles(
         &mut self,
-        tree: &mut HashMap<Uuid, Shape>,
+        tree: &mut HashMap<Uuid, &mut Shape>,
         modifiers: &HashMap<Uuid, Matrix>,
     ) {
         for (uuid, matrix) in modifiers {
-            if let Some(shape) = tree.get(uuid) {
+            if let Some(shape) = tree.get_mut(uuid) {
                 let mut shape: Shape = shape.clone();
                 shape.apply_transform(matrix);
                 self.update_tile_for(&shape);

render-wasm/src/shapes/modifiers.rs

@@ -16,7 +16,7 @@ use crate::uuid::Uuid;
 
 fn propagate_children(
     shape: &Shape,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     parent_bounds_before: &Bounds,
     parent_bounds_after: &Bounds,
     transform: Matrix,
@@ -83,7 +83,7 @@ fn propagate_children(
 
 fn calculate_group_bounds(
     shape: &Shape,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> Option<Bounds> {
@@ -303,19 +303,20 @@ mod tests {
 
     #[test]
     fn test_propagate_shape() {
-        let mut shapes = HashMap::<Uuid, Shape>::new();
+        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, child);
+        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);
-        shapes.insert(parent_id, parent.clone());
+        let mut parent_clone = parent.clone();
+        shapes.insert(parent_id, &mut parent_clone);
 
         let mut transform = Matrix::scale((2.0, 1.5));
         let x = parent.selrect.x();
@@ -341,17 +342,17 @@ mod tests {
 
     #[test]
     fn test_group_bounds() {
-        let mut shapes = HashMap::<Uuid, Shape>::new();
+        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, child1);
+        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, child2);
+        shapes.insert(child2_id, &mut child2);
 
         let parent_id = Uuid::new_v4();
         let mut parent = Shape::new(parent_id);
@@ -359,7 +360,8 @@ mod tests {
         parent.add_child(child1_id);
         parent.add_child(child2_id);
         parent.set_selrect(0.0, 0.0, 3.0, 3.0);
-        shapes.insert(parent_id, parent.clone());
+        let mut parent_clone = parent.clone();
+        shapes.insert(parent_id, &mut parent_clone);
 
         let bounds =
             calculate_group_bounds(&parent, &shapes, &HashMap::new(), &HashMap::new()).unwrap();

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

@@ -178,7 +178,7 @@ fn initialize_tracks(
     layout_bounds: &Bounds,
     layout_axis: &LayoutAxis,
     flex_data: &FlexData,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> Vec<TrackData> {
@@ -420,7 +420,7 @@ fn calculate_track_data(
     layout_data: &LayoutData,
     flex_data: &FlexData,
     layout_bounds: &Bounds,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> Vec<TrackData> {
@@ -551,7 +551,7 @@ pub fn reflow_flex_layout(
     shape: &Shape,
     layout_data: &LayoutData,
     flex_data: &FlexData,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &mut HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> VecDeque<Modifier> {

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

@@ -40,7 +40,7 @@ fn calculate_tracks(
     grid_data: &GridData,
     layout_bounds: &Bounds,
     cells: &Vec<GridCell>,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &HashMap<Uuid, Bounds>,
 ) -> Vec<TrackData> {
     let layout_size = if is_column {
@@ -105,7 +105,7 @@ fn set_auto_base_size(
     column: bool,
     tracks: &mut Vec<TrackData>,
     cells: &Vec<GridCell>,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &HashMap<Uuid, Bounds>,
 ) {
     for cell in cells {
@@ -156,7 +156,7 @@ fn set_auto_multi_span(
     column: bool,
     tracks: &mut Vec<TrackData>,
     cells: &Vec<GridCell>,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &HashMap<Uuid, Bounds>,
 ) {
     // Remove groups with flex (will be set in flex_multi_span)
@@ -230,7 +230,7 @@ fn set_flex_multi_span(
     column: bool,
     tracks: &mut Vec<TrackData>,
     cells: &Vec<GridCell>,
-    shapes: &HashMap<Uuid, Shape>,
+    shapes: &HashMap<Uuid, &mut Shape>,
     bounds: &HashMap<Uuid, Bounds>,
 ) {
     // Remove groups without flex
@@ -509,7 +509,7 @@ fn cell_bounds(
 fn create_cell_data<'a>(
     layout_bounds: &Bounds,
     children: &IndexSet<Uuid>,
-    shapes: &'a HashMap<Uuid, Shape>,
+    shapes: &'a HashMap<Uuid, &mut Shape>,
     cells: &Vec<GridCell>,
     column_tracks: &Vec<TrackData>,
     row_tracks: &Vec<TrackData>,
@@ -618,7 +618,7 @@ pub fn reflow_grid_layout<'a>(
     shape: &Shape,
     layout_data: &LayoutData,
     grid_data: &GridData,
-    shapes: &'a HashMap<Uuid, Shape>,
+    shapes: &'a HashMap<Uuid, &mut Shape>,
     bounds: &mut HashMap<Uuid, Bounds>,
     structure: &HashMap<Uuid, Vec<StructureEntry>>,
 ) -> VecDeque<Modifier> {

render-wasm/src/state.rs

@@ -7,6 +7,54 @@ use crate::shapes::Shape;
 use crate::shapes::StructureEntry;
 use crate::uuid::Uuid;
 
+/// 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
+    shapes: Vec<Box<Shape>>,
+    counter: usize,
+}
+
+impl ShapesPool {
+    pub fn new() -> Self {
+        ShapesPool {
+            shapes: vec![],
+            counter: 0,
+        }
+    }
+
+    pub fn initialize(&mut self, capacity: usize) {
+        self.counter = 0;
+        self.shapes = Vec::with_capacity(capacity);
+        for _ in 0..capacity {
+            self.shapes.push(Box::new(Shape::new(Uuid::nil())));
+        }
+    }
+
+    pub fn add_shape(&mut self, id: Uuid) -> &mut Shape {
+        if self.counter >= self.shapes.len() {
+            self.shapes.push(Box::new(Shape::new(Uuid::nil())));
+        }
+        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.
@@ -16,9 +64,10 @@ pub(crate) struct State<'a> {
     pub render_state: RenderState,
     pub current_id: Option<Uuid>,
     pub current_shape: Option<&'a mut Shape>,
-    pub shapes: HashMap<Uuid, Shape>,
+    pub shapes: HashMap<Uuid, &'a mut Shape>,
     pub modifiers: HashMap<Uuid, skia::Matrix>,
     pub structure: HashMap<Uuid, Vec<StructureEntry>>,
+    pub shapes_pool: ShapesPool,
 }
 
 impl<'a> State<'a> {
@@ -30,6 +79,7 @@ impl<'a> State<'a> {
             shapes: HashMap::with_capacity(capacity),
             modifiers: HashMap::new(),
             structure: HashMap::new(),
+            shapes_pool: ShapesPool::new(),
         }
     }
 
@@ -61,13 +111,17 @@ impl<'a> State<'a> {
         Ok(())
     }
 
+    pub fn init_shapes_pool(&mut self, capacity: usize) {
+        self.shapes_pool.initialize(capacity);
+    }
+
     pub fn use_shape(&'a mut self, id: Uuid) {
         if !self.shapes.contains_key(&id) {
-            let new_shape = Shape::new(id);
+            let new_shape = self.shapes_pool.add_shape(id);
             self.shapes.insert(id, new_shape);
         }
         self.current_id = Some(id);
-        self.current_shape = self.shapes.get_mut(&id);
+        self.current_shape = self.shapes.get_mut(&id).map(|r| &mut **r);
     }
 
     pub fn delete_shape(&mut self, id: Uuid) {