From b85419088c1a0e4e96e090b1062f3841e41c6ee8 Mon Sep 17 00:00:00 2001
From: alicesaidhi <166900055+alicesaidhi@users.noreply.github.com>
Date: Tue, 20 Aug 2024 13:40:29 +0200
Subject: [PATCH] allow fragments in implicit effects for children (#29)

* fix fragments in effects

* add test case

* more descriptive test
---
 src/apply.luau  |  3 ++-
 src/bind.luau   | 13 +++++++++++--
 test/tests.luau | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/src/apply.luau b/src/apply.luau
index 13e8cab..0f4586e 100644
--- a/src/apply.luau
+++ b/src/apply.luau
@@ -11,6 +11,7 @@ local graph = require(script.Parent.graph)
 type Node<T> = graph.Node<T>
 
 type Array<V> = { V }
+type ArrayOrV<V> = {ArrayOrV<V>} | V
 type Map<K, V> = { [K]: V }
 
 local free_caches: {
@@ -124,7 +125,7 @@ local function apply<T>(instance: T & Instance, properties: { [unknown]: unknown
                 end    
             elseif type(property) == "number" then
                 if type(value) == "function" then
-                    bind.children(instance, value :: () -> Instance | Array<Instance>) -- bind children
+                    bind.children(instance, value :: () -> ArrayOrV<Instance>) -- bind children
                 elseif type(value) == "table" then
                     if is_action(value) then
                         table.insert(actions[(value :: any).priority], (value :: any).callback :: () -> ()) -- add action to buffer
diff --git a/src/bind.luau b/src/bind.luau
index dd34ce3..3a016a6 100644
--- a/src/bind.luau
+++ b/src/bind.luau
@@ -38,6 +38,7 @@ type ChildrenBinding = {
     children: () -> Instance | { Instance }
 }
 
+type ArrayOrV<V> = V | { V }
 local function update_children_effect(p: ChildrenBinding)
     local cur_children_set: { [Instance]: true } = p.cur_children_set -- cache of all children parented before update
     local new_child_set: { [Instance]: true } = p.new_children_set -- cache of all children parented after update
@@ -48,8 +49,14 @@ local function update_children_effect(p: ChildrenBinding)
         new_children = { new_children }
     end
 
-    if new_children then
-        for _, child in next, new_children :: { Instance } do
+    local function process_child(child: ArrayOrV<Instance>)
+        if type(child) == "table" then
+            for _, child in next, child do
+                process_child(child)
+            end
+        else
+            if new_child_set[child] then return end -- stops redundant reparenting
+
             new_child_set[child] = true -- record child set from this update
             if not cur_children_set[child] then
                 child.Parent = p.instance -- if child wasn't already parented then parent it
@@ -59,6 +66,8 @@ local function update_children_effect(p: ChildrenBinding)
         end
     end
 
+    process_child(new_children)
+
     for child in next, cur_children_set do
         child.Parent = nil -- unparent all children that weren't in the new children set
     end
diff --git a/test/tests.luau b/test/tests.luau
index c0f69d1..adda8da 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -945,6 +945,42 @@ TEST("create()", wrap_root(function()
         CHECK((f2 :: any).a == 2)
     end
 
+    do CASE "nested children effect"
+        local a = create "Frame" { Name = "a" }
+        local b = create "Frame" { Name = "b" }
+        local c = create "Frame" { Name = "c" }
+        local d = create "Frame" { Name = "d" }
+        local e = create "Frame" { Name = "e" }
+        
+        local children = source {
+            a,
+            { b, c, { d } },
+            { { e } }
+        }
+        
+        local obj = create "Frame" {
+            children
+        }
+
+        CHECK(obj:FindFirstChild("a"))
+        CHECK(obj:FindFirstChild("b"))
+        CHECK(obj:FindFirstChild("c"))
+        CHECK(obj:FindFirstChild("d"))
+        CHECK(obj:FindFirstChild("e"))
+
+        children {
+            b,
+            { c, a },
+            { { d } }
+        }
+
+        CHECK(obj:FindFirstChild("a"))
+        CHECK(obj:FindFirstChild("b"))
+        CHECK(obj:FindFirstChild("c"))
+        CHECK(obj:FindFirstChild("d"))
+        CHECK(not obj:FindFirstChild("e"))
+    end
+
     do CASE "garbage collection test"
         local wref