Allow returning a node from decorate modifier to replace a node

src/core/util.ts

@@ -1,12 +1,12 @@
 import { Handle } from './Destroyable';
-import { DNode, RenderResult } from './interfaces';
+import { DNode, RenderResult, VNode, WNode } from './interfaces';
 import { isWNode, isVNode } from './vdom';
 
 const slice = Array.prototype.slice;
 const hasOwnProperty = Object.prototype.hasOwnProperty;
 
 export interface Modifier<T extends DNode> {
-	(dNode: T, breaker: () => void): void;
+	(dNode: T, breaker: () => void): void | VNode | WNode | null | string | boolean;
 }
 
 export interface Predicate<T extends DNode> {
@@ -368,20 +368,50 @@ export function decorate(
 		shallow = optionsOrModifier.shallow || false;
 	}
 
-	let nodes = Array.isArray(dNodes) ? [...dNodes] : [dNodes];
+	let shouldDrain = false;
 	function breaker() {
-		nodes = [];
+		shouldDrain = true;
+	}
+
+	function modifyNode(modifier: Modifier<DNode>, node: DNode) {
+		const modifiedNode = modifier(node, breaker);
+		return modifiedNode === undefined ? node : modifiedNode;
 	}
-	while (nodes.length) {
-		const node = nodes.shift();
+
+	let nodes = Array.isArray(dNodes) ? [...dNodes] : [dNodes];
+	for (let i = 0; i < nodes.length; i++) {
+		const node = nodes[i];
 		if (node && node !== true) {
-			if (!shallow && (isWNode(node) || isVNode(node)) && node.children) {
-				nodes = [...nodes, ...node.children];
-			}
 			if (!predicate || predicate(node)) {
-				modifier(node, breaker);
+				nodes[i] = modifyNode(modifier, node);
 			}
 		}
+		if (shouldDrain) {
+			break;
+		}
 	}
-	return dNodes;
+
+	const rootNodes = nodes.slice();
+
+	if (!shallow) {
+		while (nodes.length) {
+			const node = nodes.shift();
+			if ((isWNode(node) || isVNode(node)) && node.children) {
+				for (let i = 0; i < node.children.length; i++) {
+					const child = node.children[i];
+					if (child && child !== true) {
+						if (!predicate || predicate(child)) {
+							node.children[i] = modifyNode(modifier, child);
+						}
+					}
+					nodes.push(node.children[i]);
+					if (shouldDrain) {
+						nodes = [];
+					}
+				}
+			}
+		}
+	}
+
+	return Array.isArray(dNodes) ? rootNodes : rootNodes[0];
 }

tests/core/unit/d.ts

@@ -407,6 +407,51 @@ registerSuite('d', {
 				assert.strictEqual(children[4], 'my text');
 				assert.isNull(children[5]);
 			}
+		},
+		'returning a new node replaces the node, and new children will be checked'() {
+			const childOne = v('div', { id: 'child-one' });
+			const nodeOne = w(WidgetBase, {}, [childOne]);
+			const nodeTwo = 'text node';
+			const childTwo = v('div');
+			const nodeThree = v('div', { id: 'node-three' }, [childTwo]);
+			const nodeFour = v('div', { id: 'node-four' });
+			const nodes = [nodeOne, nodeTwo, nodeThree, nodeFour];
+			const newGrandChild = v('div', { id: 'grand-child' });
+			const newChildOne = w(WidgetBase, {}, [newGrandChild]);
+			const newChildFour = v('div', { id: 'new-child-four' });
+			const newNodeFour = v('div', { id: 'new-node-four' }, [newChildFour]);
+			const newNodes = decorate(nodes, {
+				modifier: (node: DNode) => {
+					if (node === 'text node') {
+						return 'new text node';
+					}
+
+					if (isVNode(node) || isWNode(node)) {
+						if (node.properties.id === 'child-one') {
+							return newChildOne;
+						}
+
+						if (node.properties.id === 'grand-child') {
+							node.properties = { id: 'new-grand-child-id' };
+						}
+
+						if (node.properties.id === 'node-four') {
+							return newNodeFour;
+						}
+
+						if (node.properties.id === 'node-three' || node.properties.id === 'new-child-four') {
+							return false;
+						}
+					}
+				}
+			});
+			assert.strictEqual(newNodes.length, 4);
+			assert.strictEqual(nodeOne.children[0], newChildOne);
+			assert.strictEqual(newGrandChild.properties.id, 'new-grand-child-id');
+			assert.strictEqual(newNodes[1], 'new text node');
+			assert.isFalse(newNodes[2]);
+			assert.strictEqual(newNodes[3], newNodeFour);
+			assert.isFalse(newNodeFour.children![0]);
 		}
 	},
 	'isVNode and isWNode': {