Merge pull request #940 from Demonthos/make-widgets-cross-framework

packages/native-core/src/custom_element.rs

@@ -0,0 +1,180 @@
+//! A custom element is a controlled element that renders to a shadow DOM. This allows you to create elements that act like widgets without relying on a specific framework.
+//!
+//! Each custom element is registered with a element name and namespace with [`RealDom::register_custom_element`] or [`RealDom::register_custom_element_with_factory`]. Once registered, they will be created automatically when the element is added to the DOM.
+
+// Used in doc links
+#[allow(unused)]
+use crate::real_dom::RealDom;
+
+use std::sync::{Arc, RwLock};
+
+use rustc_hash::FxHashMap;
+use shipyard::Component;
+
+use crate::{
+    node::{FromAnyValue, NodeType},
+    node_ref::AttributeMask,
+    prelude::{NodeImmutable, NodeMut},
+    tree::TreeMut,
+    NodeId,
+};
+
+pub(crate) struct CustomElementRegistry<V: FromAnyValue + Send + Sync> {
+    builders: FxHashMap<(&'static str, Option<&'static str>), CustomElementBuilder<V>>,
+}
+
+impl<V: FromAnyValue + Send + Sync> Default for CustomElementRegistry<V> {
+    fn default() -> Self {
+        Self {
+            builders: FxHashMap::default(),
+        }
+    }
+}
+
+impl<V: FromAnyValue + Send + Sync> CustomElementRegistry<V> {
+    pub fn register<F, U>(&mut self)
+    where
+        F: CustomElementFactory<U, V>,
+        U: CustomElementUpdater<V>,
+    {
+        self.builders.insert(
+            (F::NAME, F::NAMESPACE),
+            CustomElementBuilder {
+                create: |node| Box::new(F::create(node)),
+            },
+        );
+    }
+
+    pub fn add_shadow_dom(&self, mut node: NodeMut<V>) {
+        let element_tag = if let NodeType::Element(el) = &*node.node_type() {
+            Some((el.tag.clone(), el.namespace.clone()))
+        } else {
+            None
+        };
+        if let Some((tag, ns)) = element_tag {
+            if let Some(builder) = self.builders.get(&(tag.as_str(), ns.as_deref())) {
+                let boxed_custom_element = { (builder.create)(node.reborrow()) };
+
+                let shadow_roots = boxed_custom_element.roots();
+
+                let light_id = node.id();
+                node.real_dom_mut().tree_mut().create_subtree(
+                    light_id,
+                    shadow_roots,
+                    boxed_custom_element.slot(),
+                );
+
+                let boxed_custom_element = CustomElementManager {
+                    inner: Arc::new(RwLock::new(boxed_custom_element)),
+                };
+
+                node.insert(boxed_custom_element);
+            }
+        }
+    }
+}
+
+struct CustomElementBuilder<V: FromAnyValue + Send + Sync> {
+    create: fn(NodeMut<V>) -> Box<dyn CustomElementUpdater<V>>,
+}
+
+/// A controlled element that renders to a shadow DOM.
+///
+/// Register with [`RealDom::register_custom_element`]
+///
+/// This is a simplified custom element trait for elements that can create themselves. For more granular control, implement [`CustomElementFactory`] and [`CustomElementUpdater`] instead.
+pub trait CustomElement<V: FromAnyValue + Send + Sync = ()>: Send + Sync + 'static {
+    /// The tag of the element
+    const NAME: &'static str;
+
+    /// The namespace of the element
+    const NAMESPACE: Option<&'static str> = None;
+
+    /// Create a new element *without mounting* it.
+    /// The node passed in is the light DOM node. The element should not modify the light DOM node, but it can get the [`NodeMut::real_dom_mut`] from the node to create new nodes.
+    fn create(light_root: NodeMut<V>) -> Self;
+
+    /// The root node of the custom element. These roots must be not change once the element is created.
+    fn roots(&self) -> Vec<NodeId>;
+
+    /// The slot to render children of the element into. The slot must be not change once the element is created.
+    fn slot(&self) -> Option<NodeId> {
+        None
+    }
+
+    /// Update the custom element's shadow tree with the new attributes.
+    /// Called when the attributes of the custom element are changed.
+    fn attributes_changed(&mut self, light_node: NodeMut<V>, attributes: &AttributeMask);
+}
+
+/// A factory for creating custom elements
+///
+/// Register with [`RealDom::register_custom_element_with_factory`]
+pub trait CustomElementFactory<W: CustomElementUpdater<V>, V: FromAnyValue + Send + Sync = ()>:
+    Send + Sync + 'static
+{
+    /// The tag of the element
+    const NAME: &'static str;
+
+    /// The namespace of the element
+    const NAMESPACE: Option<&'static str> = None;
+
+    /// Create a new element *without mounting* it.
+    /// The node passed in is the light DOM node. The element should not modify the light DOM node, but it can get the [`NodeMut::real_dom_mut`] from the node to create new nodes.
+    fn create(dom: NodeMut<V>) -> W;
+}
+
+impl<W: CustomElement<V>, V: FromAnyValue + Send + Sync> CustomElementFactory<W, V> for W {
+    const NAME: &'static str = W::NAME;
+
+    const NAMESPACE: Option<&'static str> = W::NAMESPACE;
+
+    fn create(node: NodeMut<V>) -> Self {
+        Self::create(node)
+    }
+}
+
+/// A trait for updating custom elements
+pub trait CustomElementUpdater<V: FromAnyValue + Send + Sync = ()>: Send + Sync + 'static {
+    /// Update the custom element's shadow tree with the new attributes.
+    /// Called when the attributes of the custom element are changed.
+    fn attributes_changed(&mut self, light_root: NodeMut<V>, attributes: &AttributeMask);
+
+    /// The root node of the custom element. These roots must be not change once the element is created.
+    fn roots(&self) -> Vec<NodeId>;
+
+    /// The slot to render children of the element into. The slot must be not change once the element is created.
+    fn slot(&self) -> Option<NodeId> {
+        None
+    }
+}
+
+impl<W: CustomElement<V>, V: FromAnyValue + Send + Sync> CustomElementUpdater<V> for W {
+    fn attributes_changed(&mut self, light_root: NodeMut<V>, attributes: &AttributeMask) {
+        self.attributes_changed(light_root, attributes);
+    }
+
+    fn roots(&self) -> Vec<NodeId> {
+        self.roots()
+    }
+
+    fn slot(&self) -> Option<NodeId> {
+        self.slot()
+    }
+}
+
+/// A dynamic trait object wrapper for [`CustomElementUpdater`]
+#[derive(Component, Clone)]
+pub(crate) struct CustomElementManager<V: FromAnyValue = ()> {
+    inner: Arc<RwLock<Box<dyn CustomElementUpdater<V>>>>,
+}
+
+impl<V: FromAnyValue + Send + Sync> CustomElementManager<V> {
+    /// Update the custom element based on attributes changed.
+    pub fn on_attributes_changed(&self, light_root: NodeMut<V>, attributes: &AttributeMask) {
+        self.inner
+            .write()
+            .unwrap()
+            .attributes_changed(light_root, attributes);
+    }
+}

