diff --git a/apis/libtui.lua b/apis/libtui.lua
index c74e034..0657997 100644
--- a/apis/libtui.lua
+++ b/apis/libtui.lua
@@ -1,4 +1,4 @@
-local _VERSION = '0.1.0';
+local _VERSION = '0.1.1';
 
 local NODE_TEXT = 'text';
 local NODE_BUTTON = 'button';
@@ -260,26 +260,45 @@ local function buttonLabel(props)
   return '[ ' .. tostring(props.text or props.label or '') .. ' ]';
 end
 
+local function flexOf(child)
+  local f = (child.props or {}).flex;
+  if f == true then
+    return 1;
+  end
+  if type(f) == 'number' and f > 0 then
+    return f;
+  end
+  return nil;
+end
+
 local naturalSize;
 
 local function childrenNaturalSize(children, direction, gap)
   local width = 0;
   local height = 0;
+  local mainCount = 0;
 
-  for index, child in ipairs(children) do
+  for _, child in ipairs(children) do
+    local hasFlex = flexOf(child) ~= nil;
     local childWidth, childHeight = naturalSize(child);
 
     if direction == 'row' then
-      width = width + childWidth;
       height = math.max(height, childHeight);
-      if index > 1 then
-        width = width + gap;
+      if not hasFlex then
+        if mainCount > 0 then
+          width = width + gap;
+        end
+        width = width + childWidth;
+        mainCount = mainCount + 1;
       end
     else
       width = math.max(width, childWidth);
-      height = height + childHeight;
-      if index > 1 then
-        height = height + gap;
+      if not hasFlex then
+        if mainCount > 0 then
+          height = height + gap;
+        end
+        height = height + childHeight;
+        mainCount = mainCount + 1;
       end
     end
   end
@@ -336,41 +355,38 @@ local function layoutChildren(children, rect, direction, gap)
   local flexSize = 0;
   local axisSize = direction == 'row' and rect.w or rect.h;
   local layouts = {};
-  local remaining;
+  local lastFlexIndex = nil;
+  local flexes = {};
 
-  for _, child in ipairs(children) do
-    local props = child.props or {};
-    if props.flex then
-      flexSize = flexSize + props.flex;
+  for index, child in ipairs(children) do
+    local flex = flexOf(child);
+    flexes[index] = flex;
+    if flex then
+      flexSize = flexSize + flex;
+      lastFlexIndex = index;
     else
       fixedSize = fixedSize + childAxisSize(child, direction);
     end
   end
 
-  remaining = axisSize - fixedSize - math.max(#children - 1, 0) * gap;
+  local remaining = axisSize - fixedSize - math.max(#children - 1, 0) * gap;
   if remaining < 0 then
     remaining = 0;
   end
 
   local cursor = direction == 'row' and rect.x or rect.y;
-  local lastFlexIndex = nil;
   local usedFlexSize = 0;
 
-  for index, child in ipairs(children) do
-    if (child.props or {}).flex then
-      lastFlexIndex = index;
-    end
-  end
-
   for index, child in ipairs(children) do
     local props = child.props or {};
+    local flex = flexes[index];
     local size;
 
-    if props.flex then
+    if flex then
       if index == lastFlexIndex then
         size = remaining - usedFlexSize;
       else
-        size = math.floor(remaining * props.flex / flexSize);
+        size = math.floor(remaining * flex / flexSize);
         usedFlexSize = usedFlexSize + size;
       end
     else
@@ -548,12 +564,27 @@ local function createTui(eventloop)
   end
 
   function api.rerender()
+    if root == nil then
+      return;
+    end
     safeRedraw();
   end
 
+  local function restoreTerminal()
+    if not previousState then
+      return;
+    end
+    term.setTextColor(previousState.color);
+    term.setBackgroundColor(previousState.bgColor);
+    term.clear();
+    term.setCursorPos(previousState.cursorX, previousState.cursorY);
+    term.setCursorBlink(previousState.cursorBlink);
+    previousState = nil;
+  end
+
   function api.render(nextRoot)
-    root = nextRoot;
     finalEvent = nil;
+    root = nextRoot;
 
     previousState = {
       color = term.getTextColor(),
@@ -572,13 +603,7 @@ local function createTui(eventloop)
     end
 
     eventloop.onStart(safeRedraw);
-    eventloop.onStop(function()
-      term.setTextColor(previousState.color);
-      term.setBackgroundColor(previousState.bgColor);
-      term.clear();
-      term.setCursorPos(previousState.cursorX, previousState.cursorY);
-      term.setCursorBlink(previousState.cursorBlink);
-    end);
+    eventloop.onStop(restoreTerminal);
 
     eventloop.register('mouse_click', function(button, x, y)
       for index = #clickables, 1, -1 do
@@ -606,6 +631,7 @@ local function createTui(eventloop)
 
     local ok, reason = pcall(eventloop.startLoop);
     if not ok then
+      pcall(restoreTerminal);
       finalEvent = createErrorEvent(reason);
     end
 
@@ -616,18 +642,17 @@ local function createTui(eventloop)
   end
 
   api.Text = makeText;
+  api.text = makeText;
   api.Button = makeButton;
+  api.button = makeButton;
   api.Box = makeBox;
+  api.box = makeBox;
   api.List = makeList;
+  api.list = makeList;
   api.Fragment = makeFragment;
   api.version = _VERSION;
   api.eventloop = eventloop;
 
-  api.text = makeText;
-  api.button = makeButton;
-  api.box = makeBox;
-  api.list = makeList;
-
   return api;
 end