handle children recursively, including keyed children

Author: tigerabrodi

src/lib/__tests__/diff.children.test.ts

@@ -180,7 +180,7 @@ describe('diffNodes - Children cases', () => {
     const ComponentFn = (props: Props): VNode => ({
       kind: 'static',
       type: 'div',
-      props: {},
+      props,
       children: [],
     })
 
@@ -219,24 +219,204 @@ describe('diffNodes - Children cases', () => {
     expect(patches).toEqual([])
   })
 
-  // Text and number nodes
-  test('handles text node changes')
-  test('handles number node changes')
-  test('handles mixed node types (elements, text, numbers)')
-
-  // Nested structures
-  test('recursively diffs nested children')
-  test('handles deeply nested structure changes')
-  test('correctly patches nested children props changes')
-
-  // Edge cases
-  test('handles empty children array to non-empty')
-  test('handles non-empty children array to empty')
-  test('handles null/undefined children')
-  test('handles children type changes (text to element)')
-
-  // Mixed node kinds
-  test('handles mix of static and component children')
-  test('handles mix of regular and memo component children')
-  test('preserves memoization in nested children')
+  test('handles multiple children additions and removals', () => {
+    const oldNode: StaticVNode = {
+      kind: 'static',
+      type: 'div',
+      props: {},
+      children: [
+        {
+          kind: 'static',
+          type: 'span',
+          props: { id: '1' },
+          children: [],
+          key: '1',
+        },
+        {
+          kind: 'static',
+          type: 'span',
+          props: { id: '2' },
+          children: [],
+          key: '2',
+        },
+      ],
+    }
+
+    const newNode: StaticVNode = {
+      kind: 'static',
+      type: 'div',
+      props: {},
+      children: [
+        {
+          kind: 'static',
+          type: 'span',
+          props: { id: '3' },
+          children: [],
+          key: '3', // New key
+        },
+        {
+          kind: 'static',
+          type: 'span',
+          props: { id: '1' },
+          children: [],
+          key: '1', // Same key as before
+        },
+      ],
+    }
+
+    const patches = diffNodes(oldNode, newNode)
+
+    // Should generate:
+    // - REORDER for id:1 span, fromIndex: 0 toIndex: 1
+    // - REMOVE for id:2 span
+    // - INSERT for id:3 span
+    expect(patches).toEqual([
+      {
+        type: 'REORDER',
+        node: oldNode.children[0], // id:1 span
+        fromIndex: 0,
+        toIndex: 1,
+      },
+      {
+        type: 'REMOVE',
+        node: oldNode.children[1], // id:2 span
+      },
+      {
+        type: 'INSERT',
+        node: newNode.children[0], // id:3 span
+        index: 0,
+      },
+    ])
+  })
+
+  // Nested structures - Important for recursive diffing
+  test('recursively diffs nested children', () => {
+    const oldNode: StaticVNode = {
+      kind: 'static',
+      type: 'div',
+      props: {},
+      children: [
+        {
+          kind: 'static',
+          type: 'div',
+          props: {},
+          children: [
+            {
+              kind: 'static',
+              type: 'span',
+              props: { value: 'old' },
+              children: [],
+              key: 'nested',
+            },
+          ],
+          key: 'parent',
+        },
+      ],
+    }
+
+    const newNode: StaticVNode = {
+      kind: 'static',
+      type: 'div',
+      props: {},
+      children: [
+        {
+          kind: 'static',
+          type: 'div',
+          props: {},
+          children: [
+            {
+              kind: 'static',
+              type: 'span',
+              props: { value: 'new' },
+              children: [],
+              key: 'nested',
+            },
+          ],
+          key: 'parent',
+        },
+      ],
+    }
+
+    const patches = diffNodes(oldNode, newNode)
+
+    expect(patches).toEqual([
+      {
+        type: 'PROPS',
+        node: oldNode.children[0].children[0],
+        newProps: { value: 'new' },
+      },
+    ])
+  })
+
+  // Mixed node kinds - Tests how different node types interact
+  test('handles mix of regular and memo component children', () => {
+    const ComponentFn = (props: Props): VNode => ({
+      kind: 'static',
+      type: 'div',
+      props,
+      children: [],
+    })
+
+    const MemoComponentFn = (props: Props): VNode => ({
+      kind: 'static',
+      type: 'span',
+      props,
+      children: [],
+    })
+
+    const oldNode: StaticVNode = {
+      kind: 'static',
+      type: 'div',
+      props: {},
+      children: [
+        {
+          kind: 'regular',
+          type: ComponentFn,
+          props: { value: 'old' },
+          children: [],
+          key: 'regular',
+        },
+        {
+          kind: 'memo',
+          type: MemoComponentFn,
+          props: { value: 'old' },
+          children: [],
+          compare: (prev, next) => prev.value === next.value,
+          key: 'memo',
+        },
+      ],
+    }
+
+    const newNode: StaticVNode = {
+      kind: 'static',
+      type: 'div',
+      props: {},
+      children: [
+        {
+          kind: 'regular',
+          type: ComponentFn,
+          props: { value: 'new' },
+          children: [],
+          key: 'regular',
+        },
+        {
+          kind: 'memo',
+          type: MemoComponentFn,
+          props: { value: 'old' }, // Same value, shouldn't update
+          children: [],
+          compare: (prev, next) => prev.value === next.value,
+          key: 'memo',
+        },
+      ],
+    }
+
+    const patches = diffNodes(oldNode, newNode)
+    expect(patches).toEqual([
+      {
+        type: 'PROPS',
+        node: oldNode.children[0],
+        newProps: { value: 'new' },
+      },
+    ])
+  })
 })

