diff --git a/src/elements/div.rs b/src/elements/div.rs
index b119b1e..d34fe63 100644
--- a/src/elements/div.rs
+++ b/src/elements/div.rs
@@ -15,8 +15,6 @@ impl Element for Arc<Div> {
 
     fn after_render(&mut self, scope: &mut crate::render_scope::RenderScope) {
         let transform = scope.get_transform().clone();
-
-        scope.draw();
         let (w, h) = scope.get_parent_size();
         scope.set_parent_size(transform.width, transform.height);
 
diff --git a/src/elements/flex.rs b/src/elements/flex.rs
index 5659ecd..5183c53 100644
--- a/src/elements/flex.rs
+++ b/src/elements/flex.rs
@@ -1,36 +1,34 @@
 use std::sync::{Arc, Mutex};
 
-use crate::{
-    extensions::Handler,
-    style::RawTransform,
-    widget::{Element, Widget},
-    RenderWrapperEvent,
-};
+use crate::widget::{Element, Widget};
 
 pub struct FlexRow {
     color: u32,
     gap: u16,
     children: Mutex<Vec<Arc<Widget>>>,
-}
-pub struct FlexCol {
-    transform: Mutex<RawTransform>,
-    color: u32,
-    gap: u16,
+    size: (u16, u16),
 }
 
-impl Element for Arc<FlexRow> {
+pub struct FlexCol {
+    color: u32,
+    gap: u16,
+    children: Mutex<Vec<Arc<Widget>>>,
+    size: (u16, u16),
+}
+
+impl Element for FlexRow {
     fn render(&mut self, scope: &mut crate::render_scope::RenderScope) {
-        let (width, height) = scope.get_size_or_parent();
+        let (width, height) = scope.get_size_or(self.size.0, self.size.1);
         scope.draw_rect(width, height, self.color);
     }
 
     fn after_render(&mut self, scope: &mut crate::render_scope::RenderScope) {
-        let transform = scope.get_transform().clone();
-
-        scope.draw();
+        let mut transform = scope.get_transform().clone();
         let (w, h) = scope.get_parent_size();
         scope.set_parent_size(transform.width, transform.height);
 
+        let mut v = 0;
+
         for elem in self.children.lock().unwrap().iter() {
             scope.clear();
             if let Some(t) = elem.get() {
@@ -41,64 +39,24 @@ impl Element for Arc<FlexRow> {
                 scope.set_transform(&t);
             }
             let t = scope.get_transform_mut();
+            transform.width = transform.width.max(t.width);
+            transform.height = transform.height.max(v + t.height);
             t.x += transform.x;
-            t.y += transform.y;
+            t.y += transform.y + v;
+            v += t.height + self.gap;
 
             scope.draw();
+
             elem.0.lock().unwrap().after_render(scope);
-            scope.clear();
         }
         scope.set_parent_size(w, h);
+        self.size = (transform.width, transform.height);
     }
 
     fn draw_child(&self, element: &Arc<Widget>) {
         self.children.lock().unwrap().push(element.clone());
     }
 
-    // fn draw_child(&self, element: &Arc<Widget>) {
-    //     let r = self.clone();
-    //     element.inject(move |w| {
-    //         let r = r.clone();
-    //         w.component(Handler::new(move |elem, e: &RenderWrapperEvent| {
-    //             let mut scope = e.get_scope();
-    //             scope.clear();
-
-    //             let mut transform = r.transform.lock().unwrap();
-
-    //             let (w, h) = scope.get_parent_size();
-    //             scope.set_parent_size(transform.width, transform.height);
-
-    //             if let Some(t) = elem.get() {
-    //                 scope.set_transform(&t);
-    //             }
-
-    //             elem.0.lock().unwrap().render(scope);
-
-    //             if let Some(t) = elem.get() {
-    //                 scope.set_transform(&t);
-    //             }
-
-    //             let elem_transform = scope.get_transform_mut();
-
-    //             transform.height += if transform.height == 0 {
-    //                 elem_transform.height
-    //             } else {
-    //                 elem_transform.height + r.gap
-    //             };
-    //             transform.width = transform.width.max(elem_transform.width);
-
-    //             elem_transform.x = transform.x;
-    //             elem_transform.y = transform.y;
-    //             transform.y += elem_transform.height + r.gap;
-
-    //             scope.draw();
-    //             elem.0.lock().unwrap().after_render(&mut scope);
-
-    //             scope.set_parent_size(w, h);
-    //         }))
-    //     })
-    // }
-
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }
@@ -109,75 +67,55 @@ impl Element for Arc<FlexRow> {
 }
 
 impl FlexRow {
-    pub fn new(color: u32, gap: u16) -> Arc<Self> {
-        Arc::new(Self {
+    pub fn new(color: u32, gap: u16) -> Self {
+        Self {
             children: Mutex::new(Vec::new()),
+            size: (0, 0),
             color,
             gap,
-        })
+        }
     }
 }
 
-impl Element for Arc<FlexCol> {
+impl Element for FlexCol {
     fn render(&mut self, scope: &mut crate::render_scope::RenderScope) {
-        let transform = self.transform.lock().unwrap();
-        let (w, h) = scope.get_size_or(transform.width, transform.height);
-        scope.draw_rect(w, h, self.color);
+        let (width, height) = scope.get_size_or(self.size.0, self.size.1);
+        scope.draw_rect(width, height, self.color);
     }
 
     fn after_render(&mut self, scope: &mut crate::render_scope::RenderScope) {
-        let mut t = self.transform.lock().unwrap();
-        let st = scope.get_transform();
+        let mut transform = scope.get_transform().clone();
+        let (w, h) = scope.get_parent_size();
+        scope.set_parent_size(transform.width, transform.height);
 
-        t.x = st.x;
-        t.y = st.y;
+        let mut v = 0;
 
-        t.width = 0;
-        t.height = 0;
+        for elem in self.children.lock().unwrap().iter() {
+            scope.clear();
+            if let Some(t) = elem.get() {
+                scope.set_transform(&t);
+            }
+            elem.0.lock().unwrap().render(scope);
+            if let Some(t) = elem.get() {
+                scope.set_transform(&t);
+            }
+            let t = scope.get_transform_mut();
+            transform.width = transform.width.max(v + t.width);
+            transform.height = transform.height.max(t.height);
+            t.x += transform.x + v;
+            t.y += transform.y;
+            v += t.width + self.gap;
+
+            scope.draw();
+
+            elem.0.lock().unwrap().after_render(scope);
+        }
+        scope.set_parent_size(w, h);
+        self.size = (transform.width, transform.height);
     }
 
     fn draw_child(&self, element: &Arc<Widget>) {
-        let r = self.clone();
-        element.inject(move |w| {
-            let r = r.clone();
-            w.component(Handler::new(move |elem, e: &RenderWrapperEvent| {
-                let mut scope = e.get_scope();
-                scope.clear();
-
-                let mut transform = r.transform.lock().unwrap();
-
-                let (w, h) = scope.get_parent_size();
-                scope.set_parent_size(transform.width, transform.height);
-
-                if let Some(t) = elem.get() {
-                    scope.set_transform(&t);
-                }
-
-                elem.0.lock().unwrap().render(scope);
-
-                if let Some(t) = elem.get() {
-                    scope.set_transform(&t);
-                }
-
-                let elem_transform = scope.get_transform_mut();
-
-                transform.width += if transform.width == 0 {
-                    elem_transform.width
-                } else {
-                    elem_transform.width + r.gap
-                };
-                transform.height = transform.height.max(elem_transform.height);
-
-                elem_transform.x = transform.x;
-                elem_transform.y = transform.y;
-                transform.x += elem_transform.width + r.gap;
-
-                scope.draw();
-                elem.0.lock().unwrap().after_render(&mut scope);
-
-                scope.set_parent_size(w, h);
-            }))
-        });
+        self.children.lock().unwrap().push(element.clone());
     }
 
     fn as_any(&self) -> &dyn std::any::Any {
@@ -190,11 +128,12 @@ impl Element for Arc<FlexCol> {
 }
 
 impl FlexCol {
-    pub fn new(color: u32, gap: u16) -> Arc<Self> {
-        Arc::new(Self {
-            transform: Mutex::new(RawTransform::new()),
+    pub fn new(color: u32, gap: u16) -> Self {
+        Self {
+            children: Mutex::new(Vec::new()),
+            size: (0, 0),
             color,
             gap,
-        })
+        }
     }
 }
diff --git a/src/frontend.rs b/src/frontend.rs
index f6c761e..91fb248 100644
--- a/src/frontend.rs
+++ b/src/frontend.rs
@@ -37,9 +37,6 @@ impl Rsx {
                         w.dependency_box(d);
                     }
                     child.draw_parent(screen, Some(w.clone()));
-                    if let Some(parent) = &parent {
-                        parent.0.lock().unwrap().draw_child(&w);
-                    }
                 }
             }
         }
diff --git a/src/macros.rs b/src/macros.rs
index 88c429a..3285798 100644
--- a/src/macros.rs
+++ b/src/macros.rs
@@ -157,7 +157,7 @@ macro_rules! rsx_inner {
         $crate::rsx_inner! { $r, $($rest)* };
     };
 
-    ($r:expr, $expand:ident($($inner:tt)*) $($rest:tt)*) => {
+    ($r:expr, $expand:ident => ($($inner:tt)*) $($rest:tt)*) => {
         $r.expand(&mut $expand($($inner)*));
         $crate::rsx_inner! { $r, $($rest)* };
     };
diff --git a/src/main.rs b/src/main.rs
index d5a7698..534ef3c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,11 @@
-use osui::{elements::div::Div, frontend::Rsx, rsx, style::Transform, Screen};
+use osui::{
+    elements::{div::Div, flex::FlexCol},
+    frontend::Rsx,
+    rsx,
+    state::use_state,
+    style::Transform,
+    Screen,
+};
 
 fn main() -> std::io::Result<()> {
     let screen = Screen::new();
@@ -10,15 +17,33 @@ fn main() -> std::io::Result<()> {
 
 fn app() -> Rsx {
     rsx! {
-        @Transform::new();
+        @Transform::center();
+        FlexCol {
+            test => (1)
+            test => (3)
+        } (0x0000ff, 4)
+    }
+}
+
+fn test(inc: i32) -> Rsx {
+    let count = use_state(0);
+
+    std::thread::spawn({
+        let count = count.clone();
+        move || loop {
+            std::thread::sleep(std::time::Duration::from_millis(300));
+            **count.get() += inc;
+        }
+    });
+
+    rsx! {
+        @Transform::new().dimensions(40, 3);
         Div {
-            @Transform::center().dimensions(40, 3);
+            @Transform::center().dimensions(8, 1);
             Div {
-                @Transform::center().dimensions(8, 1);
-                Div {
-                    "hello"
-                } (0x00ff00)
-            } (0xff0000)
-        } (0x0000ff)
+                %count
+                "{count}"
+            } (0x00ff00)
+        } (0xff0000)
     }
 }