diff --git a/engine/breeze/ui/ui.c b/engine/breeze/ui/ui.c
index 92b0b7f..a90c9ce 100644
--- a/engine/breeze/ui/ui.c
+++ b/engine/breeze/ui/ui.c
@@ -101,6 +101,7 @@ void bzUIBegin(BzUI *ui, i32 width, i32 height) {
     bzArrayPush(ui->nodeStack, ui->root);
 
     bzUINodeClearLinks(ui->root);
+    bzMemSet(&ui->root->style, 0, sizeof(ui->root->style));
 
     ui->root->semanticSize[BZ_UI_AXIS_X] = (BzUISize){
         .kind = BZ_UI_SIZE_PIXELS,
@@ -249,17 +250,15 @@ static void calculatePositions(BzUINode *node, f32 x, f32 y) {
     node->computedPosition[BZ_UI_AXIS_X] = x;
     node->computedPosition[BZ_UI_AXIS_Y] = y;
 
-    BzUINode *child = node->first;
-    while (child != NULL) {
-        switch (child->layout.type) {
-            case BZ_UI_LAYOUT_FLEX_BOX:
-                calculatePositionsFlexBox(child);
-                break;
-            default:
+    switch (node->layout.type) {
+        case BZ_UI_LAYOUT_FLEX_BOX:
+            calculatePositionsFlexBox(node);
+            break;
+        default:
+            for (BzUINode *child = node->first; child; child = child->next) {
                 calculatePositions(child, x, y);
-                break;
-        }
-        child = child->next;
+            }
+            break;
     }
 }
 
@@ -425,21 +424,12 @@ BzUINode *bzUIPopParent(BzUI *ui) {
     return node;
 }
 
-void bzUIPushLayout(BzUI *ui, BzUILayout layout) {
-    BzUINode *node = bzUINodeMake(ui, ui->idCount++, BZ_UI_NONE);
-    node->semanticSize[BZ_UI_AXIS_X] = (BzUISize) {
-        .kind  = BZ_UI_SIZE_PARENT_PERCENT,
-        .value = 1.0f,
-    };
-    node->semanticSize[BZ_UI_AXIS_Y] = node->semanticSize[BZ_UI_AXIS_X];
-
-    node->layout = layout;
-    bzUIPushParent(ui, node);
-}
-void bzUIPopLayout(BzUI *ui) {
-    BzUINode *last = bzArrayGet(ui->nodeStack, bzArraySize(ui->nodeStack) - 1);
-    BZ_ASSERT(last->flags == BZ_UI_NONE);
-    bzUIPopParent(ui);
+void bzUISetParentLayout(BzUI *ui, BzUILayout layout) {
+    i32 stackSize = bzArraySize(ui->nodeStack);
+    BZ_ASSERT(stackSize > 0);
+    BzUINode *last = bzArrayGet(ui->nodeStack, stackSize - 1);
+    BZ_ASSERT(last);
+    last->layout = layout;
 }
 
 BzUIInteraction bzUIGetInteraction(BzUI *ui, BzUINode *node) {
diff --git a/engine/breeze/ui/ui.h b/engine/breeze/ui/ui.h
index 87ed766..22b6eff 100644
--- a/engine/breeze/ui/ui.h
+++ b/engine/breeze/ui/ui.h
@@ -119,22 +119,16 @@ void  bzUIDestroy(BzUI *ui);
 void bzUIBegin(BzUI *ui, i32 width, i32 height);
 void bzUIEnd(BzUI *ui);
 
-// Widget construction
 
 BzUINode *bzUINodeMake(BzUI *ui, BzUIKey key, BzUIFlags flags);
 
 BzUINode *bzUIPushParent(BzUI *ui, BzUINode *node);
 BzUINode *bzUIPopParent(BzUI *ui);
 
-// Widget layout
-void bzUIPushLayout(BzUI *ui, BzUILayout layout);
-void bzUIPopLayout(BzUI *ui);
+void bzUISetParentLayout(BzUI *ui, BzUILayout layout);
 
-// Widget interaction
 BzUIInteraction bzUIGetInteraction(BzUI *ui, BzUINode *node);
 
-// UI
-
 bool bzUIButton(BzUI *ui, const char *string, BzUIStyle *style);
 
 #endif //BREEZE_UI_CORE_H
diff --git a/engine/tests/ui_test.c b/engine/tests/ui_test.c
index fc9d63c..474b34c 100644
--- a/engine/tests/ui_test.c
+++ b/engine/tests/ui_test.c
@@ -19,7 +19,7 @@ void render(float dt, int *game) {
 
 
     bzUIBegin(ui, GetScreenWidth(), GetScreenHeight());
-    bzUIPushLayout(ui, *&(BzUILayout) {
+    bzUISetParentLayout(ui, *&(BzUILayout) {
             .type = BZ_UI_LAYOUT_FLEX_BOX,
             .flags = BZ_UI_FLEX_DIR_COLUMN | BZ_UI_FLEX_ALIGN_CENTER | BZ_UI_FLEX_JUSTIFY_CENTER,
     });