Add cleanup for indexes() and values()

More testing required.

src/cleanup.luau

@@ -3,34 +3,48 @@ if not game then script = require "test/relative-string" end
 local ref_to_id = {} :: { [string]: number }
 local id_to_ref = {} :: { [number]: string }
 local cleanup_callbacks = {} :: { [number]: () -> () } -- always dense
-local cleanup_callers = {} :: { [number]: () -> () } -- can be sparse
+local cleanup_lifetime = {} :: { [number]: unknown } -- can be sparse
-setmetatable(cleanup_callers :: any, { __mode = "v" })
+setmetatable(cleanup_lifetime :: any, { __mode = "v" })
+
+local manual_mode = {
+    caller = false :: false | () -> (),
+    callbacks = {} :: { () -> () }
+}
 
 -- todo: rare case where mem address is reused by another function on same line
 -- does this case handle itself?
 
-local function cleanup(callback: () -> ())
+local function cleanup_ref(ref: string, lifetime: unknown, callback: () -> ())
-    local caller = debug.info(2, "f") :: () -> ()
-    local line = debug.info(2, "l") :: number
-    local ref = tostring(caller) .. line
-
     local id = ref_to_id[ref]
 
     if id then
         cleanup_callbacks[id]()
+        cleanup_lifetime[id] = lifetime -- rare case where ref is reused while lifetime is nil
     else
         id = #cleanup_callbacks + 1
         ref_to_id[ref] = id
         id_to_ref[id :: any] = ref -- todo
-        cleanup_callers[id :: any] = caller -- todo
+        cleanup_lifetime[id :: any] = lifetime -- todo
     end
 
     cleanup_callbacks[id] = callback
 end
 
+local function cleanup(callback: () -> ())
+    local lifetime = debug.info(2, "f") -- `caller of cleanup() is lifetime of cleanup`
+    local line = debug.info(2, "l")
+
+    if manual_mode.caller == lifetime then
+        table.insert(manual_mode.callbacks, callback)
+    else
+        local ref = tostring(lifetime) .. line
+        cleanup_ref(ref, lifetime, callback)
+    end
+end
+
 local function clean_garbage()
     for id = #cleanup_callbacks, 1, -1 do
-        if cleanup_callers[id] == nil then -- caller was garbage collected
+        if cleanup_lifetime[id] == nil then -- lifetime was garbage collected
             local callback = cleanup_callbacks[id]
 
             do -- swap and pop
@@ -39,8 +53,8 @@ local function clean_garbage()
                 cleanup_callbacks[id] = cleanup_callbacks[max_id]
                 cleanup_callbacks[max_id] = nil
 
-                cleanup_callers[id] = cleanup_callers[max_id]
+                cleanup_lifetime[id] = cleanup_lifetime[max_id]
-                cleanup_callers[max_id] = nil
+                cleanup_lifetime[max_id] = nil
 
                 local ref = id_to_ref[id]
                 local max_ref = id_to_ref[max_id]
@@ -58,4 +72,16 @@ local function clean_garbage()
     end
 end
 
-return function() return cleanup, clean_garbage end
+local manual_cleanup_mode = function(caller: () -> ()?)
+    if caller == nil then
+        local clone = table.clone(manual_mode.callbacks)
+        manual_mode.caller = false
+        table.clear(manual_mode.callbacks)
+        return clone
+    else
+        manual_mode.caller = caller
+    end
+    return manual_mode.callbacks
+end :: ( (caller: (...any) -> ()) -> () ) & ( (nil) -> { () -> () } )
+
+return function() return cleanup, clean_garbage, manual_cleanup_mode, cleanup_ref end

src/maps.luau

@@ -1,8 +1,11 @@
 if not game then script = require "test/relative-string" end
 
+-- todo: more testing needed regarding `cleanup()` usage
+
 local throw = require(script.Parent.throw)
 local flags = require(script.Parent.flags)
 local graph = require(script.Parent.graph)
+local _, _, manual_cleanup_mode, cleanup_ref = require(script.Parent.cleanup)()
 type Node<T> = graph.Node<T>
 local create = graph.create
 local set = graph.set
@@ -28,35 +31,48 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
     local remove_queue = {} :: { K }
     local output_array = {} :: { VO }
 
+    local cleanups = {} :: Map<K, { () -> () }>
+
     local function recompute(data)
         -- queue removed values
-        for k in next, input_cache do
+        for i in next, input_cache do
-            if data[k] == nil then
+            if data[i] == nil then
-                table.insert(remove_queue, k)
+                table.insert(remove_queue, i)
             end
         end
 
         -- remove queued values
