diff --git a/src/UI/Revery_UI.re b/src/UI/Revery_UI.re
index a6df1054c..0e7d3bd57 100644
--- a/src/UI/Revery_UI.re
+++ b/src/UI/Revery_UI.re
@@ -44,7 +44,15 @@ let start =
 
   let rootNode = (new viewNode)();
   let mouseCursor: Mouse.Cursor.t = Mouse.Cursor.make();
-  let ui = UiContainer.create(window, rootNode, mouseCursor, createOptions);
+  let container = React.Container.createContainer(rootNode);
+  let ui =
+    UiContainer.create(
+      window,
+      rootNode,
+      container,
+      mouseCursor,
+      createOptions,
+    );
 
   let _ =
     Revery_Core.Event.subscribe(
diff --git a/src/UI/UiContainer.re b/src/UI/UiContainer.re
index 963b32555..36fac014e 100644
--- a/src/UI/UiContainer.re
+++ b/src/UI/UiContainer.re
@@ -23,16 +23,16 @@ module Options = {
 
 type t = {
   rootNode: ViewNode.viewNode,
-  /* container: UiReact.t, */
+  container: ref(UiReact.Container.t),
   window: Window.t,
   mouseCursor: Mouse.Cursor.t,
   options: Options.t,
 };
 
-let create = (window, rootNode, mouseCursor, options) => {
+let create = (window, rootNode, container, mouseCursor, options) => {
   window,
   rootNode,
-  /* container, */
+  container: ref(container),
   mouseCursor,
   options,
 };
diff --git a/src/UI/UiReact.re b/src/UI/UiReact.re
index e7a06d108..c81c2a02f 100644
--- a/src/UI/UiReact.re
+++ b/src/UI/UiReact.re
@@ -3,4 +3,47 @@ module Reconciler = UiReconciler.Reconciler;
 let onStale = Reconciler.onStale;
 
 module React = Brisk_reconciler.Make(Reconciler);
+
+module Container = {
+  type state = {
+    rendered: React.RenderedElement.t,
+    previousElement: React.syntheticElement,
+  };
+
+  type t = {
+    node: UiReconciler.reveryNode,
+    state: option(state),
+  };
+
+  let createContainer: UiReconciler.reveryNode => t =
+    n => {node: n, state: None};
+
+  let updateContainer: (t, React.syntheticElement) => t =
+    ({node, state}, element) => {
+      let newRendered =
+        switch (state) {
+        | None =>
+          let updates = React.RenderedElement.render(node, element);
+          React.RenderedElement.executeHostViewUpdates(updates) |> ignore;
+          updates |> React.RenderedElement.executePendingEffects;
+        | Some(s) =>
+          let nextElement =
+            React.RenderedElement.update(
+              ~previousElement=s.previousElement,
+              ~renderedElement=s.rendered,
+              element,
+            )
+            |> React.RenderedElement.flushPendingUpdates;
+          React.RenderedElement.executeHostViewUpdates(nextElement) |> ignore;
+          nextElement |> React.RenderedElement.executePendingEffects;
+        };
+
+      let ret: t = {
+        node,
+        state: Some({rendered: newRendered, previousElement: element}),
+      };
+      ret;
+    };
+};
+
 include React;
diff --git a/src/UI/UiRender.re b/src/UI/UiRender.re
index d82d1a9ef..0d42259ff 100644
--- a/src/UI/UiRender.re
+++ b/src/UI/UiRender.re
@@ -18,44 +18,13 @@ open RenderPass;
 
 let _projection = Mat4.create();
 
-let lastUpdate: ref(option(UiReact.RenderedElement.t)) = ref(None);
-let previousElement: ref(option(UiReact.syntheticElement)) = ref(None);
-
 let render = (container: UiContainer.t, component: UiReact.syntheticElement) => {
-  let {rootNode, window, options, _} = container;
+  let {rootNode, window, container, options, _} = container;
 
   AnimationTicker.tick();
 
   /* Perform reconciliation */
-  /* TODO:
-   * - Refactor this to something like a 'Container' to better manage the state
-   * - Perf: Better logic to determine if we should just update or flush pending updates - both are not always needed.
-   */
-  let latest =
-    switch (lastUpdate^, previousElement^) {
-    | (None, None) =>
-      let updates = UiReact.RenderedElement.render(rootNode, component);
-      UiReact.RenderedElement.executeHostViewUpdates(updates) |> ignore;
-      let updates = UiReact.RenderedElement.executePendingEffects(updates);
-      Some(updates);
-    | (Some(v), Some(previousElement)) =>
-      let nextElement =
-        UiReact.RenderedElement.update(
-          ~previousElement,
-          ~renderedElement=v,
-          component,
-        )
-        |> UiReact.RenderedElement.flushPendingUpdates;
-
-      UiReact.RenderedElement.executeHostViewUpdates(nextElement) |> ignore;
-
-      let ret = nextElement |> UiReact.RenderedElement.executePendingEffects;
-      Some(ret);
-    | _ => None
-    };
-
-  previousElement := Some(component);
-  lastUpdate := latest;
+  container := UiReact.Container.updateContainer(container^, component);
 
   /* Layout */
   let size = Window.getSize(window);
@@ -84,7 +53,6 @@ let render = (container: UiContainer.t, component: UiReact.syntheticElement) =>
   };
 
   /* Render */
-
   Mat4.ortho(
     _projection,
     0.0,