From 52ea2e46900d10fc3199719964eb2b6f24ac2fc2 Mon Sep 17 00:00:00 2001
From: Aaron Smith <83140718+centau@users.noreply.github.com>
Date: Mon, 11 Sep 2023 17:41:05 +0100
Subject: [PATCH]

---
 src/bind.luau   |  41 +++-------------
 src/derive.luau |   8 ++--
 src/graph.luau  | 122 ++++++++++++++++++++++++++++++++++--------------
 src/maps.luau   |  20 ++++----
 src/root.luau   |   5 ++
 src/spring.luau |   4 +-
 src/watch.luau  |  12 ++---
 test/tests.luau |  33 ++++++++++---
 todo.md         |   4 +-
 9 files changed, 147 insertions(+), 102 deletions(-)

diff --git a/src/bind.luau b/src/bind.luau
index d5ba90d..fe1605d 100644
--- a/src/bind.luau
+++ b/src/bind.luau
@@ -9,38 +9,7 @@ local create_node = graph.create_node
 local get_scope = graph.get_scope
 local open_scope = graph.open_scope
 local close_scope = graph.close_scope
-local track = graph.track
-local add_child = graph.add_child
-
---[[
-
-Roblox instances in Luau are referenced using a kind of userdata proxy,
-this proxy can be garbage collected independently from the actual instance, even
-if the instance is still parented. Since reactive bindings allow the garbage
-collection of instances, this proxy can can garbage collected while the instance
-is still parented, causing the binding to be lost and no longer update the
-instance on changes.
-
-Vide's solution to this is to hold the proxy in memory as long as the instance
-is parented to the datamodel by using `GetPropertyChanged("Parent")` to add or
-remove the proxy from a table whose sole purpose is to strongly reference
-proxies.
-
-todo: investigate behavior in case B is parented to A, and A has no parent or reference, and B has a binding.
-
-]]
-
-type Binding = {
-    instance: Instance,
-    property: string,
-    source: () -> unknown
-}
-
-local function binder(b: Binding)
-    (b.instance :: any)[b.property] = b.source()
-end
-
-
+local set_owner = graph.set_owner
 
 -- todo: replace with throw's method
 local root do
@@ -65,7 +34,7 @@ local function traceback(skips: number) -- ensures trace begins outside of any v
     return debug.traceback(nil, s)
 end
 
-function create_binding<T>(updater: (T) -> (), binding_data: T)
+function create_binding<T>(updater: (T) -> T, binding_data: T)
     -- if flags.strict then
     --     -- wrap setter in function with stack inspection for better error msgs
     --     local fn = setter
@@ -84,7 +53,7 @@ function create_binding<T>(updater: (T) -> (), binding_data: T)
     local owner = get_scope()
     assert(owner)
 
-    add_child(owner, binding)
+    set_owner(binding, owner)
     open_scope(binding)
 
     updater(binding_data)
