5 files changed, 20 insertions, 33 deletions
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