src/lib/diff.ts

@@ -82,43 +82,100 @@ export function diffNodes(oldNode: VNode, newNode: VNode): Array<Patch> {
     const oldChildren = oldNode.children
     const newChildren = newNode.children
 
-    // maxLength so that we go over all children
-    const maxLength = Math.max(oldChildren.length, newChildren.length)
+    // Get all keys and map to index for both old and new children
+    const oldKeys = new Map<string | number, number>()
+    const newKeys = new Map<string | number, number>()
+
+    // Track which indices we've handled
+    // When handling children without keys, we need to know which ones we've already handled
+    const handledIndices = new Set<number>()
+
+    // Gather all keys for both old and new children
+    oldChildren.forEach((child, index) => {
+      if (child.key) oldKeys.set(child.key, index)
+    })
 
-    for (let i = 0; i < maxLength; i++) {
-      // Either both exists or one of them doesn't
-      // We know because of maxLength, either of them will always exist till the end
-      const oldChild = oldChildren[i]
-      const newChild = newChildren[i]
+    newChildren.forEach((child, index) => {
+      if (child.key) newKeys.set(child.key, index)
+    })
+
+    // Handle keyed nodes first
+    oldChildren.forEach((oldChild, indexInOldVDom) => {
+      if (!oldChild.key) return // Skip keyless for now
 
-      // If old child doesn't exist, we know that newChild exists
-      if (!oldChild) {
+      if (!newKeys.has(oldChild.key)) {
+        // Key no longer exists -> REMOVE
         patches.push({
-          type: 'INSERT',
-          node: newChild, // Already VNode which is either static or component node
-          index: i,
+          type: 'REMOVE',
+          node: oldChild,
         })
+      } else {
+        const indexInNewVDom = newKeys.get(oldChild.key)!
+        const hasMoved = indexInNewVDom !== indexInOldVDom
+        if (hasMoved) {
+          // Position changed -> REORDER
+          patches.push({
+            type: 'REORDER',
+            node: oldChild,
+            fromIndex: indexInOldVDom,
+            toIndex: indexInNewVDom,
+          })
+        }
+
+        const newChild = newChildren[indexInNewVDom]
 
-        // No need to do any more work here
-        // Continue to the next child
-        continue
+        // Handle the keyed child recursively using both its old and new child
+        patches.push(...diffNodes(oldChild, newChild))
+
+        handledIndices.add(indexInNewVDom)
       }
+    })
+
+    newChildren.forEach((newChild, indexInNewVDom) => {
+      if (!newChild.key) return // Skip keyless for now
 
-      if (!newChild) {
+      if (!oldKeys.has(newChild.key)) {
+        // New key -> INSERT
         patches.push({
-          type: 'REMOVE',
-          node: oldChild, // Already VNode which is either static or component node
+          type: 'INSERT',
+          node: newChild,
+          index: indexInNewVDom,
         })
 
-        // No need to do any more work here
-        continue
+        handledIndices.add(indexInNewVDom)
       }
+    })
+
+    // Now handle keyless nodes by position
+    // But only for indices we haven't handled yet!
+    const maxLength = Math.max(oldChildren.length, newChildren.length)
+
+    for (let index = 0; index < maxLength; index++) {
+      const hasAlreadyBeenHandled = handledIndices.has(index)
+      if (hasAlreadyBeenHandled) continue
+
+      const oldChild = oldChildren[index]
+      const newChild = newChildren[index]
 
-      // Both exist! Need to diff them
-      // We know that both are VNode
-      // diffNodes will return patches for the children
-      const patch = diffNodes(oldChild, newChild)
-      patches.push(...patch)
+      if (!oldChild && newChild) {
+        // New child for this position (index)
+        // We know it won't collide with the keyed children because we're always skipping them
+        patches.push({
+          type: 'INSERT',
+          node: newChild,
+          index: index,
+        })
+      } else if (oldChild && !newChild) {
+        // Old child exists, but new child doesn't
+        // Remove the old child
+        patches.push({
+          type: 'REMOVE',
+          node: oldChild,
+        })
+      } else if (oldChild && newChild) {
+        // Both exist, diff them recursively
+        patches.push(...diffNodes(oldChild, newChild))
+      }
     }
   }