From 7330c0e9d474ca8368178de143f6505794aa8655 Mon Sep 17 00:00:00 2001
From: Adoo <Adoo@outlook.com>
Date: Tue, 24 Dec 2024 22:58:15 +0800
Subject: [PATCH] =?UTF-8?q?fix(core):=20=F0=9F=90=9B=20there=20may=20be=20?=
 =?UTF-8?q?a=20missing=20update=20of=20the=20track=20ID=20in=20the=20class?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 core/src/builtin_widgets/class.rs | 60 +++++++++++++++++++++++--------
 core/src/pipe.rs                  |  8 +----
 core/src/widget_tree/widget_id.rs | 13 ++++++-
 3 files changed, 58 insertions(+), 23 deletions(-)

diff --git a/core/src/builtin_widgets/class.rs b/core/src/builtin_widgets/class.rs
index 519b95211..1c460d2e4 100644
--- a/core/src/builtin_widgets/class.rs
+++ b/core/src/builtin_widgets/class.rs
@@ -88,7 +88,7 @@ pub fn empty_cls(w: Widget) -> Widget { w }
 /// A collection comprises the implementations of the `ClassName`, offering the
 /// implementation of `Class` within its descendants.
 
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub struct Classes {
   pub(crate) store: ahash::HashMap<ClassName, ClassImpl>,
 }
@@ -332,6 +332,9 @@ impl ClassNode {
 
     // Retain the original widget ID.
     let [new, old] = tree.get_many_mut(&[n_orig, orig_id]);
+    // The "new" node is transferred to the "old" node, and the tracking ID within
+    // it needs to be updated.
+    new.update_track_id(orig_id);
     std::mem::swap(new, old);
     if new_id == n_orig {
       // If applying the class does not generate additional widgets, the original
@@ -339,18 +342,12 @@ impl ClassNode {
       new_id = orig_id;
     } else {
       n_orig.insert_after(orig_id, tree);
-      tree.remove_subtree(n_orig);
     }
+    tree.remove_subtree(n_orig);
 
     if child_id != new_id {
       // update the DynamicWidgetId out of the class node when id changed.
-      let mut v = SmallVec::new();
-      class_node.query_all_write(&QueryId::of::<TrackId>(), &mut v);
-      v.into_iter()
-        .filter_map(QueryHandle::into_ref)
-        .for_each(|handle: QueryRef<'_, TrackId>| {
-          handle.set(Some(new_id));
-        });
+      class_node.update_track_id(new_id);
     }
 
     new_id.wrap_node(tree, |node| {
@@ -388,12 +385,6 @@ impl ClassNode {
       }
     }
 
-    new_id
-      .query_all_iter::<TrackId>(tree)
-      .for_each(|wid| {
-        wid.set(Some(new_id));
-      });
-
     id_info.borrow_mut().gen_range = GenRange::Single(new_id);
     let marker = tree.dirty_marker();
     marker.mark(new_id);
@@ -721,4 +712,43 @@ mod tests {
     wnd.draw_frame();
     wnd.assert_root_size(Size::splat(120.));
   }
+
+  #[test]
+  fn fix_track_id_in_new_class() {
+    reset_test_env!();
+
+    class_names!(TRACK_ID);
+    let mut classes = initd_classes();
+    classes.insert(TRACK_ID, |w| {
+      let mut w = FatObj::new(w);
+      rdl! {
+        @Container {
+          size: Size::new(100., 100.),
+          @ $w {
+            on_performed_layout: move |e| {
+              let id = $w.track_id().get().unwrap();
+              assert!(!id.is_dropped(e.tree()));
+            }
+          }
+        }
+      }
+      .into_widget()
+    });
+
+    let (cls, w_cls) = split_value(EMPTY);
+
+    let mut wnd = TestWindow::new(fn_widget! {
+      let cls = cls.clone_watcher();
+      classes.clone().with_child(fn_widget! {
+        @Container {
+          size: Size::new(100., 100.),
+          class: pipe!(*$cls),
+        }
+      })
+    });
+
+    wnd.draw_frame();
+    *w_cls.write() = TRACK_ID;
+    wnd.draw_frame();
+  }
 }
diff --git a/core/src/pipe.rs b/core/src/pipe.rs
index 6186c2adc..4149c0305 100644
--- a/core/src/pipe.rs
+++ b/core/src/pipe.rs
@@ -739,13 +739,7 @@ impl PipeNode {
     let [old, new] = tree.get_many_mut(&[old, new_id]);
     std::mem::swap(old, new);
 
-    let mut wids = SmallVec::new();
-    new.query_all(&QueryId::of::<TrackId>(), &mut wids);
-    wids.into_iter().for_each(|wid| {
-      if let Some(wid) = wid.into_ref::<TrackId>() {
-        wid.set(Some(new_id));
-      }
-    });
+    new.update_track_id(new_id);
 
     std::mem::swap(&mut self.as_mut().data, old);
     *old = old_node;
diff --git a/core/src/widget_tree/widget_id.rs b/core/src/widget_tree/widget_id.rs
index 4a4eb55de..728cc6606 100644
--- a/core/src/widget_tree/widget_id.rs
+++ b/core/src/widget_tree/widget_id.rs
@@ -2,7 +2,7 @@ use std::{convert::Infallible, rc::Rc};
 
 use indextree::{Node, NodeId};
 use rxrust::ops::box_it::CloneableBoxOp;
-use smallvec::smallvec;
+use smallvec::{SmallVec, smallvec};
 
 use super::*;
 use crate::{
@@ -45,6 +45,17 @@ impl TrackId {
   pub(crate) fn set(&self, id: Option<WidgetId>) { *self.0.write() = id; }
 }
 
+impl dyn RenderQueryable {
+  pub(crate) fn update_track_id(&self, new_id: WidgetId) {
+    let mut handles = SmallVec::new();
+    self.query_all(&QueryId::of::<TrackId>(), &mut handles);
+    handles
+      .into_iter()
+      .filter_map(QueryHandle::into_ref::<TrackId>)
+      .for_each(move |q| q.set(Some(new_id)));
+  }
+}
+
 impl Default for TrackId {
   fn default() -> Self { Self(Stateful::new(None)) }
 }