Fix mount use after removal/Simplify mounts a bit (#2834)

* remove clone mounted; mounted nodes should never leak into components?

* Remove the mount from vnodes after unmounting them

* Fix use after removal in iterator diffing

* remove logs

* Re-fix keyed list diffing

* Use new.len for diffing keyed lens and remove redundant check

* simplify core list diffing logic a bit

* remove println

* add proper asserts for new tests

---------

Co-authored-by: Jonathan Kelley <jkelleyrtp@gmail.com>

packages/core/src/diff/component.rs

@@ -9,11 +9,11 @@ use crate::{
         ElementRef, MountId, ScopeOrder, SuspenseBoundaryProps, SuspenseBoundaryPropsWithOwner,
         VComponent, WriteMutations,
     },
-    nodes::VNode,
+    nodes::{AsVNode, VNode},
     prelude::SuspenseContext,
     scopes::ScopeId,
     virtual_dom::VirtualDom,
-    RenderReturn,
+    Element,
 };
 
 impl VirtualDom {
@@ -36,11 +36,11 @@ impl VirtualDom {
         &mut self,
         to: Option<&mut M>,
         scope: ScopeId,
-        new_nodes: RenderReturn,
+        new_nodes: Element,
     ) {
         self.runtime.clone().with_scope_on_stack(scope, || {
             // We don't diff the nodes if the scope is suspended or has an error
-            let Ok(new_real_nodes) = &new_nodes.node else {
+            let Ok(new_real_nodes) = &new_nodes else {
                 return;
             };
             let scope_state = &mut self.scopes[scope.0];
@@ -51,7 +51,8 @@ impl VirtualDom {
             // If it is suspended, we need to diff it but write the mutations nothing
             // Note: It is important that we still diff the scope even if it is suspended, because the scope may render other child components which may change between renders
             let mut render_to = to.filter(|_| self.runtime.scope_should_render(scope));
-            old.diff_node(new_real_nodes, self, render_to.as_deref_mut());
+            old.as_vnode()
+                .diff_node(new_real_nodes, self, render_to.as_deref_mut());
 
             self.scopes[scope.0].last_rendered_node = Some(new_nodes);
 
@@ -69,7 +70,7 @@ impl VirtualDom {
         &mut self,
         to: Option<&mut M>,
         scope: ScopeId,
-        new_nodes: RenderReturn,
+        new_nodes: Element,
         parent: Option<ElementRef>,
     ) -> usize {
         self.runtime.clone().with_scope_on_stack(scope, || {
@@ -79,7 +80,9 @@ impl VirtualDom {
             let mut render_to = to.filter(|_| self.runtime.scope_should_render(scope));
 
             // Create the node
-            let nodes = new_nodes.create(self, parent, render_to.as_deref_mut());
+            let nodes = new_nodes
+                .as_vnode()
+                .create(self, parent, render_to.as_deref_mut());
 
             // Then set the new node as the last rendered node
             self.scopes[scope.0].last_rendered_node = Some(new_nodes);
@@ -104,8 +107,12 @@ impl VirtualDom {
 
         // Remove the component from the dom
         if let Some(node) = self.scopes[scope_id.0].last_rendered_node.as_ref() {
-            node.clone_mounted()
-                .remove_node_inner(self, to, destroy_component_state, replace_with)
+            node.clone().as_vnode().remove_node_inner(
+                self,
+                to,
+                destroy_component_state,
+                replace_with,
+            )
         };
 
         if destroy_component_state {

packages/core/src/diff/iterator.rs

@@ -133,29 +133,17 @@ impl VirtualDom {
         let new_middle = &new[left_offset..(new.len() - right_offset)];
 
         debug_assert!(
-            !((old_middle.len() == new_middle.len()) && old_middle.is_empty()),
-            "keyed children must have the same number of children"
+            !old_middle.is_empty(),
+            "Old middle returned from `diff_keyed_ends` should not be empty"
+        );
+        debug_assert!(
+            !new_middle.is_empty(),
+            "New middle returned from `diff_keyed_ends` should not be empty"
         );
 
+        // A few nodes in the middle were removed, just remove the old nodes
         if new_middle.is_empty() {
-            // remove the old elements
             self.remove_nodes(to, old_middle, None);
-        } else if old_middle.is_empty() {
-            // there were no old elements, so just create the new elements
-            // we need to find the right "foothold" though - we shouldn't use the "append" at all
-            if left_offset == 0 {
-                // insert at the beginning of the old list
-                let foothold = &old[old.len() - right_offset];
-                self.create_and_insert_before(to, new_middle, foothold, parent);
-            } else if right_offset == 0 {
-                // insert at the end  the old list
-                let foothold = old.last().unwrap();
-                self.create_and_insert_after(to, new_middle, foothold, parent);
-            } else {
-                // inserting in the middle
-                let foothold = &old[left_offset - 1];
-                self.create_and_insert_after(to, new_middle, foothold, parent);
-            }
         } else {
             self.diff_keyed_middle(to, old_middle, new_middle, parent);
         }
@@ -187,14 +175,7 @@ impl VirtualDom {
         // If that was all of the old children, then create and append the remaining
         // new children and we're finished.
         if left_offset == old.len() {
-            self.create_and_insert_after(to, &new[left_offset..], old.last().unwrap(), parent);
-            return None;
-        }
-
-        // And if that was all of the new children, then remove all of the remaining
-        // old children and we're finished.
-        if left_offset == new.len() {
-            self.remove_nodes(to, &old[left_offset..], None);
+            self.create_and_insert_after(to, &new[left_offset..], &new[left_offset - 1], parent);
             return None;
         }
 
@@ -209,6 +190,35 @@ impl VirtualDom {
             right_offset += 1;
         }
 
+        // If that was all of the old children, then create and prepend the remaining
+        // new children and we're finished.
+        if right_offset == old.len() {
+            self.create_and_insert_before(
+                to,
+                &new[..new.len() - right_offset],
+                &new[new.len() - right_offset],
+                parent,
+            );
+            return None;
+        }
+
+        // If the right offset + the left offset is the same as the new length, then we just need to remove the old nodes
+        if right_offset + left_offset == new.len() {
+            self.remove_nodes(to, &old[left_offset..old.len() - right_offset], None);
+            return None;
+        }
+
+        // If the right offset + the left offset is the same as the old length, then we just need to add the new nodes
+        if right_offset + left_offset == old.len() {
+            self.create_and_insert_before(
+                to,
+                &new[left_offset..new.len() - right_offset],
+                &new[new.len() - right_offset],
+                parent,
+            );
+            return None;
+        }
+
         Some((left_offset, right_offset))
     }
 

packages/core/src/diff/node.rs

@@ -35,8 +35,6 @@ impl VNode {
             return self.replace(std::slice::from_ref(new), parent, dom, to);
         }
 
-        let mount_id = self.mount.get();
-
         self.move_mount_to(new, dom);
 
         // If the templates are the same, we don't need to do anything, except copy over the mount information
@@ -52,6 +50,7 @@ impl VNode {
         }
 
         // Now diff the dynamic nodes
+        let mount_id = new.mount.get();
         for (dyn_node_idx, (old, new)) in self
             .dynamic_nodes
             .iter()
@@ -64,7 +63,7 @@ impl VNode {
 
     fn move_mount_to(&self, new: &VNode, dom: &mut VirtualDom) {
         // Copy over the mount information
-        let mount_id = self.mount.get();
+        let mount_id = self.mount.take();
         new.mount.set(mount_id);
 
         if mount_id.mounted() {
@@ -72,7 +71,7 @@ impl VNode {
             let mount = &mut mounts[mount_id.0];
 
             // Update the reference to the node for bubbling events
-            mount.node = new.clone_mounted();
+            mount.node = new.clone();
         }
     }
 
@@ -293,6 +292,7 @@ impl VNode {
         self.reclaim_roots(mount, dom, to, destroy_component_state, replace_with);
 
         if destroy_component_state {
+            let mount = self.mount.take();
             // Remove the mount information
             dom.runtime.mounts.borrow_mut().remove(mount.0);
         }
@@ -416,7 +416,7 @@ impl VNode {
         dom: &mut VirtualDom,
         to: &mut impl WriteMutations,
     ) {
-        let mount_id = self.mount.get();
+        let mount_id = new.mount.get();
         for (idx, (old_attrs, new_attrs)) in self
             .dynamic_attrs
             .iter()
@@ -535,7 +535,7 @@ impl VNode {
             self.mount.set(mount);
             tracing::trace!(?self, ?mount, "creating template");
             entry.insert(VNodeMount {
-                node: self.clone_mounted(),
+                node: self.clone(),
                 parent,
                 root_ids: vec![ElementId(0); template.roots.len()].into_boxed_slice(),
                 mounted_attributes: vec![ElementId(0); template.attr_paths.len()]
@@ -559,7 +559,8 @@ impl VNode {
                     .as_usize()
                     .expect("node should already be mounted"),
             ),
-            "Node mount should be valid"
+            "Tried to find mount {:?} in dom.mounts, but it wasn't there",
+            self.mount.get()
         );
         let mount = self.mount.get();