packages/native-core/src/lib.rs

@@ -7,6 +7,7 @@ use std::hash::BuildHasherDefault;
 use node_ref::NodeMask;
 use rustc_hash::FxHasher;
 
+pub mod custom_element;
 #[cfg(feature = "dioxus")]
 pub mod dioxus;
 #[cfg(feature = "layout-attributes")]
@@ -18,6 +19,7 @@ mod passes;
 pub mod real_dom;
 pub mod tree;
 pub mod utils;
+
 pub use shipyard::EntityId as NodeId;
 
 pub mod exports {

packages/native-core/src/node_ref.rs

@@ -60,7 +60,7 @@ impl<'a, V: FromAnyValue> NodeView<'a, V> {
             NodeType::Element(ElementNode { attributes, .. }) => Some(
                 attributes
                     .iter()
-                    .filter(move |(attr, _)| self.mask.attritutes.contains_attribute(&attr.name))
+                    .filter(move |(attr, _)| self.mask.attritutes.contains(&attr.name))
                     .map(|(attr, val)| OwnedAttributeView {
                         attribute: attr,
                         value: val,
@@ -107,8 +107,8 @@ pub enum AttributeMask {
 }
 
 impl AttributeMask {
-    /// Check if the given attribute is visible
-    fn contains_attribute(&self, attr: &str) -> bool {
+    /// Check if the mask contains the given attribute
+    pub fn contains(&self, attr: &str) -> bool {
         match self {
             AttributeMask::All => true,
             AttributeMask::Some(attrs) => attrs.contains(attr),
@@ -187,25 +187,50 @@ impl NodeMask {
         self.attritutes = self.attritutes.union(&attributes);
     }
 
+    /// Get the mask for the attributes
+    pub fn attributes(&self) -> &AttributeMask {
+        &self.attritutes
+    }
+
     /// Set the mask to view the tag
     pub fn set_tag(&mut self) {
         self.tag = true;
     }
 
+    /// Get the mask for the tag
+    pub fn tag(&self) -> bool {
+        self.tag
+    }
+
     /// Set the mask to view the namespace
     pub fn set_namespace(&mut self) {
         self.namespace = true;
     }
 
+    /// Get the mask for the namespace
+    pub fn namespace(&self) -> bool {
+        self.namespace
+    }
+
     /// Set the mask to view the text
     pub fn set_text(&mut self) {
         self.text = true;
     }
 
+    /// Get the mask for the text
+    pub fn text(&self) -> bool {
+        self.text
+    }
+
     /// Set the mask to view the listeners
     pub fn set_listeners(&mut self) {
         self.listeners = true;
     }
+
+    /// Get the mask for the listeners
+    pub fn listeners(&self) -> bool {
+        self.listeners
+    }
 }
 
 /// A builder for a mask that controls what attributes are visible.

packages/native-core/src/node_watcher.rs

@@ -1,17 +1,19 @@
 //! Helpers for watching for changes in the DOM tree.
 
-use crate::{node::FromAnyValue, prelude::*};
+use crate::{node::FromAnyValue, node_ref::AttributeMask, prelude::*};
 
 /// A trait for watching for changes in the DOM tree.
 pub trait NodeWatcher<V: FromAnyValue + Send + Sync> {
     /// Called after a node is added to the tree.
-    fn on_node_added(&self, _node: NodeMut<V>) {}
+    fn on_node_added(&mut self, _node: NodeMut<V>) {}
     /// Called before a node is removed from the tree.
-    fn on_node_removed(&self, _node: NodeMut<V>) {}
+    fn on_node_removed(&mut self, _node: NodeMut<V>) {}
     /// Called after a node is moved to a new parent.
-    fn on_node_moved(&self, _node: NodeMut<V>) {}
-    // /// Called after the text content of a node is changed.
-    // fn on_text_changed(&self, _node: NodeMut<V>) {}
-    // /// Called after an attribute of an element is changed.
-    // fn on_attribute_changed(&self, _node: NodeMut<V>, attribute: &str) {}
+    fn on_node_moved(&mut self, _node: NodeMut<V>) {}
+}
+
+/// A trait for watching for changes to attributes of an element.
+pub trait AttributeWatcher<V: FromAnyValue + Send + Sync> {
+    /// Called before update_state is called on the RealDom
+    fn on_attributes_changed(&self, _node: NodeMut<V>, _attributes: &AttributeMask) {}
 }

packages/native-core/src/passes.rs

@@ -104,6 +104,9 @@ pub trait State<V: FromAnyValue + Send + Sync = ()>: Any + Send + Sync {
     /// This is a mask of what aspects of the node are required to update this state
     const NODE_MASK: NodeMaskBuilder<'static>;
 
+    /// Does the state traverse into the shadow dom or pass over it. This should be true for layout and false for styles
+    const TRAVERSE_SHADOW_DOM: bool = false;
+
     /// Update this state in a node, returns if the state was updated
     fn update<'a>(
         &mut self,
@@ -150,6 +153,7 @@ pub trait State<V: FromAnyValue + Send + Sync = ()>: Any + Send + Sync {
             dependants: Default::default(),
             mask: node_mask,
             pass_direction: pass_direction::<V, Self>(),
+            enter_shadow_dom: Self::TRAVERSE_SHADOW_DOM,
             workload: Self::workload_system,
             phantom: PhantomData,
         }
@@ -228,28 +232,42 @@ pub fn run_pass<V: FromAnyValue + Send + Sync>(
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct Dependant {
+    pub(crate) type_id: TypeId,
+    pub(crate) enter_shadow_dom: bool,
+}
+
 /// The states that depend on this state
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 pub struct Dependants {
     /// The states in the parent direction that should be invalidated when this state is invalidated
-    pub parent: Vec<TypeId>,
+    pub(crate) parent: Vec<Dependant>,
     /// The states in the child direction that should be invalidated when this state is invalidated
-    pub child: Vec<TypeId>,
+    pub(crate) child: Vec<Dependant>,
     /// The states in the node direction that should be invalidated when this state is invalidated
-    pub node: Vec<TypeId>,
+    pub(crate) node: Vec<TypeId>,
 }
 
 impl Dependants {
     fn mark_dirty(&self, dirty: &DirtyNodeStates, id: NodeId, tree: &impl TreeRef, height: u16) {
-        for dependant in &self.child {
-            for id in tree.children_ids(id) {
-                dirty.insert(*dependant, id, height + 1);
+        for &Dependant {
+            type_id,
+            enter_shadow_dom,
+        } in &self.child
+        {
+            for id in tree.children_ids_advanced(id, enter_shadow_dom) {
+                dirty.insert(type_id, id, height + 1);
             }
         }
 
-        for dependant in &self.parent {
-            if let Some(id) = tree.parent_id(id) {
-                dirty.insert(*dependant, id, height - 1);
+        for &Dependant {
+            type_id,
+            enter_shadow_dom,
+        } in &self.parent
+        {
+            if let Some(id) = tree.parent_id_advanced(id, enter_shadow_dom) {
+                dirty.insert(type_id, id, height - 1);
             }
         }
 
@@ -269,6 +287,7 @@ pub struct TypeErasedState<V: FromAnyValue + Send = ()> {
     pub(crate) mask: NodeMask,
     pub(crate) workload: fn(TypeId, Arc<Dependants>, PassDirection) -> WorkloadSystem,
     pub(crate) pass_direction: PassDirection,
+    pub(crate) enter_shadow_dom: bool,
     phantom: PhantomData<V>,
 }
 

packages/native-core/src/real_dom.rs

@@ -10,12 +10,16 @@ use std::collections::VecDeque;
 use std::ops::{Deref, DerefMut};
 use std::sync::{Arc, RwLock};
 
+use crate::custom_element::{
+    CustomElement, CustomElementFactory, CustomElementManager, CustomElementRegistry,
+    CustomElementUpdater,
+};
 use crate::node::{
     ElementNode, FromAnyValue, NodeType, OwnedAttributeDiscription, OwnedAttributeValue, TextNode,
 };
 use crate::node_ref::{NodeMask, NodeMaskBuilder};
-use crate::node_watcher::NodeWatcher;
-use crate::passes::{DirtyNodeStates, PassDirection, TypeErasedState};
+use crate::node_watcher::{AttributeWatcher, NodeWatcher};
+use crate::passes::{Dependant, DirtyNodeStates, PassDirection, TypeErasedState};
 use crate::prelude::AttributeMaskBuilder;
 use crate::tree::{TreeMut, TreeMutView, TreeRef, TreeRefView};
 use crate::NodeId;
@@ -49,6 +53,7 @@ impl Deref for DirtyNodesResult {
 pub(crate) struct NodesDirty<V: FromAnyValue + Send + Sync> {
     passes_updated: FxHashMap<NodeId, FxHashSet<TypeId>>,
     nodes_updated: FxHashMap<NodeId, NodeMask>,
+    nodes_created: FxHashSet<NodeId>,
     pub(crate) passes: Box<[TypeErasedState<V>]>,
 }
 
@@ -92,6 +97,7 @@ impl<V: FromAnyValue + Send + Sync> NodesDirty<V> {
 }
 
 type NodeWatchers<V> = Arc<RwLock<Vec<Box<dyn NodeWatcher<V> + Send + Sync>>>>;
+type AttributeWatchers<V> = Arc<RwLock<Vec<Box<dyn AttributeWatcher<V> + Send + Sync>>>>;
 
 /// A Dom that can sync with the VirtualDom mutations intended for use in lazy renderers.
 /// The render state passes from parent to children and or accumulates state from children to parents.
@@ -108,8 +114,10 @@ pub struct RealDom<V: FromAnyValue + Send + Sync = ()> {
     nodes_listening: FxHashMap<String, FxHashSet<NodeId>>,
     pub(crate) dirty_nodes: NodesDirty<V>,
     node_watchers: NodeWatchers<V>,
+    attribute_watchers: AttributeWatchers<V>,
     workload: ScheduledWorkload,
     root_id: NodeId,
+    custom_elements: Arc<RwLock<CustomElementRegistry<V>>>,
     phantom: std::marker::PhantomData<V>,
 }
 
@@ -123,19 +131,25 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
             let (current, before) = before.split_last_mut().unwrap();
             for state in before.iter_mut().chain(after.iter_mut()) {
                 let dependants = Arc::get_mut(&mut state.dependants).unwrap();
+
+                let current_dependant = Dependant {
+                    type_id: current.this_type_id,
+                    enter_shadow_dom: current.enter_shadow_dom,
+                };
+
                 // If this node depends on the other state as a parent, then the other state should update its children of the current type when it is invalidated
                 if current
                     .parent_dependancies_ids
                     .contains(&state.this_type_id)
-                    && !dependants.child.contains(&current.this_type_id)
+                    && !dependants.child.contains(&current_dependant)
                 {
-                    dependants.child.push(current.this_type_id);
+                    dependants.child.push(current_dependant);
                 }
                 // If this node depends on the other state as a child, then the other state should update its parent of the current type when it is invalidated
                 if current.child_dependancies_ids.contains(&state.this_type_id)
-                    && !dependants.parent.contains(&current.this_type_id)
+                    && !dependants.parent.contains(&current_dependant)
                 {
-                    dependants.parent.push(current.this_type_id);
+                    dependants.parent.push(current_dependant);
                 }
                 // If this node depends on the other state as a sibling, then the other state should update its siblings of the current type when it is invalidated
                 if current.node_dependancies_ids.contains(&state.this_type_id)
@@ -146,15 +160,19 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
             }
             // If the current state depends on itself, then it should update itself when it is invalidated
             let dependants = Arc::get_mut(&mut current.dependants).unwrap();
+            let current_dependant = Dependant {
+                type_id: current.this_type_id,
+                enter_shadow_dom: current.enter_shadow_dom,
+            };
             match current.pass_direction {
                 PassDirection::ChildToParent => {
-                    if !dependants.parent.contains(&current.this_type_id) {
-                        dependants.parent.push(current.this_type_id);
+                    if !dependants.parent.contains(&current_dependant) {
+                        dependants.parent.push(current_dependant);
                     }
                 }
                 PassDirection::ParentToChild => {
-                    if !dependants.child.contains(&current.this_type_id) {
-                        dependants.child.push(current.this_type_id);
+                    if !dependants.child.contains(&current_dependant) {
+                        dependants.child.push(current_dependant);
                     }
                 }
                 _ => {}
@@ -191,10 +209,13 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
                 passes_updated,
                 nodes_updated,
                 passes: tracked_states,
+                nodes_created: [root_id].into_iter().collect(),
             },
             node_watchers: Default::default(),
+            attribute_watchers: Default::default(),
             workload,
             root_id,
+            custom_elements: Default::default(),
             phantom: std::marker::PhantomData,
         }
     }
@@ -211,8 +232,12 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
 
     /// Create a new node of the given type in the dom and return a mutable reference to it.
     pub fn create_node(&mut self, node: impl Into<NodeType<V>>) -> NodeMut<'_, V> {
-        let id = self.world.add_entity(node.into());
+        let node = node.into();
+        let is_element = matches!(node, NodeType::Element(_));
+
+        let id = self.world.add_entity(node);
         self.tree_mut().create_node(id);
+
         self.dirty_nodes
             .passes_updated
             .entry(id)
@@ -220,10 +245,17 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
             .extend(self.dirty_nodes.passes.iter().map(|x| x.this_type_id));
         self.dirty_nodes
             .mark_dirty(id, NodeMaskBuilder::ALL.build());
-        let watchers = self.node_watchers.clone();
-        for watcher in &*watchers.read().unwrap() {
-            watcher.on_node_added(NodeMut::new(id, self));
+        self.dirty_nodes.nodes_created.insert(id);
+
+        // Create a custom element if needed
+        if is_element {
+            let custom_elements = self.custom_elements.clone();
+            custom_elements
+                .read()
+                .unwrap()
+                .add_shadow_dom(NodeMut::new(id, self));
         }
+
         NodeMut::new(id, self)
     }
 
@@ -284,8 +316,48 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
         &mut self,
         ctx: SendAnyMap,
     ) -> (FxDashSet<NodeId>, FxHashMap<NodeId, NodeMask>) {
+        let nodes_created = std::mem::take(&mut self.dirty_nodes.nodes_created);
+
+        // call node watchers
+        {
+            let watchers = self.node_watchers.clone();
+
+            // ignore watchers if they are already being modified
+            if let Ok(mut watchers) = watchers.try_write() {
+                for id in &nodes_created {
+                    for watcher in &mut *watchers {
+                        watcher.on_node_added(NodeMut::new(*id, self));
+                    }
+                }
+            };
+        }
+
         let passes = std::mem::take(&mut self.dirty_nodes.passes_updated);
         let nodes_updated = std::mem::take(&mut self.dirty_nodes.nodes_updated);
+
+        for (node_id, mask) in &nodes_updated {
+            if self.contains(*node_id) {
+                // call attribute watchers but ignore watchers if they are already being modified
+                let watchers = self.attribute_watchers.clone();
+                if let Ok(mut watchers) = watchers.try_write() {
+                    for watcher in &mut *watchers {
+                        watcher.on_attributes_changed(
+                            self.get_mut(*node_id).unwrap(),
+                            mask.attributes(),
+                        );
+                    }
+                };
+
+                // call custom element watchers
+                let node = self.get_mut(*node_id).unwrap();
+                let custom_element_manager =
+                    node.get::<CustomElementManager<V>>().map(|x| x.clone());
+                if let Some(custom_element_manager) = custom_element_manager {
+                    custom_element_manager.on_attributes_changed(node, mask.attributes());
+                }
+            }
+        }
+
         let dirty_nodes =
             DirtyNodeStates::with_passes(self.dirty_nodes.passes.iter().map(|p| p.this_type_id));
         let tree = self.tree_ref();
@@ -312,27 +384,42 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
     }
 
     /// Traverses the dom in a depth first manner, calling the provided function on each node.
-    pub fn traverse_depth_first(&self, mut f: impl FnMut(NodeRef<V>)) {
+    /// If `enter_shadow_dom` is true, then the traversal will enter shadow doms in the tree.
+    pub fn traverse_depth_first_advanced(
+        &self,
+        enter_shadow_dom: bool,
+        mut f: impl FnMut(NodeRef<V>),
+    ) {
         let mut stack = vec![self.root_id()];
         let tree = self.tree_ref();
         while let Some(id) = stack.pop() {
             if let Some(node) = self.get(id) {
                 f(node);
-                let children = tree.children_ids(id);
+                let children = tree.children_ids_advanced(id, enter_shadow_dom);
                 stack.extend(children.iter().copied().rev());
             }
         }
     }
 
+    /// Traverses the dom in a depth first manner, calling the provided function on each node.
+    pub fn traverse_depth_first(&self, f: impl FnMut(NodeRef<V>)) {
+        self.traverse_depth_first_advanced(true, f)
+    }
+
     /// Traverses the dom in a breadth first manner, calling the provided function on each node.
-    pub fn traverse_breadth_first(&self, mut f: impl FnMut(NodeRef<V>)) {
+    /// If `enter_shadow_dom` is true, then the traversal will enter shadow doms in the tree.
+    pub fn traverse_breadth_first_advanced(
+        &self,
+        enter_shadow_doms: bool,
+        mut f: impl FnMut(NodeRef<V>),
+    ) {
         let mut queue = VecDeque::new();
         queue.push_back(self.root_id());
         let tree = self.tree_ref();
         while let Some(id) = queue.pop_front() {
             if let Some(node) = self.get(id) {
                 f(node);
-                let children = tree.children_ids(id);
+                let children = tree.children_ids_advanced(id, enter_shadow_doms);
                 for id in children {
                     queue.push_back(id);
                 }
@@ -340,12 +427,22 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
         }
     }
 
+    /// Traverses the dom in a breadth first manner, calling the provided function on each node.
+    pub fn traverse_breadth_first(&self, f: impl FnMut(NodeRef<V>)) {
+        self.traverse_breadth_first_advanced(true, f);
+    }
+
     /// Traverses the dom in a depth first manner mutably, calling the provided function on each node.
-    pub fn traverse_depth_first_mut(&mut self, mut f: impl FnMut(NodeMut<V>)) {
+    /// If `enter_shadow_dom` is true, then the traversal will enter shadow doms in the tree.
+    pub fn traverse_depth_first_mut_advanced(
+        &mut self,
+        enter_shadow_doms: bool,
+        mut f: impl FnMut(NodeMut<V>),
+    ) {
         let mut stack = vec![self.root_id()];
         while let Some(id) = stack.pop() {
             let tree = self.tree_ref();
-            let mut children = tree.children_ids(id);
+            let mut children = tree.children_ids_advanced(id, enter_shadow_doms);
             drop(tree);
             children.reverse();
             if let Some(node) = self.get_mut(id) {
@@ -356,13 +453,23 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
         }
     }
 
+    /// Traverses the dom in a depth first manner mutably, calling the provided function on each node.
+    pub fn traverse_depth_first_mut(&mut self, f: impl FnMut(NodeMut<V>)) {
+        self.traverse_depth_first_mut_advanced(true, f)
+    }
+
     /// Traverses the dom in a breadth first manner mutably, calling the provided function on each node.
-    pub fn traverse_breadth_first_mut(&mut self, mut f: impl FnMut(NodeMut<V>)) {
+    /// If `enter_shadow_dom` is true, then the traversal will enter shadow doms in the tree.
+    pub fn traverse_breadth_first_mut_advanced(
+        &mut self,
+        enter_shadow_doms: bool,
+        mut f: impl FnMut(NodeMut<V>),
+    ) {
         let mut queue = VecDeque::new();
         queue.push_back(self.root_id());
         while let Some(id) = queue.pop_front() {
             let tree = self.tree_ref();
-            let children = tree.children_ids(id);
+            let children = tree.children_ids_advanced(id, enter_shadow_doms);
             drop(tree);
             if let Some(node) = self.get_mut(id) {
                 f(node);
@@ -373,11 +480,27 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
         }
     }
 
+    /// Traverses the dom in a breadth first manner mutably, calling the provided function on each node.
+    pub fn traverse_breadth_first_mut(&mut self, f: impl FnMut(NodeMut<V>)) {
+        self.traverse_breadth_first_mut_advanced(true, f);
+    }
+
     /// Adds a [`NodeWatcher`] to the dom. Node watchers are called whenever a node is created or removed.
     pub fn add_node_watcher(&mut self, watcher: impl NodeWatcher<V> + 'static + Send + Sync) {
         self.node_watchers.write().unwrap().push(Box::new(watcher));
     }
 
+    /// Adds an [`AttributeWatcher`] to the dom. Attribute watchers are called whenever an attribute is changed.
+    pub fn add_attribute_watcher(
+        &mut self,
+        watcher: impl AttributeWatcher<V> + 'static + Send + Sync,
+    ) {
+        self.attribute_watchers
+            .write()
+            .unwrap()
+            .push(Box::new(watcher));
+    }
+
     /// Returns a reference to the underlying world. Any changes made to the world will not update the reactive system.
     pub fn raw_world(&self) -> &World {
         &self.world
@@ -387,6 +510,20 @@ impl<V: FromAnyValue + Send + Sync> RealDom<V> {
     pub fn raw_world_mut(&mut self) -> &mut World {
         &mut self.world
     }
+
+    /// Registers a new custom element.
+    pub fn register_custom_element<E: CustomElement<V>>(&mut self) {
+        self.register_custom_element_with_factory::<E, E>()
+    }
+
+    /// Registers a new custom element with a custom factory.
+    pub fn register_custom_element_with_factory<F, U>(&mut self)
+    where
+        F: CustomElementFactory<U, V>,
+        U: CustomElementUpdater<V>,
+    {
+        self.custom_elements.write().unwrap().register::<F, U>()
+    }
 }
 
 /// A reference to a tracked component in a node.
@@ -458,6 +595,14 @@ pub trait NodeImmutable<V: FromAnyValue + Send + Sync = ()>: Sized {
             .then(|| ViewEntry::new(view, self.id()))
     }
 
+    /// Get the ids of the children of the current node, if enter_shadow_dom is true and the current node is a shadow slot, the ids of the nodes under the node the shadow slot is attached to will be returned
+    #[inline]
+    fn children_ids_advanced(&self, id: NodeId, enter_shadow_dom: bool) -> Vec<NodeId> {
+        self.real_dom()
+            .tree_ref()
+            .children_ids_advanced(id, enter_shadow_dom)
+    }
+
     /// Get the ids of the children of the current node
     #[inline]
     fn child_ids(&self) -> Vec<NodeId> {
@@ -476,6 +621,14 @@ pub trait NodeImmutable<V: FromAnyValue + Send + Sync = ()>: Sized {
             .collect()
     }
 
+    /// Get the id of the parent of the current node, if enter_shadow_dom is true and the current node is a shadow root, the node the shadow root is attached to will be returned
+    #[inline]
+    fn parent_id_advanced(&self, id: NodeId, enter_shadow_dom: bool) -> Option<NodeId> {
+        self.real_dom()
+            .tree_ref()
+            .parent_id_advanced(id, enter_shadow_dom)
+    }
+
     /// Get the id of the parent of the current node
     #[inline]
     fn parent_id(&self) -> Option<NodeId> {
@@ -585,6 +738,14 @@ impl<'a, V: FromAnyValue + Send + Sync> NodeImmutable<V> for NodeMut<'a, V> {
 }
 
 impl<'a, V: FromAnyValue + Send + Sync> NodeMut<'a, V> {
+    /// Reborrow the node mutably
+    pub fn reborrow(&mut self) -> NodeMut<'_, V> {
+        NodeMut {
+            id: self.id,
+            dom: self.dom,
+        }
+    }
+
     /// Get the real dom this node was created in mutably
     #[inline(always)]
     pub fn real_dom_mut(&mut self) -> &mut RealDom<V> {
@@ -741,6 +902,7 @@ impl<'a, V: FromAnyValue + Send + Sync> NodeMut<'a, V> {
         }
         let id = self.id();
         self.dom.tree_mut().replace(id, new);
+        self.remove();
     }
 
     /// Add an event listener
@@ -798,7 +960,7 @@ impl<'a, V: FromAnyValue + Send + Sync> NodeMut<'a, V> {
     /// mark that this node was removed for the incremental system
     fn mark_removed(&mut self) {
         let watchers = self.dom.node_watchers.clone();
-        for watcher in &*watchers.read().unwrap() {
+        for watcher in &mut *watchers.write().unwrap() {
             watcher.on_node_removed(NodeMut::new(self.id(), self.dom));
         }
     }
@@ -806,9 +968,12 @@ impl<'a, V: FromAnyValue + Send + Sync> NodeMut<'a, V> {
     /// mark that this node was moved for the incremental system
     fn mark_moved(&mut self) {
         let watchers = self.dom.node_watchers.clone();
-        for watcher in &*watchers.read().unwrap() {
-            watcher.on_node_moved(NodeMut::new(self.id(), self.dom));
-        }
+        // ignore watchers if the we are inside of a watcher
+        if let Ok(mut watchers) = watchers.try_write() {
+            for watcher in &mut *watchers {
+                watcher.on_node_moved(NodeMut::new(self.id(), self.dom));
+            }
+        };
     }
 
     /// Get a mutable reference to the type of the current node
@@ -925,6 +1090,15 @@ pub struct ElementNodeMut<'a, V: FromAnyValue + Send + Sync = ()> {
     dirty_nodes: &'a mut NodesDirty<V>,
 }
 
+impl std::fmt::Debug for ElementNodeMut<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ElementNodeMut")
+            .field("id", &self.id)
+            .field("element", &*self.element)
+            .finish()
+    }
+}
+
 impl<V: FromAnyValue + Send + Sync> ElementNodeMut<'_, V> {
     /// Get the current element
     fn element(&self) -> &ElementNode<V> {
@@ -1016,6 +1190,14 @@ impl<V: FromAnyValue + Send + Sync> ElementNodeMut<'_, V> {
         self.element_mut().attributes.get_mut(name)
     }
 
+    /// Get an attribute of the element
+    pub fn get_attribute(
+        &self,
+        name: &OwnedAttributeDiscription,
+    ) -> Option<&OwnedAttributeValue<V>> {
+        self.element().attributes.get(name)
+    }
+
     /// Get the set of all events the element is listening to
     pub fn listeners(&self) -> &FxHashSet<String> {
         &self.element().listeners

packages/native-core/src/tree.rs

@@ -4,11 +4,25 @@ use crate::NodeId;
 use shipyard::{Component, EntitiesViewMut, Get, View, ViewMut};
 use std::fmt::Debug;
 
+/// A shadow tree reference inside of a tree. This tree is isolated from the main tree.
+#[derive(PartialEq, Eq, Clone, Debug, Component)]
+pub struct ShadowTree {
+    /// The root of the shadow tree
+    pub shadow_roots: Vec<NodeId>,
+    /// The node that children of the super tree should be inserted under.
+    pub slot: Option<NodeId>,
+}
+
 /// A node in a tree.
 #[derive(PartialEq, Eq, Clone, Debug, Component)]
 pub struct Node {
     parent: Option<NodeId>,
     children: Vec<NodeId>,
+    child_subtree: Option<ShadowTree>,
+    /// If this node is a slot in a shadow_tree, this is node whose child_subtree is that shadow_tree.
+    slot_for_light_tree: Option<NodeId>,
+    /// If this node is a root of a shadow_tree, this is the node whose child_subtree is that shadow_tree.
+    root_for_light_tree: Option<NodeId>,
     height: u16,
 }
 
@@ -19,10 +33,51 @@ pub type TreeMutView<'a> = (EntitiesViewMut<'a>, ViewMut<'a, Node>);
 
 /// A immutable view of a tree.
 pub trait TreeRef {
+    /// Get the id of the parent of the current node, if enter_shadow_dom is true and the current node is a shadow root, the node the shadow root is attached to will be returned
+    #[inline]
+    fn parent_id_advanced(&self, id: NodeId, enter_shadow_dom: bool) -> Option<NodeId> {
+        // If this node is the root of a shadow_tree, return the node the shadow_tree is attached
+        let root_for_light_tree = self.root_for_light_tree(id);
+        match (root_for_light_tree, enter_shadow_dom) {
+            (Some(id), true) => Some(id),
+            _ => {
+                let parent_id = self.parent_id(id);
+                if enter_shadow_dom {
+                    // If this node is attached via a slot, return the slot as the parent instead of the light tree parent
+                    parent_id.map(|id| {
+                        self.shadow_tree(id)
+                            .and_then(|tree| tree.slot)
+                            .unwrap_or(id)
+                    })
+                } else {
+                    parent_id
+                }
+            }
+        }
+    }
     /// The parent id of the node.
     fn parent_id(&self, id: NodeId) -> Option<NodeId>;
+    /// Get the ids of the children of the current node, if enter_shadow_dom is true and the current node is a shadow slot, the ids of the nodes under the node the shadow slot is attached to will be returned
+    #[inline]
+    fn children_ids_advanced(&self, id: NodeId, enter_shadow_dom: bool) -> Vec<NodeId> {
+        let shadow_tree = self.shadow_tree(id);
+        let slot_of_light_tree = self.slot_for_light_tree(id);
+        match (shadow_tree, slot_of_light_tree, enter_shadow_dom) {
+            // If this node is a shadow root, return the shadow roots
+            (Some(tree), _, true) => tree.shadow_roots.clone(),
+            // If this node is a slot, return the children of the node the slot is attached to
+            (None, Some(id), true) => self.children_ids(id),
+            _ => self.children_ids(id),
+        }
+    }
     /// The children ids of the node.
     fn children_ids(&self, id: NodeId) -> Vec<NodeId>;
+    /// The shadow tree tree under the node.
+    fn shadow_tree(&self, id: NodeId) -> Option<&ShadowTree>;
+    /// The node that contains the shadow tree this node is a slot for
+    fn slot_for_light_tree(&self, id: NodeId) -> Option<NodeId>;
+    /// The node that contains the shadow tree this node is a root of
+    fn root_for_light_tree(&self, id: NodeId) -> Option<NodeId>;
     /// The height of the node.
     fn height(&self, id: NodeId) -> Option<u16>;
     /// Returns true if the node exists.
@@ -31,10 +86,8 @@ pub trait TreeRef {
 
 /// A mutable view of a tree.
 pub trait TreeMut: TreeRef {
-    /// Removes the node and all of its children.
+    /// Removes the node and its children from the tree but do not delete the entities.
     fn remove(&mut self, id: NodeId);
-    /// Removes the node and all of its children.
-    fn remove_single(&mut self, id: NodeId);
     /// Adds a new node to the tree.
     fn create_node(&mut self, id: NodeId);
     /// Adds a child to the node.
@@ -45,6 +98,10 @@ pub trait TreeMut: TreeRef {
     fn insert_before(&mut self, old_id: NodeId, new_id: NodeId);
     /// Inserts a node after another node.
     fn insert_after(&mut self, old_id: NodeId, new_id: NodeId);
+    /// Creates a new shadow tree.
+    fn create_subtree(&mut self, id: NodeId, shadow_roots: Vec<NodeId>, slot: Option<NodeId>);
+    /// Remove any shadow tree.
+    fn remove_subtree(&mut self, id: NodeId);
 }
 
 impl<'a> TreeRef for TreeRefView<'a> {
@@ -65,16 +122,47 @@ impl<'a> TreeRef for TreeRefView<'a> {
     fn contains(&self, id: NodeId) -> bool {
         self.get(id).is_ok()
     }
+
+    fn shadow_tree(&self, id: NodeId) -> Option<&ShadowTree> {
+        self.get(id).ok()?.child_subtree.as_ref()
+    }
+
+    fn slot_for_light_tree(&self, id: NodeId) -> Option<NodeId> {
+        self.get(id).ok()?.slot_for_light_tree
+    }
+
+    fn root_for_light_tree(&self, id: NodeId) -> Option<NodeId> {
+        self.get(id).ok()?.root_for_light_tree
+    }
 }
 
 impl<'a> TreeMut for TreeMutView<'a> {
     fn remove(&mut self, id: NodeId) {
         fn recurse(tree: &mut TreeMutView<'_>, id: NodeId) {
-            let children = tree.children_ids(id);
+            let (light_tree, children) = {
+                let node = (&mut tree.1).get(id).unwrap();
+                (node.slot_for_light_tree, std::mem::take(&mut node.children))
+            };
+
             for child in children {
                 recurse(tree, child);
             }
+
+            // If this node is a slot in a shadow_tree, remove it from the shadow_tree.
+            if let Some(light_tree) = light_tree {
+                let root_for_light_tree = (&mut tree.1).get(light_tree).unwrap();
+
+                if let Some(shadow_tree) = &mut root_for_light_tree.child_subtree {
+                    shadow_tree.slot = None;
+                }
+
+                debug_assert!(
+                    root_for_light_tree.children.is_empty(),
+                    "ShadowTree root should have no children when slot is removed."
+                );
+            }
         }
+
         {
             let mut node_data_mut = &mut self.1;
             if let Some(parent) = node_data_mut.get(id).unwrap().parent {
@@ -86,16 +174,6 @@ impl<'a> TreeMut for TreeMutView<'a> {
         recurse(self, id);
     }
 
-    fn remove_single(&mut self, id: NodeId) {
-        {
-            let mut node_data_mut = &mut self.1;
-            if let Some(parent) = node_data_mut.get(id).unwrap().parent {
-                let parent = (&mut node_data_mut).get(parent).unwrap();
-                parent.children.retain(|&child| child != id);
-            }
-        }
-    }
-
     fn create_node(&mut self, id: NodeId) {
         let (entities, node_data_mut) = self;
         entities.add_component(
@@ -105,19 +183,21 @@ impl<'a> TreeMut for TreeMutView<'a> {
                 parent: None,
                 children: Vec::new(),
                 height: 0,
+                child_subtree: None,
+                slot_for_light_tree: None,
+                root_for_light_tree: None,
             },
         );
     }
 
     fn add_child(&mut self, parent: NodeId, new: NodeId) {
-        let height;
         {
             let mut node_state = &mut self.1;
             (&mut node_state).get(new).unwrap().parent = Some(parent);
             let parent = (&mut node_state).get(parent).unwrap();
             parent.children.push(new);
-            height = parent.height + 1;
         }
+        let height = child_height((&self.1).get(parent).unwrap(), self);
         set_height(self, new, height);
     }
 
@@ -133,27 +213,29 @@ impl<'a> TreeMut for TreeMutView<'a> {
                         break;
                     }
                 }
-                let height = parent.height + 1;
+                let height = child_height((&self.1).get(parent_id).unwrap(), self);
                 set_height(self, new_id, height);
             }
         }
-        // remove the old node
         self.remove(old_id);
     }
 
     fn insert_before(&mut self, old_id: NodeId, new_id: NodeId) {
-        let mut node_state = &mut self.1;
-        let old_node = node_state.get(old_id).unwrap();
-        let parent_id = old_node.parent.expect("tried to insert before root");
-        (&mut node_state).get(new_id).unwrap().parent = Some(parent_id);
-        let parent = (&mut node_state).get(parent_id).unwrap();
+        let parent_id = {
+            let old_node = self.1.get(old_id).unwrap();
+            old_node.parent.expect("tried to insert before root")
+        };
+        {
+            (&mut self.1).get(new_id).unwrap().parent = Some(parent_id);
+        }
+        let parent = (&mut self.1).get(parent_id).unwrap();
         let index = parent
             .children
             .iter()
             .position(|child| *child == old_id)
             .unwrap();
         parent.children.insert(index, new_id);
-        let height = parent.height + 1;
+        let height = child_height((&self.1).get(parent_id).unwrap(), self);
         set_height(self, new_id, height);
     }
 
@@ -169,21 +251,121 @@ impl<'a> TreeMut for TreeMutView<'a> {
             .position(|child| *child == old_id)
             .unwrap();
         parent.children.insert(index + 1, new_id);
-        let height = parent.height + 1;
+        let height = child_height((&self.1).get(parent_id).unwrap(), self);
         set_height(self, new_id, height);
     }
+
+    fn create_subtree(&mut self, id: NodeId, shadow_roots: Vec<NodeId>, slot: Option<NodeId>) {
+        let (_, node_data_mut) = self;
+
+        let light_root_height;
+        {
+            let shadow_tree = ShadowTree {
+                shadow_roots: shadow_roots.clone(),
+                slot,
+            };
+
+            let light_root = node_data_mut
+                .get(id)
+                .expect("tried to create shadow_tree with non-existent id");
+
+            light_root.child_subtree = Some(shadow_tree);
+            light_root_height = light_root.height;
+
+            if let Some(slot) = slot {
+                let slot = node_data_mut
+                    .get(slot)
+                    .expect("tried to create shadow_tree with non-existent slot");
+                slot.slot_for_light_tree = Some(id);
+            }
+        }
+
+        // Now that we have created the shadow_tree, we need to update the height of the shadow_tree roots
+        for root in shadow_roots {
+            (&mut self.1).get(root).unwrap().root_for_light_tree = Some(id);
+            set_height(self, root, light_root_height + 1);
+        }
+    }
+
+    fn remove_subtree(&mut self, id: NodeId) {
+        let (_, node_data_mut) = self;
+
+        if let Ok(node) = node_data_mut.get(id) {
+            if let Some(shadow_tree) = node.child_subtree.take() {
+                // Remove the slot's link to the shadow_tree
+                if let Some(slot) = shadow_tree.slot {
+                    let slot = node_data_mut
+                        .get(slot)
+                        .expect("tried to remove shadow_tree with non-existent slot");
+                    slot.slot_for_light_tree = None;
+                }
+
+                let node = node_data_mut.get(id).unwrap();
+
+                // Reset the height of the light root's children
+                let height = node.height;
+                for child in node.children.clone() {
+                    println!("child: {:?}", child);
+                    set_height(self, child, height + 1);
+                }
+
+                // Reset the height of the shadow roots
+                for root in &shadow_tree.shadow_roots {
+                    set_height(self, *root, 0);
+                }
+            }
+        }
+    }
+}
+
+fn child_height(parent: &Node, tree: &impl TreeRef) -> u16 {
+    match &parent.child_subtree {
+        Some(shadow_tree) => {
+            if let Some(slot) = shadow_tree.slot {
+                tree.height(slot)
+                    .expect("Attempted to read a slot that does not exist")
+                    + 1
+            } else {
+                panic!("Attempted to read the height of a child of a node with a shadow tree, but the shadow tree does not have a slot. Every shadow tree attached to a node with children must have a slot.")
+            }
+        }
+        None => parent.height + 1,
+    }
 }
 
 /// Sets the height of a node and updates the height of all its children
 fn set_height(tree: &mut TreeMutView<'_>, node: NodeId, height: u16) {
-    let children = {
+    let (shadow_tree, light_tree, children) = {
         let mut node_data_mut = &mut tree.1;
         let node = (&mut node_data_mut).get(node).unwrap();
         node.height = height;
-        node.children.clone()
+
+        (
+            node.child_subtree.clone(),
+            node.slot_for_light_tree,
+            node.children.clone(),
+        )
     };
-    for child in children {
-        set_height(tree, child, height + 1);
+
+    // If the children are actually part of a shadow_tree, there height is determined by the height of the shadow_tree
+    if let Some(shadow_tree) = shadow_tree {
+        // Set the height of the shadow_tree roots
+        for &shadow_root in &shadow_tree.shadow_roots {
+            set_height(tree, shadow_root, height + 1);
+        }
+    } else {
+        // Otherwise, we just set the height of the children to be one more than the height of the parent
+        for child in children {
+            set_height(tree, child, height + 1);
+        }
+    }
+
+    // If this nodes is a slot for a shadow_tree, we need to go to the super tree and update the height of its children
+    if let Some(light_tree) = light_tree {
+        let children = (&tree.1).get(light_tree).unwrap().children.clone();
+        for child in children {
+            set_height(tree, child, height + 1);
+        }
     }
 }
 
@@ -209,6 +391,21 @@ impl<'a> TreeRef for TreeMutView<'a> {
     fn contains(&self, id: NodeId) -> bool {
         self.1.get(id).is_ok()
     }
+
+    fn shadow_tree(&self, id: NodeId) -> Option<&ShadowTree> {
+        let node_data = &self.1;
+        node_data.get(id).ok()?.child_subtree.as_ref()
+    }
+
+    fn slot_for_light_tree(&self, id: NodeId) -> Option<NodeId> {
+        let node_data = &self.1;
+        node_data.get(id).ok()?.slot_for_light_tree
+    }
+
+    fn root_for_light_tree(&self, id: NodeId) -> Option<NodeId> {
+        let node_data = &self.1;
+        node_data.get(id).ok()?.root_for_light_tree
+    }
 }
 
 #[test]
@@ -235,6 +432,85 @@ fn creation() {
     assert_eq!(tree.children_ids(parent_id), &[child_id]);
 }
 
+#[test]
+fn shadow_tree() {
+    use shipyard::World;
+    #[derive(Component)]
+    struct Num(i32);
+
+    let mut world = World::new();
+    // Create main tree
+    let parent_id = world.add_entity(Num(1i32));
+    let child_id = world.add_entity(Num(0i32));
+
+    // Create shadow tree
+    let shadow_parent_id = world.add_entity(Num(2i32));
+    let shadow_child_id = world.add_entity(Num(3i32));
+
+    let mut tree = world.borrow::<TreeMutView>().unwrap();
+
+    tree.create_node(parent_id);
+    tree.create_node(child_id);
+
+    tree.add_child(parent_id, child_id);
+
+    tree.create_node(shadow_parent_id);
+    tree.create_node(shadow_child_id);
+
+    tree.add_child(shadow_parent_id, shadow_child_id);
+
+    // Check that both trees are correct individually
+    assert_eq!(tree.height(parent_id), Some(0));
+    assert_eq!(tree.height(child_id), Some(1));
+    assert_eq!(tree.parent_id(parent_id), None);
+    assert_eq!(tree.parent_id(child_id).unwrap(), parent_id);
+    assert_eq!(tree.children_ids(parent_id), &[child_id]);
+
+    assert_eq!(tree.height(shadow_parent_id), Some(0));
+    assert_eq!(tree.height(shadow_child_id), Some(1));
+    assert_eq!(tree.parent_id(shadow_parent_id), None);
+    assert_eq!(tree.parent_id(shadow_child_id).unwrap(), shadow_parent_id);
+    assert_eq!(tree.children_ids(shadow_parent_id), &[shadow_child_id]);
+
+    // Add shadow tree to main tree
+    tree.create_subtree(parent_id, vec![shadow_parent_id], Some(shadow_child_id));
+
+    assert_eq!(tree.height(parent_id), Some(0));
+    assert_eq!(tree.height(shadow_parent_id), Some(1));
+    assert_eq!(tree.height(shadow_child_id), Some(2));
+    assert_eq!(tree.height(child_id), Some(3));
+    assert_eq!(
+        tree.1
+            .get(parent_id)
+            .unwrap()
+            .child_subtree
+            .as_ref()
+            .unwrap()
+            .shadow_roots,
+        &[shadow_parent_id]
+    );
+    assert_eq!(
+        tree.1.get(shadow_child_id).unwrap().slot_for_light_tree,
+        Some(parent_id)
+    );
+
+    // Remove shadow tree from main tree
+    tree.remove_subtree(parent_id);
+
+    // Check that both trees are correct individually
+    assert_eq!(tree.height(parent_id), Some(0));
+    assert_eq!(tree.height(child_id), Some(1));
+    assert_eq!(tree.parent_id(parent_id), None);
+    assert_eq!(tree.parent_id(child_id).unwrap(), parent_id);
+    assert_eq!(tree.children_ids(parent_id), &[child_id]);
+
+    assert_eq!(tree.height(shadow_parent_id), Some(0));
+    assert_eq!(tree.height(shadow_child_id), Some(1));
+    assert_eq!(tree.parent_id(shadow_parent_id), None);
+    assert_eq!(tree.parent_id(shadow_child_id).unwrap(), shadow_parent_id);
+    assert_eq!(tree.children_ids(shadow_parent_id), &[shadow_child_id]);
+}
+
 #[test]
 fn insertion() {
     use shipyard::World;

packages/native-core/src/utils/persistant_iterator.rs

@@ -60,7 +60,7 @@ struct PersistantElementIterUpdater<V> {
 }
 
 impl<V: FromAnyValue + Sync + Send> NodeWatcher<V> for PersistantElementIterUpdater<V> {
-    fn on_node_moved(&self, node: NodeMut<V>) {
+    fn on_node_moved(&mut self, node: NodeMut<V>) {
         // if any element is moved, update its parents in the stack
         let mut stack = self.stack.lock().unwrap();
         let moved = node.id();
@@ -78,7 +78,7 @@ impl<V: FromAnyValue + Sync + Send> NodeWatcher<V> for PersistantElementIterUpda
         }
     }
 
-    fn on_node_removed(&self, node: NodeMut<V>) {
+    fn on_node_removed(&mut self, node: NodeMut<V>) {
         // if any element is removed in the chain, remove it and its children from the stack
         let mut stack = self.stack.lock().unwrap();
         let removed = node.id();