10 files changed, 32 insertions, 47 deletions
diff --git a/include/cru/ui/controls/Button.hpp b/include/cru/ui/controls/Button.hpp
index e8285507..5619ec89 100644
--- a/include/cru/ui/controls/Button.hpp
+++ b/include/cru/ui/controls/Button.hpp
@@ -28,9 +28,6 @@ class Button : public ContentControl {
   const ButtonStyle& GetStyle() const { return style_; }
   void SetStyle(ButtonStyle style);
 
- protected:
-  void OnChildChanged(Control* old_child, Control* new_child) override;
-
  private:
   std::unique_ptr<render::BorderRenderObject> render_object_{};
 
diff --git a/include/cru/ui/controls/Container.hpp b/include/cru/ui/controls/Container.hpp
index d9cb8aec..18958837 100644
--- a/include/cru/ui/controls/Container.hpp
+++ b/include/cru/ui/controls/Container.hpp
@@ -19,9 +19,6 @@ class Container : public ContentControl {
 
   render::RenderObject* GetRenderObject() const override;
 
- protected:
-  void OnChildChanged(Control* old_child, Control* new_child) override;
-
  private:
   std::unique_ptr<render::BorderRenderObject> render_object_;
 };
diff --git a/include/cru/ui/controls/ContentControl.hpp b/include/cru/ui/controls/ContentControl.hpp
index 47720a87..1bdaf7e4 100644
--- a/include/cru/ui/controls/ContentControl.hpp
+++ b/include/cru/ui/controls/ContentControl.hpp
@@ -1,6 +1,8 @@
 #pragma once
 #include "Control.hpp"
 
+#include "cru/ui/render/RenderObject.hpp"
+
 namespace cru::ui::controls {
 class ContentControl : public Control {
  protected:
@@ -19,8 +21,18 @@ class ContentControl : public Control {
  protected:
   virtual void OnChildChanged(Control* old_child, Control* new_child);
 
+  render::RenderObject* GetContainerRenderObject() const {
+    return container_render_object_;
+  }
+  void SetContainerRenderObject(render::RenderObject* ro) {
+    container_render_object_ = ro;
+  }
+
  private:
   using Control::AddChild;
   using Control::RemoveChild;
+
+ private:
+  render::RenderObject* container_render_object_ = nullptr;
 };
 }  // namespace cru::ui::controls
diff --git a/include/cru/ui/controls/FlexLayout.hpp b/include/cru/ui/controls/FlexLayout.hpp
index a6c6a40c..4f6abfdb 100644
--- a/include/cru/ui/controls/FlexLayout.hpp
+++ b/include/cru/ui/controls/FlexLayout.hpp
@@ -34,10 +34,6 @@ class FlexLayout : public LayoutControl {
   FlexChildLayoutData GetChildLayoutData(Control* control);
   void SetChildLayoutData(Control* control, FlexChildLayoutData data);
 
- protected:
-  void OnAddChild(Control* child, Index position) override;
-  void OnRemoveChild(Control* child, Index position) override;
-
  private:
   std::shared_ptr<render::FlexLayoutRenderObject> render_object_;
 };
diff --git a/include/cru/ui/controls/StackLayout.hpp b/include/cru/ui/controls/StackLayout.hpp
index 373b4681..aa9440c2 100644
--- a/include/cru/ui/controls/StackLayout.hpp
+++ b/include/cru/ui/controls/StackLayout.hpp
@@ -21,10 +21,6 @@ class StackLayout : public LayoutControl {
 
   render::RenderObject* GetRenderObject() const override;
 
- protected:
-  void OnAddChild(Control* child, Index position) override;
-  void OnRemoveChild(Control* child, Index position) override;
-
  private:
   std::shared_ptr<render::StackLayoutRenderObject> render_object_;
 };
diff --git a/src/ui/controls/Button.cpp b/src/ui/controls/Button.cpp
index b7407ec2..39c4b961 100644
--- a/src/ui/controls/Button.cpp
+++ b/src/ui/controls/Button.cpp
@@ -31,6 +31,8 @@ Button::Button() : click_detector_(this) {
 
   render_object_ = std::make_unique<render::BorderRenderObject>();
   render_object_->SetAttachedControl(this);
+  SetContainerRenderObject(render_object_.get());
+
   Set(render_object_.get(), style_.normal);
   render_object_->SetBorderEnabled(true);
 
@@ -62,11 +64,4 @@ Button::~Button() = default;
 render::RenderObject* Button::GetRenderObject() const {
   return render_object_.get();
 }
-
-void Button::OnChildChanged(Control* old_child, Control* new_child) {
-  if (old_child != nullptr) render_object_->RemoveChild(0);
-  if (new_child != nullptr)
-    render_object_->AddChild(new_child->GetRenderObject(), 0);
-}
-
 }  // namespace cru::ui::controls
diff --git a/src/ui/controls/Container.cpp b/src/ui/controls/Container.cpp
index 8b15c566..30129f64 100644
--- a/src/ui/controls/Container.cpp
+++ b/src/ui/controls/Container.cpp
@@ -2,17 +2,19 @@
 
 #include "cru/platform/graphics/Factory.hpp"
 #include "cru/ui/render/BorderRenderObject.hpp"
+#include "cru/ui/render/RenderObject.hpp"
 
 namespace cru::ui::controls {
 Container::Container() {
   render_object_ = std::make_unique<render::BorderRenderObject>();
   render_object_->SetBorderEnabled(false);
+  render_object_->SetAttachedControl(this);
+  SetContainerRenderObject(render_object_.get());
 }
 
 Container::~Container() = default;
 
-void Container::OnChildChanged(Control*, Control* new_child) {
-  render_object_->RemoveChild(0);
-  render_object_->AddChild(new_child->GetRenderObject(), 0);
+render::RenderObject* Container::GetRenderObject() const {
+  return render_object_.get();
 }
 }  // namespace cru::ui::controls
diff --git a/src/ui/controls/ContentControl.cpp b/src/ui/controls/ContentControl.cpp
index 653882c0..8c6f0b00 100644
--- a/src/ui/controls/ContentControl.cpp
+++ b/src/ui/controls/ContentControl.cpp
@@ -19,7 +19,13 @@ void ContentControl::SetChild(Control* child) {
 }
 
 void ContentControl::OnChildChanged(Control* old_child, Control* new_child) {
-  CRU_UNUSED(old_child)
-  CRU_UNUSED(new_child)
+  if (container_render_object_) {
+    if (old_child) {
+      container_render_object_->RemoveChild(0);
+    }
+    if (new_child) {
+      container_render_object_->AddChild(new_child->GetRenderObject(), 0);
+    }
+  }
 }
 }  // namespace cru::ui::controls
diff --git a/src/ui/controls/FlexLayout.cpp b/src/ui/controls/FlexLayout.cpp
index 05f6999f..e390241f 100644
--- a/src/ui/controls/FlexLayout.cpp
+++ b/src/ui/controls/FlexLayout.cpp
@@ -8,6 +8,7 @@ using render::FlexLayoutRenderObject;
 FlexLayout::FlexLayout() {
   render_object_.reset(new FlexLayoutRenderObject());
   render_object_->SetAttachedControl(this);
+  SetContainerRenderObject(render_object_.get());
 }
 
 FlexLayout::~FlexLayout() = default;
@@ -68,14 +69,4 @@ void FlexLayout::SetItemCrossAlign(FlexCrossAlignment alignment) {
   if (alignment == GetItemCrossAlign()) return;
   render_object_->SetItemCrossAlign(alignment);
 }
-
-void FlexLayout::OnAddChild(Control* child, const Index position) {
-  render_object_->AddChild(child->GetRenderObject(), position);
-}
-
-void FlexLayout::OnRemoveChild(Control* child, const Index position) {
-  CRU_UNUSED(child)
-
-  render_object_->RemoveChild(position);
-}
 }  // namespace cru::ui::controls
diff --git a/src/ui/controls/StackLayout.cpp b/src/ui/controls/StackLayout.cpp
index ce500b79..89968571 100644
--- a/src/ui/controls/StackLayout.cpp
+++ b/src/ui/controls/StackLayout.cpp
@@ -1,12 +1,15 @@
 #include "cru/ui/controls/StackLayout.hpp"
+#include <memory>
 
 #include "cru/ui/render/StackLayoutRenderObject.hpp"
 
 namespace cru::ui::controls {
 using render::StackLayoutRenderObject;
 
-StackLayout::StackLayout() : render_object_(new StackLayoutRenderObject()) {
+StackLayout::StackLayout() {
+  render_object_ = std::make_unique<StackLayoutRenderObject>();
   render_object_->SetAttachedControl(this);
+  SetContainerRenderObject(render_object_.get());
 }
 
 StackLayout::~StackLayout() = default;
@@ -14,14 +17,4 @@ StackLayout::~StackLayout() = default;
 render::RenderObject* StackLayout::GetRenderObject() const {
   return render_object_.get();
 }
-
-void StackLayout::OnAddChild(Control* child, const Index position) {
-  render_object_->AddChild(child->GetRenderObject(), position);
-}
-
-void StackLayout::OnRemoveChild(Control* child, const Index position) {
-  CRU_UNUSED(child)
-
-  render_object_->RemoveChild(position);
-}
 }  // namespace cru::ui::controls