diff --git a/src/cleanup.luau b/src/cleanup.luau index 7c22058..0620626 100644 --- a/src/cleanup.luau +++ b/src/cleanup.luau @@ -1,8 +1,11 @@ if not game then script = require "test/relative-string" end -- todo: verify correct behavior in non-standard usage -local cleanup_callbacks = {} :: { [string]: () -> () } -local cleanup_callers = {} :: { [string]: () -> () } +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" }) @@ -13,30 +16,37 @@ local function cleanup(callback: () -> ()) local line = debug.info(2, "l") :: number local ref = tostring(caller) .. "\0" .. line - local fn = cleanup_callbacks[ref] - if fn then - fn() + local id = ref_to_id[ref] + + if id then + cleanup_callbacks[id]() else - cleanup_callers[ref] = caller + id = table.remove(free_ids) or #cleanup_callbacks + 1 + id_to_ref[id :: any] = ref -- todo + ref_to_id[ref] = id + cleanup_callers[id :: any] = caller -- todo end - cleanup_callbacks[ref] = callback + + cleanup_callbacks[id] = callback end -local buffer = {} +-- todo: investigate behavior if cleanup called within cleanup +-- todo: verify no memory leakage local function clean_garbage() - for ref, callback in next, cleanup_callbacks do - if cleanup_callers[ref] == nil then -- caller was garbage collected - callback() - table.insert(buffer, ref) + for id = 1, #cleanup_callbacks 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 + + local ok, err: string? = pcall(callback) + if not ok then warn(`error occured during cleanup: {err}`) end end end - - for _, ref in next, buffer do - cleanup_callbacks[ref] = nil - end - - table.clear(buffer) end return function() return cleanup, clean_garbage end diff --git a/test/benchmark.luau b/test/benchmark.luau index 0d4e383..171d481 100644 --- a/test/benchmark.luau +++ b/test/benchmark.luau @@ -250,6 +250,23 @@ end) BENCH("register cleanup", function() local cleanup = vide.cleanup + local callers = {} + + for i = 1, N do + callers[i] = function(fn, v) + fn(v) + return i + end + end + + for i = 1, START(N) do + callers[i](cleanup, function() end) + end +end) + +BENCH("repeat cleanup", function() + local cleanup = vide.cleanup + local function foo(i) cleanup(function() -- todo: why is this not causing allocations? return i @@ -261,6 +278,8 @@ BENCH("register cleanup", function() end end) +-- todo: this is sometimes never terminates + BENCH("cleanup step", function() local cleanup = vide.cleanup