@@ -100,6 +69,7 @@ type PropertyBinding = {
 
 local function update_property(p: PropertyBinding)
     (p.instance :: any)[p.property] = p.source()
+    return p
 end
 
 type ParentBinding = {
@@ -109,6 +79,7 @@ type ParentBinding = {
 
 local function update_parent(p: ParentBinding)
     p.instance.Parent = p.parent()
+    return p
 end
 
 type ChildrenBinding = {
@@ -145,6 +116,8 @@ local function update_children(p: ChildrenBinding)
 
     table.clear(cur_children_set) -- clear cache, preserve capacity
     p.cur_children_set, p.new_children_set = new_child_set, cur_children_set
+
+    return p
 end
 
 return {
diff --git a/src/derive.luau b/src/derive.luau
index 9aaa83e..3b7972f 100644
--- a/src/derive.luau
+++ b/src/derive.luau
@@ -2,8 +2,7 @@ if not game then script = require "test/relative-string" end
 
 local graph = require(script.Parent.graph)
 local create_node = graph.create_node
-local add_child = graph.add_child
-local update = graph.update
+local set_owner = graph.set_owner
 local track = graph.track
 local get_scope = graph.get_scope
 local open_scope = graph.open_scope
@@ -15,11 +14,10 @@ local function derive<T>(fn: () -> T): () -> T
 
     local node = create_node((false :: any) :: T)
     node.effect = function()
-        node.cache = fn()
+        return fn()
     end
 
-    add_child(owner, node)
-
+    set_owner(node, owner)
     open_scope(node)
 
     node.cache = fn()
diff --git a/src/graph.luau b/src/graph.luau
index 578216f..29baa5b 100644
--- a/src/graph.luau
+++ b/src/graph.luau
@@ -10,9 +10,10 @@ export type StartNode<T> = {
 
 export type Node<T> =  {
     cache: T,
+    owner: Node<T> | false,
     parents: { StartNode<T> },
     children: { Node<T> } | false,
-    effect:  (T) -> () | false,
+    effect:  ((T) -> T) | false,
     cleanups: { () -> () } | false,
 }
 
@@ -56,16 +57,26 @@ local function get_scope(): Node<unknown>?
     return scopes[scopes.n]
 end
 
+
 local function add_child<T>(parent: StartNode<any>, child: Node<any>)
     if parent.children then
         table.insert(parent.children :: { Node<T> }, child)
     else
         parent.children = { child }
     end
-
+    
     table.insert(child.parents, parent)
 end
 
+local function set_owner(node: Node<any>, owner: Node<any>)
+    node.owner = owner
+    if owner.children then
+        table.insert(owner.children :: { Node<any> }, node)
+    else
+        owner.children = { node }
+    end
+end
+
 local function open_scope<T>(node: Node<T>)
     local n = scopes.n + 1
     scopes.n = n
@@ -95,20 +106,18 @@ local function run_cleanups<T>(node: Node<T>)
     end
 end
 
-local function run_effect<T>(node: Node<T>)
-    if node.effect then
-        node.effect(node.cache)
-    end
+local function remove_child<T>(parent: StartNode<T>, child: Node<T>)
+    local children = parent.children :: {}
+    local idx = table.find(children :: {}, child)
+    
+    local n = #children
+    children[idx] = children[n]
+    children[n] = nil
 end
 
 local function unparent<T>(node: Node<T>)
     for _, parent in node.parents do
-        local children = parent.children :: {}
-        local idx = table.find(children :: {}, node)
-        
-        local n = #children
-        children[idx] = children[n]
-        children[n] = nil
+        remove_child(parent, node)
     end
 
     table.clear(node.parents)
@@ -117,6 +126,8 @@ end
 local function destroy<T>(node: Node<T>)
     run_cleanups(node)
     unparent(node)
+    if node.owner then remove_child(node.owner, node) end
+    node.owner = false
     local children = node.children :: {}
     if children then
         while children[1] do destroy(children[1]) end
@@ -125,36 +136,77 @@ end
 
 local update_queue = {} :: { Node<any> }
 
--- runs node effects, recalculates descendants and runs descendant effects
-local function rec<T>(node: StartNode<T>)
-    if not node.children then return end
+-- -- runs node effects, recalculates descendants and runs descendant effects
+-- local function rec<T>(node: StartNode<T>)
+--     if not node.children then return end
 
-    local children = node.children :: {}
+--     local children = node.children :: {}
 
-    while children[1] do
-        table.insert(update_queue, children[1])
-        rec(children[1])
-        unparent(children[1])
-    end
+--     while children[1] do
+--         table.insert(update_queue, children[1])
+--         rec(children[1])
+--         unparent(children[1])
+--     end
 
-    table.clear(children)
-end
+--     table.clear(children)
+-- end
+
+-- local function update<T>(node: StartNode<T>)
+--     --assert(#update_queue == 0, "update already in progress")
+--     -- check if recursive update
+--     local first = update_queue[1] == nil
+
+--     rec(node)
+
+--     if first then
+--         for _, n in next, update_queue do
+--             open_scope(n) -- todo
+--             run_cleanups(n)
+--             run_effect(n)
+--             close_scope()
+--         end
+
+--         table.clear(update_queue)
+--     end
+-- end
 
 local function update<T>(node: StartNode<T>)
-    --assert(#update_queue == 0, "update already in progress")
-    -- check if recursive update
-    local first = update_queue[1] == nil
+    local children = node.children :: {}
+    if not children then return end
 
-    rec(node)
+    local n0 = #update_queue
+    local first_update = n0 == 0
+    local n = n0
 
-    if first then
-        for _, n in next, update_queue do
-            open_scope(n) -- todo
-            run_cleanups(n)
-            run_effect(n)
-            close_scope()
+    do
+        local child = children[1]
+        while child do -- todo: case where child in owner context
+            unparent(child)
+
+            n += 1
+            update_queue[n] = child
+
+            child = children[1]
         end
+    end
 
+    for i = n0 + 1, n do
+        local child = update_queue[i]
+
+        local old_value = child.cache
+
+        open_scope(child)
+        run_cleanups(child)
+        local new_value = child.effect and child.effect(old_value)
+        close_scope()
+
+        if old_value ~= new_value then
+            child.cache = new_value
+            update(child)
+        end
+    end
+
+    if first_update then
         table.clear(update_queue)
     end
 end
@@ -169,7 +221,8 @@ end
 local function create_node<T>(value: T): Node<T>
     local node: Node<T> = {
         cache = value,
-        effect = false,
+        owner = false,
+        effect = false :: false,
         cleanups = false :: false,
         parents = {},
         children = false :: false
@@ -200,6 +253,7 @@ return table.freeze {
     get_scope = get_scope,
     get_stack_scope = get_stack_scope,
     add_cleanup = add_cleanup,
+    set_owner = set_owner,
     destroy = destroy,
     run_cleanups = run_cleanups,
     track = track,
diff --git a/src/maps.luau b/src/maps.luau
index 9027a6d..b8eedfb 100644
--- a/src/maps.luau
+++ b/src/maps.luau
@@ -9,6 +9,7 @@ type Scope = graph.Scope
 type StartNode<T> = graph.StartNode<T>
 local create_node = graph.create_node
 local create_start_node = graph.create_start_node
+local set_owner = graph.set_owner
 local track = graph.track
 local update = graph.update
 local get_scope = graph.get_scope
@@ -36,7 +37,6 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
     local output_cache = {} :: Map<K, VO>
     local input_nodes = {} :: Map<K, StartNode<VI>>
     local remove_queue = {} :: { K }
-    local output_array = {} :: { VO }
 
     local scopes = {} :: Map<K, Scope>
 
@@ -71,8 +71,8 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
                     local scope = create_node(false)
                     scopes[i] = scope
 
+                    set_owner(scope, owner)
                     open_scope(scope)
-                    track(owner)
 
                     local node = create_start_node(v)
                     input_nodes[i] = node
@@ -93,8 +93,7 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
 
         close_scope()
 
-        -- output elements
-        table.clear(output_array)
+        local output_array = table.create(#scopes)
         for _, v in next, output_cache do
             table.insert(output_array, v)
         end
@@ -105,7 +104,7 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
 
     local output = create_node(false :: any)
     output.effect = function()
-        update_children(input())
+        return update_children(input())
     end
 
     open_scope(output)
@@ -130,7 +129,6 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
 
     local output_cache = {} :: Map<VI, VO>
     local input_nodes = {} :: Map<VI, StartNode<K>>
-    local output_array = {} :: { VO }
 
     local scopes = {} :: Map<VI, Scope>
 
@@ -147,7 +145,7 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
             end
         end
 
-        open_scope(root)
+        open_scope(owner)
     
         -- process data
         for i, v in next, data do
@@ -159,6 +157,7 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
                 local scope = create_node(false)
                 scopes[v] = scope
 
+                set_owner(scope, owner)
                 open_scope(scope)
 
                 local node = create_start_node(i)
@@ -193,13 +192,10 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
         table.clear(cur_input_cache)
         cur_input_cache_up, new_input_cache_up = new_input_cache, cur_input_cache
 
-        -- output elements
-        table.clear(output_array)
-
+        local output_array = table.create(#scopes)
         for _, v in next, output_cache do
             table.insert(output_array, v)
         end
-
         check_primitives(output_array)
 
         return output_array
@@ -207,7 +203,7 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
 
     local output = create_node(false :: any)
     output.effect = function()
-        update_children(input())
+        return update_children(input())
     end
 
     open_scope(output)
diff --git a/src/root.luau b/src/root.luau
index e495e34..d795d4d 100644
--- a/src/root.luau
+++ b/src/root.luau
@@ -11,6 +11,8 @@ local close_scope = graph.close_scope
 local get_scope = graph.get_scope
 local destroy = graph.destroy
 
+local refs = {}
+
 local function root<T>(fn: () -> T): (T, () -> ())
     --assert(not get_scope())
     local node = create_node(false)
@@ -21,7 +23,10 @@ local function root<T>(fn: () -> T): (T, () -> ())
 
     close_scope()
 
+    refs[node] = true
+
     return v, function()
+        refs[node] = nil
         destroy(node)
     end
 end
diff --git a/src/spring.luau b/src/spring.luau
index 32d3d04..9b4e35f 100644
--- a/src/spring.luau
+++ b/src/spring.luau
@@ -31,7 +31,7 @@ local get_scope = graph.get_scope
 local open_scope = graph.open_scope
 local close_scope = graph.close_scope
 local update = graph.update
-local add_child = graph.add_child
+local set_owner = graph.set_owner
 local track = graph.track
 
 local UPDATE_RATE = 120
@@ -156,7 +156,7 @@ local function spring<T>(source: () -> T, period: number?, damping_ratio: number
     local updater = create_node(false)
     updater.effect = true :: any -- todo
 
-    add_child(owner, updater)
+    set_owner(updater, owner)
     open_scope(updater)
 
     local initial_value = source()
diff --git a/src/watch.luau b/src/watch.luau
index b9a28ea..9213631 100644
--- a/src/watch.luau
+++ b/src/watch.luau
@@ -5,23 +5,21 @@ local create_node = graph.create_node
 local get_scope = graph.get_scope
 local open_scope = graph.open_scope
 local close_scope = graph.close_scope
-local add_child = graph.add_child
-local track = graph.track
+local set_owner = graph.set_owner
 
-local function watch(effect: () -> ())
+local function watch<T>(effect: (T) -> T, initial_value: T)
     local owner = get_scope()
     assert(owner)
 
     local node = create_node(false)
     node.effect = effect
 
-    add_child(owner, node)
-
+    set_owner(node, owner)
     open_scope(node)
 
-    effect()
+    effect(initial_value)
 
     close_scope()
 end
 
-return watch
+return watch :: (<T>(effect: (T) -> T, initial_value: T) -> ()) & ((effect: () -> ()) -> ())
diff --git a/test/tests.luau b/test/tests.luau
index 11c2b49..574bdc6 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -241,7 +241,6 @@ TEST("graph", function()
         CHECK(table.find(get_children(root), scope1 :: Node<any>))
 
         destroy(scope1)
-        CHECK(table.find(get_children(root), scope1 :: Node<any>))
         CHECK(cleaned.scope1)
         CHECK(cleaned.bind1)
         scope1 = NIL
@@ -252,8 +251,6 @@ TEST("graph", function()
         CHECK(#get_children(selected) == 1)
     end
 
-    -- todo: further tests
-
     do CASE "nodes garbage collection"
         local wref = weak { create_node(1) }
         destroy(wref[1])
@@ -330,9 +327,11 @@ TEST("source()", wrap_root(function()
 end))
 
 TEST("derive()", wrap_root(function()
+    local root = vide.root
     local source = vide.source
     local derive = vide.derive
     local watch = vide.watch
+    local cleanup = vide.cleanup
 
     do CASE "derive new value on source change"
         local a = source(1)
@@ -367,7 +366,7 @@ TEST("derive()", wrap_root(function()
         local num = source(0)
 
         local is_even = derive(function()
-            return num() % 2 == 0
+            return bit32.band(num(), 0b01) == 0
         end)
 
         local count = 0
@@ -417,15 +416,37 @@ TEST("derive()", wrap_root(function()
         CHECK(count == 4)
     end
 
+    do CASE "owner not disconnected"
+        local count = 0
+        local a = source(0)
+
+        local _, destroy = root(function()
+
+            local b = derive(function()
+                cleanup(function()
+                    count += 1
+                end)
+
+                return a()
+            end)
+        end)
+
+        CHECK(count == 0)
+        a(1) -- b clears parents (should not clear owner)
+        CHECK(count == 1)
+        destroy()
+        CHECK(count == 2)
+    end
+
     do CASE "garbage collection"
         -- check that `b` does not allow gc of `a`        
         local a = source(1)
 
-        local b = derive(function()
+        local _b = derive(function()
             return a()
         end)
 
-        b = NIL
+        _b = NIL
 
         local wref = weak { a }
 
diff --git a/todo.md b/todo.md
index 393b7a7..24f6b03 100644
--- a/todo.md
+++ b/todo.md
@@ -3,8 +3,8 @@
 - better error reporting and stack traces in strict mode
 - auto-enable of strict mode depending on compiler optimizaton level
 - investigate if weak table iteration can be invalidated
-- define behavior of `cleanup()` in `untrack()` scopes
-  - 
+- property binding optimization
+  - would no longer allow `cleanup()` usage in binding scopes
 - solution to nested reactivity, see: SolidJS stores
 - SolidJS control flow components
 - equality checking of derived sources