From 0105e33fac658491b4acde0d33b8a8dd93233a56 Mon Sep 17 00:00:00 2001
From: Aaron Smith <83140718+centau@users.noreply.github.com>
Date: Tue, 22 Aug 2023 11:46:56 +0100
Subject: [PATCH] Limit `cleanup()` calls to once per scope

---
 docs/api/reactivity-utility.md | 15 ++++++++++++++-
 docs/api/strict-mode.md        |  1 +
 src/cleanup.luau               | 20 +++++++++++++++++---
 test/tests.luau                | 31 +++++++++++++++++++++++++++++++
 todo.md                        |  1 -
 5 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/docs/api/reactivity-utility.md b/docs/api/reactivity-utility.md
index 604e507..9e441c2 100644
--- a/docs/api/reactivity-utility.md
+++ b/docs/api/reactivity-utility.md
@@ -2,7 +2,7 @@
 
 ## cleanup()
 
-Runs a callback anytime a reactive scope is re-ran.
+Runs a callback anytime a function scope is re-ran.
 
 - **Type**
 
@@ -10,6 +10,19 @@ Runs a callback anytime a reactive scope is re-ran.
     function cleanup(callback: () -> ())
     ```
 
+- **Details**
+
+    The primary purpose of this function is to provide a means of cleaning up
+    side effects caused by source updates and `watch()` updates.
+
+    The stack is inspected to find the function that calls `cleanup()`. The
+    callback passed is called anytime the caller is re-ran, and when the caller
+    finally garbage collects.
+
+    ::: warning
+    Only one `cleanup()` call is allowed per function scope.
+    :::
+
 - **Example**
 
     ```lua
diff --git a/docs/api/strict-mode.md b/docs/api/strict-mode.md
index 0b244f8..c0f1e7e 100644
--- a/docs/api/strict-mode.md
+++ b/docs/api/strict-mode.md
@@ -17,6 +17,7 @@ Currently, strict mode will:
 4. Checks for `indexes()` and `values()` returning primitive values.
 5. Checks for duplicate nested properties at same depth.
 6. Better error reporting and stack traces.
+7. Checks for multiple `cleanup()` calls in the same function scope.
 
 By rerunning sources and watchers, any side-effects are made more apparent.
 This also helps ensure that cleanups are being handled correctly.
diff --git a/src/cleanup.luau b/src/cleanup.luau
index d4a17f1..45f2978 100644
--- a/src/cleanup.luau
+++ b/src/cleanup.luau
@@ -1,17 +1,23 @@
 if not game then script = require "test/relative-string" end
 
+local flags = require(script.Parent.flags)
+local throw = require(script.Parent.throw)
+
 local ref_to_id = {} :: { [string]: number }
 local id_to_ref = {} :: { [number]: string }
 local cleanup_callbacks = {} :: { [number]: () -> () } -- always dense
 local cleanup_lifetime = {} :: { [number]: unknown } -- can be sparse
 setmetatable(cleanup_lifetime :: any, { __mode = "v" })
 
+local debug_caller_to_line = {} :: { [() -> ()]: number }
+setmetatable(debug_caller_to_line, { __mode = "k" })
+
 local manual_mode = {
     caller = false :: false | () -> (),
     callbacks = {} :: { () -> () }
 }
 
--- todo: rare case where mem address is reused by another function on same line
+-- todo: rare case where mem address is reused by another function
 -- does this case handle itself?
 
 local function cleanup_ref(ref: string, lifetime: unknown, callback: () -> ())
@@ -32,12 +38,20 @@ 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 flags.strict then
+        local line = debug.info(2, "l")
+        local cur_line = debug_caller_to_line[lifetime]
+        if cur_line and cur_line ~= line then
+            throw "only one cleanup call is allowed per function scope"
+        end
+        debug_caller_to_line[lifetime] = line
+    end
 
     if manual_mode.caller == lifetime then
         table.insert(manual_mode.callbacks, callback)
     else
-        local ref = tostring(lifetime) .. line
+        local ref = tostring(lifetime)
         cleanup_ref(ref, lifetime, callback)
     end
 end
diff --git a/test/tests.luau b/test/tests.luau
index 80cef78..04a9157 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -547,6 +547,26 @@ TEST("cleanup()", function()
         CHECK(objB.cleaned == 2)
     end
 
+    -- this is not allowed, test to verify behavior anyways
+    do CASE "multiple cleanup"
+        local state = source(1)
+
+        local queue = {}
+    
+        watch(function()
+            state()
+            cleanup(function() table.insert(queue, 1) end)
+            cleanup(function() table.insert(queue, 2) end)
+        end)
+
+        CHECK(testkit.seq(queue, { 1 }))
+        state(2)
+        CHECK(testkit.seq(queue, { 1, 2, 1 }))
+        state(3)
+        CHECK(testkit.seq(queue, { 1, 2, 1, 2, 1 }))
+    end
+
+    --[[
     do CASE "multiple cleanup"
         local state = source(1)
 
@@ -574,6 +594,7 @@ TEST("cleanup()", function()
         --testkit.print2(queue)
         --CHECK(testkit.seq(queue, { 1, 2, 1, 2, 1, 2 }))
     end
+    ]]
 end)
 
 TEST("create()", function()
@@ -1308,6 +1329,7 @@ TEST("strict", function()
     local derive = vide.derive
     local watch = vide.watch
     local indexes, values = vide.indexes, vide.values
+    local cleanup = vide.cleanup
 
     do CASE "error on derived callback yield"
         local state = source(1)
@@ -1410,6 +1432,15 @@ TEST("strict", function()
 
         CHECK(ok)
     end
+
+    do CASE "multiple cleanup per scope"
+        local ok = pcall(function()
+            cleanup(function() end)
+            cleanup(function() end)
+        end)
+
+        CHECK(not ok)
+    end
 end)
 
 local ok = FINISH()
diff --git a/todo.md b/todo.md
index 486383e..96e07f8 100644
--- a/todo.md
+++ b/todo.md
@@ -1,7 +1,6 @@
 # todo
 
 - cleanup codebase
-- 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