7 files changed, 52 insertions, 17 deletions

diff --git a/include/cru/ui/controls/Control.h b/include/cru/ui/controls/Control.h
index 97096e35..4dad4549 100644
--- a/include/cru/ui/controls/Control.h
+++ b/include/cru/ui/controls/Control.h
@@ -17,6 +17,8 @@ namespace cru::ui::controls {
  * former one is even not used.
  */
 class CRU_UI_API Control : public Object {
+  friend class RootControl;
+
  protected:
   Control();
 
@@ -32,7 +34,7 @@ class CRU_UI_API Control : public Object {
 
   //*************** region: tree ***************
  public:
-  virtual host::WindowHost* GetWindowHost() const;
+  host::WindowHost* GetWindowHost() const { return window_host_; }
 
   Control* GetParent() const { return parent_; }
   void SetParent(Control* parent);
@@ -143,12 +145,22 @@ class CRU_UI_API Control : public Object {
  protected:
   virtual void OnParentChanged(Control* old_parent, Control* new_parent) {}
 
+  virtual void OnWindowHostChanged(host::WindowHost* old_host,
+                                   host::WindowHost* new_host) {}
+
  protected:
   virtual void OnMouseHoverChange(bool newHover) { CRU_UNUSED(newHover) }
 
  private:
+  void OnParentChangedCore(Control* old_parent, Control* new_parent);
+  void OnWindowHostChangedCore(host::WindowHost* old_host,
+                               host::WindowHost* new_host);
+
+ private:
   Control* parent_ = nullptr;
 
+  host::WindowHost* window_host_ = nullptr;
+
  private:
   bool is_mouse_over_ = false;
 
diff --git a/include/cru/ui/controls/RootControl.h b/include/cru/ui/controls/RootControl.h
index a1878836..851b4db0 100644
--- a/include/cru/ui/controls/RootControl.h
+++ b/include/cru/ui/controls/RootControl.h
@@ -18,10 +18,6 @@ class CRU_UI_API RootControl
   ~RootControl() override;
 
  public:
-  host::WindowHost* GetWindowHost() const override {
-    return window_host_.get();
-  }
-
   platform::gui::INativeWindow* GetNativeWindow();
 
  protected:
diff --git a/include/cru/ui/controls/SingleChildControl.h b/include/cru/ui/controls/SingleChildControl.h
index 6bf150c0..e733c825 100644
--- a/include/cru/ui/controls/SingleChildControl.h
+++ b/include/cru/ui/controls/SingleChildControl.h
@@ -31,7 +31,8 @@ class SingleChildControl : public Control {
       child->SetParent(this);
     }
 
-    container_render_object_->SetChild(child->GetRenderObject());
+    container_render_object_->SetChild(
+        child == nullptr ? nullptr : child->GetRenderObject());
   }
 
   render::RenderObject* GetRenderObject() const override {
diff --git a/include/cru/ui/host/WindowHost.h b/include/cru/ui/host/WindowHost.h
index 1de915e0..9fe46ac0 100644
--- a/include/cru/ui/host/WindowHost.h
+++ b/include/cru/ui/host/WindowHost.h
@@ -18,6 +18,7 @@ struct AfterLayoutEventArgs {};
 
 // The bridge between control tree and native window.
 class CRU_UI_API WindowHost : public Object {
+  friend controls::Control;
   CRU_DEFINE_CLASS_LOG_TAG(u"cru::ui::host::WindowHost")
 
  public:
diff --git a/src/ui/controls/Control.cpp b/src/ui/controls/Control.cpp
index a66b605f..43faba69 100644
--- a/src/ui/controls/Control.cpp
+++ b/src/ui/controls/Control.cpp
@@ -27,21 +27,13 @@ Control::Control() {
   });
 }
 
-Control::~Control() { ReleaseMouse(); }
-
-host::WindowHost* Control::GetWindowHost() const {
-  auto parent = GetParent();
-  if (parent) {
-    return parent->GetWindowHost();
-  }
-  return nullptr;
-}
+Control::~Control() { RemoveFromParent(); }
 
 void Control::SetParent(Control* parent) {
   if (parent_ == parent) return;
   auto old_parent = parent_;
   parent_ = parent;
-  OnParentChanged(old_parent, parent);
+  OnParentChangedCore(old_parent, parent);
 }
 
 void Control::RemoveFromParent() {
@@ -109,4 +101,34 @@ void Control::SetCursor(std::shared_ptr<ICursor> cursor) {
 std::shared_ptr<style::StyleRuleSet> Control::GetStyleRuleSet() {
   return style_rule_set_;
 }
+
+void Control::OnParentChangedCore(Control* old_parent, Control* new_parent) {
+  auto new_window_host =
+      new_parent == nullptr ? nullptr : new_parent->GetWindowHost();
+  if (window_host_ != new_window_host) {
+    auto old_host = window_host_;
+    window_host_ = new_window_host;
+    OnWindowHostChangedCore(old_host, new_window_host);
+  }
+
+  OnParentChanged(old_parent, new_parent);
+}
+
+void Control::OnWindowHostChangedCore(host::WindowHost* old_host,
+                                      host::WindowHost* new_host) {
+  if (old_host != nullptr) {
+    if (old_host->GetMouseCaptureControl() == this) {
+      old_host->CaptureMouseFor(nullptr);
+    }
+    if (old_host->GetMouseHoverControl() == this) {
+      old_host->mouse_hover_control_ = nullptr;
+    }
+  }
+
+  ForEachChild([old_host, new_host](Control* child) {
+    child->window_host_ = new_host;
+    child->OnWindowHostChangedCore(old_host, new_host);
+  });
+  OnWindowHostChanged(old_host, new_host);
+}
 }  // namespace cru::ui::controls
diff --git a/src/ui/controls/RootControl.cpp b/src/ui/controls/RootControl.cpp
index 07647024..609223a1 100644
--- a/src/ui/controls/RootControl.cpp
+++ b/src/ui/controls/RootControl.cpp
@@ -15,6 +15,9 @@ RootControl::RootControl(Control* attached_control)
   GetContainerRenderObject()->SetDefaultHorizontalAlignment(Alignment::Stretch);
   GetContainerRenderObject()->SetDefaultVertialAlignment(Alignment::Stretch);
   window_host_ = std::make_unique<host::WindowHost>(this);
+
+  Control::window_host_ = this->window_host_.get();
+
   window_host_->SetLayoutPreferToFillWindow(true);
 }
 
diff --git a/src/ui/render/SingleChildRenderObject.cpp b/src/ui/render/SingleChildRenderObject.cpp
index 7b10b343..d62a233e 100644
--- a/src/ui/render/SingleChildRenderObject.cpp
+++ b/src/ui/render/SingleChildRenderObject.cpp
@@ -2,7 +2,7 @@
 
 namespace cru::ui::render {
 void SingleChildRenderObject::SetChild(RenderObject *new_child) {
-  assert(new_child->GetParent() == nullptr);
+  assert(new_child == nullptr || new_child->GetParent() == nullptr);
   if (child_ == new_child) return;
   auto old_child = child_;
   if (child_) {