Improve cleanup() algorithm

src/cleanup.luau

@@ -1,47 +1,56 @@
 if not game then script = require "test/relative-string" end
 
--- todo: verify correct behavior in non-standard usage
-local free_ids = {} :: { number }
-local id_to_ref = {} :: { [number]: string }
 local ref_to_id = {} :: { [string]: number }
-local cleanup_callbacks = {} :: { [number]: () -> () }
-local cleanup_callers = {} :: { [number]: () -> () }
-
-setmetatable(cleanup_callers :: any, { __mode = "vs" })
+local id_to_ref = {} :: { [number]: string }
+local cleanup_callbacks = {} :: { [number]: () -> () } -- always dense
+local cleanup_callers = {} :: { [number]: () -> () } -- can be sparse
+setmetatable(cleanup_callers :: any, { __mode = "v" })
 
 -- todo: rare case where mem address is reused by another function on same line
+-- does this case handle itself?
 
 local function cleanup(callback: () -> ())
     local caller = debug.info(2, "f") :: () -> ()
     local line = debug.info(2, "l") :: number
-    local ref = tostring(caller) .. "\0" .. line
+    local ref = tostring(caller) .. line
 
     local id = ref_to_id[ref]
 
     if id then
         cleanup_callbacks[id]()
     else
-        id = table.remove(free_ids) or #cleanup_callbacks + 1
-        id_to_ref[id :: any] = ref -- todo
+        id = #cleanup_callbacks + 1
         ref_to_id[ref] = id
+        id_to_ref[id :: any] = ref -- todo
         cleanup_callers[id :: any] = caller -- todo
     end
 
     cleanup_callbacks[id] = callback
 end
 
--- todo: investigate behavior if cleanup called within cleanup
--- todo: verify no memory leakage
-
 local function clean_garbage()
-    for id = 1, #cleanup_callbacks do
+    for id = #cleanup_callbacks, 1, -1 do
         if cleanup_callers[id] == nil then -- caller was garbage collected
             local callback = cleanup_callbacks[id]
 
-            cleanup_callbacks[id] = nil
-            table.insert(free_ids, id)
-            ref_to_id[id_to_ref[id]] = nil
-            id_to_ref[id] = nil
+            do -- swap and pop
+                local max_id = #cleanup_callbacks
+
+                cleanup_callbacks[id] = cleanup_callbacks[max_id]
+                cleanup_callbacks[max_id] = nil
+
+                cleanup_callers[id] = cleanup_callers[max_id]
+                cleanup_callers[max_id] = nil
+
+                local ref = id_to_ref[id]
+                local max_ref = id_to_ref[max_id]
+
+                id_to_ref[id] = max_ref
+                id_to_ref[max_id] = nil
+
+                ref_to_id[max_ref] = id
+                ref_to_id[ref] = nil
+            end
 
             local ok, err: string? = pcall(callback)
             if not ok then warn(`error occured during cleanup: {err}`) end

test/benchmark.luau

@@ -1,8 +1,18 @@
 local BENCH, START = require("test/testkit").benchmark()
 
+-- try prevent inlining by wrapping in a closure referencing an upvalue that
+-- cannot be determined at compile-time
+local function NO_INLINE(fn)
+    local r = math.random()
+    return function(x)
+        local _ = r
+        fn(x)
+    end
+end
+
 local vide = require "src/init"
 
-local N = 2^18 -- 262144
+local N = 2^18 -- 262144[
 
 BENCH("create state", function()
     local cache = table.create(N)
@@ -247,51 +257,61 @@ BENCH("values() all remove", function()
     state(data)
 end)
 
-BENCH("register cleanup", function()
+BENCH("register new cleanup", function()
     local cleanup = vide.cleanup
 
+    local cleaner = function() end
+
     local callers = {}
 
     for i = 1, N do
         callers[i] = function(fn, v)
             fn(v)
-            return i
+            return i -- return unique upvalue to ensure unique closure
         end
     end
 
     for i = 1, START(N) do
-        callers[i](cleanup, function() end)
+        callers[i](cleanup, cleaner)
     end
 end)
 
+-- cleanup cleanups from previous benchmark
+;(collectgarbage :: any)("collect")
+vide.step(0) 
+
 BENCH("repeat cleanup", function()
     local cleanup = vide.cleanup
 
-    local function foo(i)
-        cleanup(function() -- todo: why is this not causing allocations?
-            return i
+    local foo = NO_INLINE(function(i)
+        cleanup(function()
+            -- return i -- uncomment to include overhead of closure creation
         end)
-    end
+    end)
 
     for i = 1, START(N) do
         foo(i)
     end
 end)
 
--- todo: this is sometimes never terminates
+-- cleanup cleanups from previous benchmark
+;(collectgarbage :: any)("collect")
+vide.step(0) 
 
-BENCH("cleanup step", function()
+BENCH("cleanup gc check", function()
     local cleanup = vide.cleanup
 
-    local ref = table.create(N)
+    local cleaner = function() end
+
+    local callers = table.create(N)
 
     for i = 1, N do
-        ref[i] = function()
-            cleanup(function() end)
+        callers[i] = function()
+            cleanup(cleaner)
             return i
         end
 
-        ref[i]()
+        callers[i]()
     end
 
     START(N)
@@ -299,4 +319,10 @@ BENCH("cleanup step", function()
     vide.step(0)
 end)
 
+BENCH("cleanup gc removal", function()
+    START(N)
+
+    vide.step(0) -- cleanup from previous benchmark
+end)
+
 return nil

test/tests.luau

@@ -529,6 +529,14 @@ TEST("cleanup()", function()
 
         do
             stateA = nil :: any
+            gc()
+            vide.step(0)
+        end
+
+        CHECK(objA.cleaned == 2)
+        CHECK(objB.cleaned == 1)
+
+        do
             stateB = nil :: any
             gc()
             vide.step(0)