-        for _, k in next, remove_queue do
+        for _, i in next, remove_queue do
-            input_cache[k] = nil
+            for _, callback in next, cleanups[i] do
-            output_cache[k] = nil
+                callback() -- todo: pcall
-            input_nodes[k] = nil
+            end
+
+            input_cache[i] = nil
+            output_cache[i] = nil
+            input_nodes[i] = nil
+            cleanups[i] = nil
         end
 
         table.clear(remove_queue)
 
         -- process new or changed values
-        for k, v in next, data do
+        for i, v in next, data do
-            local cv = input_cache[k]
+            local cv = input_cache[i]
 
+            if cv ~= v then
                 if cv == nil then
+                    manual_cleanup_mode(transform)
+
                     local node, get_value = create(v)
-                input_nodes[k] = node
+                    input_nodes[i] = node
-                output_cache[k] = transform(get_value, k)
+                    output_cache[i] = transform(get_value, i)
-                input_cache[k] = v
+                    input_cache[i] = v
-            elseif cv ~= v then
+
-                set(input_nodes[k], v)
+                    cleanups[i] = manual_cleanup_mode(nil)
-                input_cache[k] = v
+                else
+                    set(input_nodes[i], v)
+                    input_cache[i] = v
+                end
             end
         end
 
@@ -70,12 +86,12 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
         return output_array
     end
 
+    local output, output_get = create(nil :: any)
+
     local function derive()
         return recompute(input())
     end
 
-    local output, output_get = create(nil :: any)
-
     local nodes, value = capture(input)
 
     for _, node in next, nodes do
