Fix Fragment in For

yishn · May 18, 2024 · acc86a5 · acc86a5
1 parent 535b4ab
commit acc86a5
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 45 deletions.
diff --git a/src/intrinsic/For.test.tsx b/src/intrinsic/For.test.tsx
@@ -1,7 +1,7 @@
 import { GlobalRegistrator } from "@happy-dom/global-registrator";
 import { afterEach, beforeEach, test } from "node:test";
 import assert from "node:assert";
-import { For, useSignal, useRef, If } from "../mod.js";
+import { For, useSignal, useRef, If, ElseIf } from "../mod.js";
 import { useScope } from "../scope.js";
 
 beforeEach(() => {
@@ -99,3 +99,50 @@ test("For in If", async () => {
     "Does not leak memory",
   );
 });
+
+test("Fragment and If in For", async () => {
+  const s = useScope();
+  const [list, setList] = useSignal<string[]>(["a", "b", "c"]);
+  const ulRef = useRef<HTMLUListElement>();
+
+  document.body.append(
+    ...(
+      <ul ref={ulRef}>
+        <For each={list}>
+          {(item) => (
+            <>
+              <If condition={() => item() === "b"}>
+                <li>{item}</li>
+              </If>
+              <ElseIf condition={() => item() === "d"}>
+                <li>special</li>
+              </ElseIf>
+            </>
+          )}
+        </For>
+      </ul>
+    ).build()(),
+  );
+
+  const effectsCount = s._effects.length;
+  const subscopesCount = s._subscopes.length;
+
+  assert.strictEqual(ulRef()!.children.length, 1);
+  assert.strictEqual(ulRef()!.children[0].textContent, "b");
+
+  setList(["a", "c"]);
+  assert.strictEqual(ulRef()!.children.length, 0);
+
+  setList(["b", "b", "b"]);
+  assert.strictEqual(ulRef()!.children.length, 3);
+
+  setList(["b", "b", "b", "d"]);
+  assert.strictEqual(ulRef()!.children.length, 4);
+  assert.strictEqual(ulRef()!.children[3].textContent, "special");
+
+  assert.deepStrictEqual(
+    [s._effects.length, s._subscopes.length],
+    [effectsCount, subscopesCount],
+    "Does not leak memory",
+  );
+});
diff --git a/src/intrinsic/For.ts b/src/intrinsic/For.ts
@@ -4,6 +4,7 @@ import {
   Signal,
   SignalLike,
   useEffect,
+  useMemo,
   useSignal,
   useSubscope,
 } from "../scope.js";
@@ -32,17 +33,19 @@ export const For = <T>(props: {
     const items = MaybeSignal.upgrade(props.each ?? []);
     const anchor = renderer._node(() => document.createComment(""));
     const keyFn = props.key ?? ((_, i) => i);
-    const [nodes, setNodes] = useSignal<Node[]>([anchor], { force: true });
+    const [nodes, setNodes] = useSignal<(SignalLike<Node[]> | undefined)[]>(
+      [],
+      { force: true },
+    );
     const keyMap = new Map<unknown, KeyMeta>();
     const mutationResult = useArrayMutation(items, keyFn);
 
     const lookForAnchor = (index: number): Node => {
       for (let i = index - 1; i >= 0; i--) {
-        const key = keyFn(items()[index - 1], index - 1);
-        const nodes = keyMap.get(key)?._subnodes;
+        const subnodes = nodes()[i];
 
-        if (nodes && nodes().length > 0) {
-          return nodes()[nodes().length - 1];
+        if (subnodes && subnodes().length > 0) {
+          return subnodes()[subnodes().length - 1];
         }
       }
 
@@ -56,11 +59,7 @@ export const For = <T>(props: {
           _destroy?.();
 
           setNodes((nodes) => {
-            const index = nodes.indexOf(_subnodes?.()[0]!);
-            if (index > 0) {
-              nodes.splice(index, _subnodes?.().length);
-            }
-
+            nodes.splice(mutation._index, 1);
             return nodes;
           });
 
@@ -89,47 +88,40 @@ export const For = <T>(props: {
 
             _subnodes = props.children?.(item, index, items).build();
 
-            const itemAnchor = lookForAnchor(mutation._index);
+            let itemAnchor = lookForAnchor(mutation._index);
 
             setNodes((nodes) => {
-              const anchorIndex = nodes.indexOf(itemAnchor);
-              if (anchorIndex >= 0) {
-                nodes.splice(anchorIndex + 1, 0, ...(_subnodes?.() ?? []));
-              }
-
+              nodes.splice(mutation._index, 0, _subnodes);
               return nodes;
             });
 
-            _subnodes?.().forEach((node) =>
-              itemAnchor.parentNode?.insertBefore(node, itemAnchor.nextSibling),
-            );
+            _subnodes?.().forEach((node) => {
+              itemAnchor.parentNode?.insertBefore(node, itemAnchor.nextSibling);
+              itemAnchor = node;
+            });
           });
 
           keyMap.set(mutation._key, { _subnodes, _destroy: destroy });
         } else if (mutation._type == "m") {
           const { _subnodes } = keyMap.get(mutation._key) ?? {};
-          const itemAnchor = lookForAnchor(mutation._to);
 
           setNodes((nodes) => {
-            const index = nodes.indexOf(_subnodes?.()[0]!);
-            if (index >= 0) {
-              nodes.splice(index, _subnodes?.().length);
-            }
-
-            const anchorIndex = nodes.indexOf(itemAnchor);
-            if (anchorIndex >= 0) {
-              nodes.splice(anchorIndex + 1, 0, ...(_subnodes?.() ?? []));
-            }
-
+            nodes.splice(mutation._from, 1);
+            nodes.splice(mutation._to, 0, _subnodes);
             return nodes;
           });
 
-          _subnodes?.().forEach((node) =>
-            itemAnchor.parentNode?.insertBefore(node, itemAnchor.nextSibling),
-          );
+          let itemAnchor = lookForAnchor(mutation._to);
+
+          _subnodes?.().forEach((node) => {
+            itemAnchor.parentNode?.insertBefore(node, itemAnchor.nextSibling);
+            itemAnchor = node;
+          });
         }
       }
     }, [mutationResult]);
 
-    return nodes;
+    return useMemo(() =>
+      [anchor as Node].concat(nodes().flatMap((nodes) => nodes?.() ?? [])),
+    );
   });
diff --git a/src/intrinsic/Fragment.ts b/src/intrinsic/Fragment.ts
@@ -28,13 +28,15 @@ export const Fragment: FunctionalComponent<{
   children?: Children;
 }> = ({ children }) =>
   createTemplate(() => {
-    return useMemo(() =>
-      !Array.isArray(children)
-        ? children == null
-          ? []
-          : typeof children == "object"
-            ? children.build()()
-            : Text({ text: children }).build()()
-        : children.flatMap((children) => Fragment({ children }).build()()),
-    );
+    const arr = !Array.isArray(children)
+      ? children == null
+        ? []
+        : [
+            typeof children == "object"
+              ? children.build()
+              : Text({ text: children }).build(),
+          ]
+      : children.flatMap((children) => Fragment({ children }).build());
+
+    return useMemo(() => arr.flatMap((signal) => signal()));
   });
diff --git a/src/intrinsic/If.ts b/src/intrinsic/If.ts
@@ -45,7 +45,7 @@ export const ElseIf: FunctionalComponent<{
   return runWithRenderer({ _ifConditions: [] }, () =>
     createTemplate(() => {
       const anchor = renderer._node(() => document.createComment(""));
-      const [nodes, setNodes] = useSignal<Node[]>([anchor]);
+      const [nodes, setNodes] = useSignal<Node[]>([anchor], { force: true });
       const template = useMemo(() =>
         condition() ? Fragment({ children: props.children }) : null,
       );