diff --git a/src/ui/ui.c b/src/ui/ui.c
index f8fff26..af4aa8e 100644
--- a/src/ui/ui.c
+++ b/src/ui/ui.c
@@ -616,7 +616,6 @@ fn ui_box_t *ui__label(ui_code_loc_t loc, char *str, ...) {
     return box;
 }
 
-#define ui_tree_table() defer_block(ui_tree_table_begin(UILOC), ui_tree_table_end())
 fn void ui_tree_table_begin(ui_code_loc_t loc) {
     ui_box_t *box = ui_box(.loc = loc, .rect = ui_next_rect(ui_top_lop(), ui_top_rectp(), v2f32_null), .flags = { .draw_border = true, .children_sum_y = true });
     ui_push_top(box);
@@ -655,40 +654,54 @@ fn ui_signal_t ui_tree_table_push_expandable(ui_code_loc_t loc, char *str, ...)
     return button;
 }
 
-fn ui_box_t *ui_get_prev_box(ui_box_t *box, b32 (*match)(ui_box_t *)) {
-    for (ui_box_t *it = box; it;) {
-        if (it->last) {
-            it = it->last;
-            if (match(it)) return it;
-            else                  continue;
-        }
-        if (it->prev) {
-            it = it->prev;
-            if (match(it)) return it;
-            else                  continue;
-        }
-
-        it = it->parent;
-        if (it) it = it->prev;
+fn ui_box_t *ui_get_next_box_ex(ui_box_t *box) {
+    if (box->first) {
+        return box->first;
     }
-    return box;
+
+    if (box->next) {
+        return box->next;
+    }
+
+    for (ui_box_t *it = box->parent; it; it = it->parent) {
+        if (it->next != NULL) {
+            return it->next;
+        }
+    }
+
+    return NULL;
+}
+
+fn ui_box_t *ui_get_prev_box_ex(ui_box_t *box) {
+    if (box->last) {
+        return box->last;
+    }
+
+    if (box->prev) {
+        return box->prev;
+    }
+
+    for (ui_box_t *it = box->parent; it; it = it->parent) {
+        if (it->prev != NULL) {
+            return it->prev;
+        }
+    }
+
+    return NULL;
 }
 
 fn ui_box_t *ui_get_next_box(ui_box_t *box, b32 (*match)(ui_box_t *)) {
     for (ui_box_t *it = box; it;) {
-        if (it->first) {
-            it = it->first;
-            if (match(it)) return it;
-            else continue;
-        }
-        if (it->next) {
-            it = it->next;
-            if (match(it)) return it;
-            else continue;
-        }
+        it = ui_get_next_box_ex(it);
+        if (it && match(it)) return it;
+    }
+    return box;
+}
 
-        it = it->parent;
-        if (it) it = it->next;
+fn ui_box_t *ui_get_prev_box(ui_box_t *box, b32 (*match)(ui_box_t *)) {
+    for (ui_box_t *it = box; it;) {
+        it = ui_get_prev_box_ex(it);
+        if (it && match(it)) return it;
     }
     return box;
 }