diff --git a/render-wasm/src/main.rs b/render-wasm/src/main.rs
index cf4907a00..52e97c0e2 100644
--- a/render-wasm/src/main.rs
+++ b/render-wasm/src/main.rs
@@ -3,12 +3,9 @@ pub mod shapes;
 pub mod state;
 pub mod utils;
 
-use std::collections::HashMap;
-
 use skia_safe as skia;
 use uuid::Uuid;
 
-use crate::shapes::Shape;
 use crate::state::State;
 use crate::utils::uuid_from_u32_quartet;
 
@@ -37,34 +34,12 @@ pub unsafe extern "C" fn draw_all_shapes(zoom: f32, pan_x: f32, pan_y: f32) {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
 
     reset_canvas();
-    scale(zoom, zoom);
-    translate(pan_x, pan_y);
+    render::scale(state, zoom, zoom);
+    render::translate(state, pan_x, pan_y);
 
     render::render_shape_tree(state, Uuid::nil());
 
-    flush();
-}
-
-#[no_mangle]
-pub unsafe extern "C" fn flush() {
-    let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    state
-        .render_state
-        .gpu_state
-        .context
-        .flush_and_submit_surface(&mut state.render_state.surface, None);
-}
-
-#[no_mangle]
-pub unsafe extern "C" fn translate(dx: f32, dy: f32) {
-    let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    state.render_state.surface.canvas().translate((dx, dy));
-}
-
-#[no_mangle]
-pub unsafe extern "C" fn scale(sx: f32, sy: f32) {
-    let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    state.render_state.surface.canvas().scale((sx, sy));
+    render::flush(state);
 }
 
 #[no_mangle]
@@ -78,29 +53,18 @@ pub extern "C" fn reset_canvas() {
     state.render_state.surface.canvas().reset_matrix();
 }
 
-pub fn get_or_create_shape<'a>(shapes: &'a mut HashMap<Uuid, Shape>, id: Uuid) -> &'a mut Shape {
-    if !shapes.contains_key(&id) {
-        let new_shape = Shape::new(id);
-        shapes.insert(id, new_shape);
-    }
-
-    shapes.get_mut(&id).unwrap()
-}
-
 #[no_mangle]
 pub extern "C" fn use_shape(a: u32, b: u32, c: u32, d: u32) {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
     let id = uuid_from_u32_quartet(a, b, c, d);
-    state.current_id = Some(id);
-    let shapes = &mut state.shapes;
-    state.current_shape = Some(get_or_create_shape(shapes, id));
+    state.use_shape(id);
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn set_shape_selrect(x1: f32, y1: f32, x2: f32, y2: f32) {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
 
-    if let Some(shape) = state.current_shape.as_deref_mut() {
+    if let Some(shape) = state.current_shape() {
         shape.selrect.x1 = x1;
         shape.selrect.y1 = y1;
         shape.selrect.x2 = x2;
@@ -111,7 +75,7 @@ pub unsafe extern "C" fn set_shape_selrect(x1: f32, y1: f32, x2: f32, y2: f32) {
 #[no_mangle]
 pub unsafe extern "C" fn set_shape_rotation(rotation: f32) {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    if let Some(shape) = state.current_shape.as_deref_mut() {
+    if let Some(shape) = state.current_shape() {
         shape.rotation = rotation;
     }
 }
@@ -119,7 +83,7 @@ pub unsafe extern "C" fn set_shape_rotation(rotation: f32) {
 #[no_mangle]
 pub unsafe extern "C" fn set_shape_transform(a: f32, b: f32, c: f32, d: f32, e: f32, f: f32) {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    if let Some(shape) = state.current_shape.as_deref_mut() {
+    if let Some(shape) = state.current_shape() {
         shape.transform.a = a;
         shape.transform.b = b;
         shape.transform.c = c;
@@ -133,7 +97,7 @@ pub unsafe extern "C" fn set_shape_transform(a: f32, b: f32, c: f32, d: f32, e:
 pub extern "C" fn add_shape_child(a: u32, b: u32, c: u32, d: u32) {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
     let id = uuid_from_u32_quartet(a, b, c, d);
-    if let Some(shape) = state.current_shape.as_deref_mut() {
+    if let Some(shape) = state.current_shape() {
         shape.children.push(id);
     }
 }
@@ -141,7 +105,7 @@ pub extern "C" fn add_shape_child(a: u32, b: u32, c: u32, d: u32) {
 #[no_mangle]
 pub extern "C" fn clear_shape_children() {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    if let Some(shape) = state.current_shape.as_deref_mut() {
+    if let Some(shape) = state.current_shape() {
         shape.children.clear();
     }
 }
@@ -149,7 +113,7 @@ pub extern "C" fn clear_shape_children() {
 #[no_mangle]
 pub extern "C" fn add_shape_solid_fill(r: u8, g: u8, b: u8, a: f32) {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    if let Some(shape) = state.current_shape.as_deref_mut() {
+    if let Some(shape) = state.current_shape() {
         let alpha: u8 = (a * 0xff as f32).floor() as u8;
         let color = skia::Color::from_argb(alpha, r, g, b);
         shape.add_fill(shapes::Fill::from(color));
@@ -159,7 +123,7 @@ pub extern "C" fn add_shape_solid_fill(r: u8, g: u8, b: u8, a: f32) {
 #[no_mangle]
 pub extern "C" fn clear_shape_fills() {
     let state = unsafe { STATE.as_mut() }.expect("got an invalid state pointer");
-    if let Some(shape) = state.current_shape.as_deref_mut() {
+    if let Some(shape) = state.current_shape() {
         shape.clear_fills();
     }
 }
diff --git a/render-wasm/src/render.rs b/render-wasm/src/render.rs
index 473272bca..97fce0221 100644
--- a/render-wasm/src/render.rs
+++ b/render-wasm/src/render.rs
@@ -75,6 +75,22 @@ pub(crate) fn create_surface(gpu_state: &mut GpuState, width: i32, height: i32)
     .unwrap()
 }
 
+pub(crate) fn flush(state: &mut State) {
+    state
+        .render_state
+        .gpu_state
+        .context
+        .flush_and_submit_surface(&mut state.render_state.surface, None);
+}
+
+pub(crate) fn translate(state: &mut State, dx: f32, dy: f32) {
+    state.render_state.surface.canvas().translate((dx, dy));
+}
+
+pub(crate) fn scale(state: &mut State, sx: f32, sy: f32) {
+    state.render_state.surface.canvas().scale((sx, sy));
+}
+
 pub(crate) fn render_shape_tree(state: &mut State, id: Uuid) {
     let shape = state.shapes.get(&id).unwrap();
 
diff --git a/render-wasm/src/state.rs b/render-wasm/src/state.rs
index b52387f00..90f0b3656 100644
--- a/render-wasm/src/state.rs
+++ b/render-wasm/src/state.rs
@@ -30,4 +30,18 @@ impl<'a> State<'a> {
     pub fn set_surface(&mut self, surface: skia::Surface) {
         self.render_state.surface = surface;
     }
+
+    pub fn use_shape(&'a mut self, id: Uuid) {
+        if !self.shapes.contains_key(&id) {
+            let new_shape = Shape::new(id);
+            self.shapes.insert(id, new_shape);
+        }
+
+        self.current_id = Some(id);
+        self.current_shape = self.shapes.get_mut(&id);
+    }
+
+    pub fn current_shape(&'a mut self) -> Option<&'a mut Shape> {
+        self.current_shape.as_deref_mut()
+    }
 }