Allow creation of nested tracking scopes

It turns out that handling destruction of a nested tracking scope was
not as difficult as I thought and didn't need much changes

src/graph.luau

@@ -10,10 +10,14 @@ export type StartNode<T> = {
 
 export type Node<T> =  {
     cache: T,
-    effect:  ((T) -> T) | "owner" | "untracked",
+    effect:  ((T) -> T) | false,
     cleanups: { () -> () } | false,
-    parents: { owner: StartNode<T>?, [number]: StartNode<T> },
-    [number]: Node<T>
+
+    owned: { Node<T> } | false,
+    owner: Node<T> | false,
+
+    parents: { StartNode<T> },
+    [number]: Node<T> -- children
 }
 
 -- reactive scope stack
@@ -50,7 +54,7 @@ local function get_owning_scope(): Node<unknown>
     if not scope then
         local caller_name = debug.info(2, "n")
         return throw(`cannot use {caller_name}() in non-reactive scope, must be used within a root() or mount() callback`)
-    elseif scope.effect ~= "owner" then
+    elseif scope.effect then
         throw("reactive scope is not an owning scope; new effects cannot be created in side-effects")
     end
     return scope
@@ -62,8 +66,12 @@ local function add_child<T>(parent: StartNode<any>, child: Node<any>)
 end
 
 local function set_owner(node: Node<any>, owner: Node<any>)
-    node.parents.owner = owner
-    table.insert(owner, node)
+    node.owner = owner
+    if owner.owned then
+        table.insert(owner.owned, node)
+    else
+        owner.owned = { node }
+    end
 end
 
 local function open_scope<T>(node: Node<T>)
@@ -96,18 +104,29 @@ local function run_cleanups<T>(node: Node<T>)
     end
 end
 
+local function find_and_swap_pop<T>(t: { T }, v: T)
+    local idx = table.find(t, v)
+    assert(idx, "value not found")
+    local n = #t
+    t[idx] = t[n]
+    t[n] = nil
+end
+
 local function remove_child<T>(parent: StartNode<T>, child: Node<T>)
-    local idx = table.find(parent, child)
-    assert(idx, "child not found")
-    local n = #parent
-    parent[idx] = parent[n]
-    parent[n] = nil
+    find_and_swap_pop(parent, child)
+end
+
+local function remove_owner<T>(node: Node<T>)
+    local owner = node.owner :: Node<T>
+    if node.owner and owner.owned then
+        find_and_swap_pop(owner.owned, node)
+    end
 end
 
 local function unparent<T>(node: Node<T>)
     local parents = node.parents
 
-    for i, parent in ipairs(parents) do
+    for i, parent in next, parents do
         remove_child(parent, node)
         parents[i] = nil
     end
@@ -116,15 +135,21 @@ end
 local function destroy<T>(node: Node<T>)
     run_cleanups(node)
     unparent(node)
+    remove_owner(node)
 
-    if node.parents.owner then
-        remove_child(node.parents.owner, node)
-        node.parents.owner = nil
+    if node.owned then
+        local owned = node.owned
+        while owned[1] do destroy(owned[1]) end
     end
-
     while node[1] do destroy(node[1]) end
 end
 
+local function destroy_owned<T>(node: Node<T>)
+    if node.owned then
+        while node.owned[1] do destroy(node.owned[1]) end
+    end
+end
+
 local update_queue = { n = 0 } :: { n: number, [number]: Node<any> }
 
 local function evaluate_node<T>(node: Node<T>)
@@ -132,6 +157,8 @@ local function evaluate_node<T>(node: Node<T>)
 
     if flags.strict then
         run_cleanups(node)
+        destroy_owned(node)
+
         open_scope(node)
 
         local ok, err = check_for_yield(node.effect :: (T) -> T, cur_value)
@@ -141,7 +168,9 @@ local function evaluate_node<T>(node: Node<T>)
         if not ok then throw(err :: string) end
     end
 
-    run_cleanups(node) -- todo: move in scope?
+    run_cleanups(node)
+    destroy_owned(node)
+
     open_scope(node)
 
     local ok, new_value = pcall(node.effect :: (T) -> T, cur_value)
@@ -204,11 +233,15 @@ local function track<T>(node: StartNode<T>)
     end
 end
 
-local function create_node<T>(value: T, effect: "owner" | (T) -> T): Node<T>
+local function create_node<T>(value: T, effect: false | (T) -> T): Node<T>
     return {
         cache = value,
         effect = effect,
         cleanups = false,
+
+        owner = false,
+        owned = false,
+
         parents = {},
     }
 end

src/maps.luau

@@ -30,7 +30,7 @@ end
 local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI, K) -> VO): () -> { VO }
     local owner = get_owning_scope()
 
-    local subowner = create_node(false, "owner")
+    local subowner = create_node(false, false)
     set_owner(subowner, owner)
 
     local input_cache = {} :: Map<K, VI>
@@ -67,7 +67,7 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
 
             if cv ~= v then
                 if cv == nil then -- create new scope and run transform
-                    local scope = create_node(false, "owner")
+                    local scope = create_node(false, false)
                     scopes[i] = scope :: Node<any>
 
                     local node = create_start_node(v)
@@ -125,7 +125,7 @@ end
 local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () -> K) -> VO): () -> { VO }
     local owner = get_owning_scope()
 
-    local subowner = create_node(false, "owner")
+    local subowner = create_node(false, false)
     set_owner(subowner, owner)
     
     local cur_input_cache_up = {} :: Map<VI, K>
@@ -156,7 +156,7 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
             local cv = cur_input_cache[v]
             
             if cv == nil then -- create new scope and run transform
-                local scope = create_node(false, "owner")
+                local scope = create_node(false, false)
                 scopes[v] = scope :: Node<any>
 
                 local node = create_start_node(i)

src/root.luau

@@ -11,7 +11,7 @@ local destroy = graph.destroy
 local refs = {}
 
 local function root<T...>(fn: (destroy: () -> ()) -> T...): T...
-    local node = create_node(false, "owner")
+    local node = create_node(false, false)
 
     refs[node] = true -- prevent gc of root node
 

src/switch.luau

@@ -38,7 +38,7 @@ local function switch<T, U>(source: () -> T): (map: Map<T, ((() -> U)?)>) -> ()
                 throw("map must map a value to a function")
             end
 
-            local new_scope = create_node(false, "owner")
+            local new_scope = create_node(false, false)
             last_scope = new_scope :: Node<any>
             
             set_owner(new_scope, owner)

src/untrack.luau

@@ -13,7 +13,7 @@ local function untrack<T>(source: () -> T): T
 
     -- sources are only tracked if the node in scope has an effect
     local effect = scope.effect
-    scope.effect = "untracked"
+    scope.effect = false
 
     local ok, result = pcall(source)