@@ -84,6 +100,14 @@ local function indexes<K, VI, VO>(input: () -> Map<K, VI>, transform: (() -> VI,
 
     output.cache = recompute(value)
 
+    cleanup_ref(tostring(output), output, function()
+        for _, callbacks in next, cleanups do
+            for _, callback in next, callbacks do
+                callback() -- todo: pcall
+            end
+        end
+    end)
+
     return output_get
 end
 
@@ -96,6 +120,8 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
     local input_nodes = {} :: Map<VI, Node<K>>
     local output_array = {} :: { VO }
 
+    local cleanups = {} :: Map<VI, { () -> () }>
+
     local function recompute(data: Map<K, VI>)
         local cur_input_cache, new_input_cache = cur_input_cache_up, new_input_cache_up
 
@@ -116,9 +142,13 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
             local cv = cur_input_cache[v]
             
             if cv == nil then
+                manual_cleanup_mode(transform)
+
                 local node, get_value = create(i)
                 input_nodes[v] = node
                 output_cache[v] = transform(v, get_value)
+
+                cleanups[v] = manual_cleanup_mode(nil)
             else
                 if cv ~= i then
                     set(input_nodes[v], i)
@@ -129,8 +159,13 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
 
         -- remove old values
         for v in next, cur_input_cache do
+            for _, callback in next, cleanups[v] do
+                callback() -- todo: pcall
+            end
+
             output_cache[v] = nil
             input_nodes[v] = nil
+            cleanups[v] = nil
         end
 
         -- update buffer cache
@@ -147,12 +182,12 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
         return output_array
     end
 
+    local output, output_get = create(nil :: any)
+
     local function derive()
         return recompute(input())
     end
 
-    local output, output_get = create(nil :: any)
-
     local nodes, value = capture(input)
 
     for _, node in next, nodes do
@@ -162,6 +197,14 @@ local function values<K, VI, VO>(input: () -> Map<K, VI>, transform: (VI, () ->
 
     output.cache = recompute(value)
 
+    cleanup_ref(tostring(output), output, function()
+        for _, callbacks in next, cleanups do
+            for _, callback in next, callbacks do
+                callback() -- todo: pcall
+            end
+        end
+    end)
+
     return output_get
 end
 

test/tests.luau

@@ -869,9 +869,13 @@ TEST("create()", function()
     end
 end)
 
+-- todo: gc and cleanup call check for removed element
+
 TEST("indexes()", function()
+    local create = vide.create
     local source = vide.source
     local indexes = vide.indexes
+    local cleanup = vide.cleanup
 
     do CASE "use source"
         local input = source { 1, 2, 3 }
@@ -909,24 +913,35 @@ TEST("indexes()", function()
     do CASE "removal reflected"
         local input = source { 1, 2, 3 }
 
+        local destroyed = false
+
         local output = indexes(input, function(v, i)
-            return v 
+            local text = create "TextLabel" {
+                Text = function() return tostring(v()) end
+            }
+
+            cleanup(function()
+                destroyed = true
+            end)
+
+            return text
         end)
 
         input { 1, 2 }
 
         local t = output()
 
-        CHECK(t[1]() == 1)
+        CHECK(t[1].Text == "1")
-        CHECK(t[2]() == 2)
+        CHECK(t[2].Text == "2")
-        CHECK(t[3] == nil)
+        CHECK(t[3] == nil :: any)
+        CHECK(destroyed == true)
     end
 
     do CASE "garbage collection"
         do -- check that `output` does not allow gc of `input`
             local input = source {}
 
-            local _derived = indexes(input, function(i, v)
+            local _derived = indexes(input, function(v, i)
                 return v
             end)
 
@@ -941,8 +956,8 @@ TEST("indexes()", function()
         do -- check that `input` allows gc of `output`
             local input = source {}
 
-            local output = indexes(input, function(i, v)
+            local output = indexes(input, function(v, i)
-                return i, v
+                return v, i
             end)
 
             local wref = weak { output }
@@ -953,11 +968,54 @@ TEST("indexes()", function()
             CHECK(not wref[1])
         end
     end
+
+    do CASE "cleanup"
+        local input = source { 1, 2, 3 }
+
+        local count = table.create(3, 0)
+        local unrelated_count = 0
+
+        local function unrelated()
+            cleanup(function()
+                unrelated_count += 1
+            end)
+        end
+        
+        local output = indexes(input, function(v, i)
+            -- check that overriden cleanup scopes don't affect cleanup calls
+            -- in other function scopes
+            unrelated()
+
+            cleanup(function()
+                count[i] += 1
+            end)
+
+            return {}
+        end)
+
+        output()
+
+        CHECK(count[1] == 0)
+        CHECK(count[2] == 0)
+        CHECK(count[3] == 0)
+        CHECK(unrelated_count == 2)
+
+        output = nil :: any
+        gc()
+        vide.step(0)
+
+        CHECK(count[1] == 1)
+        CHECK(count[2] == 1)
+        CHECK(count[3] == 1)
+        CHECK(unrelated_count == 2)
+    end
 end)
 
 TEST("values()", function()
+    local create = vide.create
     local source = vide.source
     local values = vide.values
+    local cleanup = vide.cleanup
 
     do CASE "use source"
         local input = source { 1, 2, 3 }
@@ -995,17 +1053,28 @@ TEST("values()", function()
     do CASE "removal reflected"
         local input = source { 1, 2, 3 }
 
+        local destroyed = false
+
         local output = values(input, function(v, i)
-            return { v = v, i = i }
+            local text = create "TextLabel" {
+                Text = tostring(v)
+            }
+
+            cleanup(function()
+                destroyed = true
+            end)
+
+            return text
         end)
 
         input { 1, 2 }
 
         local t = output()
 
-        CHECK(t[1].v == 1)
+        CHECK(t[1].Text == "1")
-        CHECK(t[2].v == 2)
+        CHECK(t[2].Text == "2")
-        CHECK(t[3] == nil)
+        CHECK(t[3] == nil :: any)
+        CHECK(destroyed == true)
     end
 
     do CASE "removal reflected 2"
@@ -1025,6 +1094,47 @@ TEST("values()", function()
         CHECK(t[2] == nil)
         CHECK(t[3] == nil)
     end
+
+    do CASE "cleanup"
+        local input = source { 1, 2, 3 }
+
+        local count = table.create(3, 0)
+        local unrelated_count = 0
+
+        local function unrelated()
+            cleanup(function()
+                unrelated_count += 1
+            end)
+        end
+        
+        local output = values(input, function(v, i)
+            -- check that overriden cleanup scopes don't affect cleanup calls
+            -- in other function scopes
+            unrelated()
+
+            cleanup(function()
+                count[i()] += 1
+            end)
+
+            return {}
+        end)
+
+        output()
+
+        CHECK(count[1] == 0)
+        CHECK(count[2] == 0)
+        CHECK(count[3] == 0)
+        CHECK(unrelated_count == 2)
+
+        output = nil :: any
+        gc()
+        vide.step(0)
+
+        CHECK(count[1] == 1)
+        CHECK(count[2] == 1)
+        CHECK(count[3] == 1)
+        CHECK(unrelated_count == 2)
+    end
 end)
 
 TEST("spring()", function()

todo.md

@@ -2,6 +2,7 @@
 
 - cleanup codebase
 - way to optionally cleanup `values()` and `indexes()` when they gc
+- limit `cleanup()` call to once per function scope?
 - better error reporting and stack traces in strict mode
 - auto-enable of strict mode depending on compiler optimizaton level
 - check smoothness of spring at high frequency updates