src/graph.luau

@@ -2,12 +2,13 @@ if not game then script = require "test/relative-string" end
 
 local throw = require(script.Parent.throw)
 local flags = require(script.Parent.flags)
+local on_gc = require(script.Parent.on_gc)()
 
 export type Node<T> = {
     cache: T,
-    derive: () -> T,
-    effects:  { [(unknown) -> ()]: unknown }, -- weak values
-    children: { Node<T> } | false -- weak values
+    effect:  () -> (),
+    children: { Node<T> } | false, -- weak values
+    cleanups: { () -> () } | false
 }
 
 -- flag used to detect when node reference capturing is active
@@ -15,6 +16,8 @@ local reff = false
 -- array of all nodes referenced since above flag was set
 local refs = {} :: { Node<unknown> }
 
+local scopes = { n = 0 } :: { [number]: Node<unknown>, n: number }
+
 local WEAK_VALUES = { __mode = "v" }
 local EVALUATION_ERR = "error while evaluating source:\n\n"
 
@@ -46,6 +49,39 @@ local check_for_yield: <T...>(fn: (T...) -> unknown, T...) -> () do
     end
 end
 
+local function get_scope(): Node<unknown>
+    return scopes[scopes.n]
+end
+
+local function open_scope(node: Node<unknown>)
+    local n = scopes.n + 1
+    scopes.n = n
+    scopes[n] = node
+end
+
+local function close_scope()
+    local n = scopes.n
+    scopes.n = n - 1
+    scopes[n] = nil
+end
+
+local function add_cleanup(node: Node<unknown>, cleanup: () -> ())
+    if node.cleanups then
+        table.insert(node.cleanups, cleanup)
+    else
+        node.cleanups = { cleanup }
+    end
+end
+
+local function run_cleanups(node: { cleanups: { () -> () } | false})
+    if node.cleanups then
+        for _, fn in next, node.cleanups do
+            fn()
+        end
+        table.clear(node.cleanups)
+    end
+end
+
 --[[
 
 Each node side-effect is registered with a corresponding weak key.
@@ -57,21 +93,12 @@ The weak key is passed as an argument to its side-effect callback.
 
 ]]
 
-local function set_effect<T>(node: Node<unknown>, fn: (T) -> (), key: T)
-    node.effects[fn :: () -> ()] = key
+local function set_effect<T>(node: Node<unknown>, fn: () -> ())
+    node.effect = fn
 end
 
-local function run_effects(node: Node<unknown>)
-    if flags.strict then -- run effects twice if strict
-        for effect, key in next, node.effects do
-            effect(key)
-            effect(key)
-        end
-    else
-        for effect, key in next, node.effects do
-            effect(key)
-        end
-    end
+local function run_effect(node: Node<unknown>)
+    node.effect()
 end
 
 -- retrieves a node's cached value
@@ -91,15 +118,31 @@ local function set_child(parent: Node<unknown>, child: Node<unknown>)
     end
 end
 
+local function create_and_open_scope(node: Node<unknown>)
+    local parent = scopes[scopes.n]
+    if parent then
+        set_child(parent, node)
+        node.effect = function()
+            return parent
+        end
+    else
+        node.cleanups = {}
+        local cleanups = node.cleanups :: { () -> () }
+        on_gc(node, function()
+            run_cleanups({ cleanups = cleanups })
+        end)
+    end
+    open_scope(node)
+end
+
 -- runs node effects, recalculates descendants and runs descendant effects
 local function update(node: Node<unknown>)
-    run_effects(node)
+    open_scope(node)
+    run_cleanups(node)
+    run_effect(node)
+    close_scope()
     if node.children then
-        local strict = flags.strict
-
         for _, child in node.children do
-            if strict then check_for_yield(child.derive) end
-            child.cache = child.derive()
             update(child)
         end
     end
@@ -113,7 +156,9 @@ end
 
 -- links two nodes as parent-child with a function to compute a new value for child
 local function link<T>(parent: Node<unknown>, child: Node<T>, derive: () -> T)
-    child.derive = derive
+    child.effect = function()
+        child.cache = derive()
+    end
     set_child(parent, child)
 end
 
@@ -121,8 +166,6 @@ end
 local function capture<T, U>(fn: (U?) -> T, arg: U?): ({ Node<unknown> }, T)
     if reff then throw("recursive capture detected") end
 
-    if flags.strict then check_for_yield(fn, arg) end
-
     table.clear(refs)
     reff = true
 
@@ -142,10 +185,12 @@ local function capture<T, U>(fn: (U?) -> T, arg: U?): ({ Node<unknown> }, T)
 end
 
 -- captures and links any detected nodes
-local function capture_and_link<T>(child: Node<T>, fn: () -> T): T
-    local nodes, value = capture(fn, nil)
+local function capture_and_link<T>(child: Node<T>, derive: () -> T): T
+    local nodes, value = capture(derive, nil)
 
-    child.derive = fn
+    child.effect = function()
+        child.cache = derive()
+    end
     for _, parent: Node<unknown> in next, nodes do
         set_child(parent, child)
     end
@@ -156,9 +201,9 @@ end
 local function create<T>(value: T): (Node<T>, () -> T)
     local node = {
         cache = value,
-        derive  = function() return nil :: any end,
-        effects = setmetatable({}, WEAK_VALUES) :: any,
-        children = false :: false
+        effect = function() end,
+        children = false :: false,
+        cleanups = false :: false
     }
 
     local function read_node_value()
@@ -169,10 +214,17 @@ local function create<T>(value: T): (Node<T>, () -> T)
 end
 
 return table.freeze {
+    create_and_open_scope = create_and_open_scope,
+    open_scope = open_scope,
+    close_scope = close_scope,
+    get_scope = get_scope,
+    add_cleanup = add_cleanup,
+    run_cleanups = run_cleanups,
     set_effect = set_effect,
     get = get,
     set = set,
     link = link,
+    set_child = set_child,
     capture = capture,
     capture_and_link = capture_and_link,
     create = create :: (<T>(value: T) -> (Node<T>, () -> T)) & (<T>() -> (Node<T>, () -> T)),