diff --git a/src/__tests__/benchmarks/atom.bench.luau b/src/__tests__/benchmarks/atom.bench.luau index 6663caa..f688841 100644 --- a/src/__tests__/benchmarks/atom.bench.luau +++ b/src/__tests__/benchmarks/atom.bench.luau @@ -1,7 +1,7 @@ -- This file is for use by Benchmarker -- https://boatbomber.itch.io/benchmarker -local Package = game:GetService("ReplicatedStorage").Packages.charm +local Package = game:GetService("ReplicatedStorage").DevPackages.Charm local atom = require(Package.atom) local computed = require(Package.computed) @@ -11,19 +11,19 @@ local function parameterGenerator() end local benchmarks = { - atom = function(_, source: any) + atom = function(_, state: any) for _ = 1, 500 do - source(math.random(1, 2)) + state(math.random(1, 2)) end end, - computed = function(_, source: any) + computed = function(_, state: any) computed(function() - return 2 * source() + return 2 * state() end) for _ = 1, 500 do - source(math.random(1, 2)) + state(math.random(1, 2)) end end, } diff --git a/src/__tests__/capture.spec.luau b/src/__tests__/capture.spec.luau index 60b6631..bb19c5e 100644 --- a/src/__tests__/capture.spec.luau +++ b/src/__tests__/capture.spec.luau @@ -1,5 +1,6 @@ local atom = require(script.Parent.Parent.atom) local store = require(script.Parent.Parent.store) +local subscribe = require(script.Parent.Parent.subscribe) return function() it("gracefully handles yielding", function() @@ -29,7 +30,7 @@ return function() expect(result).to.equal(1) end) - it("allows nested captures", function() + it("skips atoms in inner capture calls", function() local a = atom(1) local b = atom(1) local inner @@ -41,8 +42,20 @@ return function() end) assert(outer[a], "outer did not capture a") - assert(outer[b], "outer did not capture b") + assert(outer[b] == nil, "outer captured b") assert(inner[a] == nil, "inner captured a") assert(inner[b], "inner did not capture b") end) + + it("should not capture notified dependents", function() + local state = atom(1) + + subscribe(state, function() end) + + local deps = store.capture(function() + return state(2) + end) + + expect(deps[state]).to.never.be.ok() + end) end diff --git a/src/__tests__/init.spec.luau b/src/__tests__/init.spec.luau index a7878dc..5a4c1fc 100644 --- a/src/__tests__/init.spec.luau +++ b/src/__tests__/init.spec.luau @@ -13,5 +13,7 @@ return function() afterEach(function() table.clear(store.listeners :: any) + table.clear(store.capturing.stack) + store.capturing.index = 0 end) end diff --git a/src/atom.luau b/src/atom.luau index 07aeb59..f10dbd6 100644 --- a/src/atom.luau +++ b/src/atom.luau @@ -15,8 +15,10 @@ local function atom(state: T, options: AtomOptions?): Atom local function atom(...) if select("#", ...) == 0 then - for set in next, store.capturing do - set[atom] = true + local index = store.capturing.index + + if index > 0 then + store.capturing.stack[index][atom] = true end return state diff --git a/src/store.luau b/src/store.luau index e197048..68bea45 100644 --- a/src/store.luau +++ b/src/store.luau @@ -4,10 +4,15 @@ type Set = types.Set type WeakMap = types.WeakMap local listeners: WeakMap, WeakMap<() -> (), unknown>> = setmetatable({}, { __mode = "k" }) -local capturing: Set>> = {} + local batched: Set<() -> ()> = {} local batching = false +local capturing = { + stack = {} :: { Set> }, + index = 0, +} + --[=[ Calls the given function and returns the result. If the function yields or throws an error, the thread is closed and an error is thrown. Regardless of @@ -29,7 +34,9 @@ local function try(callback: (U...) -> T | any, finally: (() -> ())?, . if coroutine.status(thread) == "suspended" then local source, line, name = debug.info(callback, "sln") + coroutine.close(thread) + error( "Yielding is not allowed in atom functions. Consider wrapping this code in a Promise or task.defer instead." .. `\nFunction defined at: {source}:{line}` @@ -37,6 +44,7 @@ local function try(callback: (U...) -> T | any, finally: (() -> ())?, . ) elseif not success then local source, line, name = debug.info(callback, "sln") + error( "An error occurred while running an atom function" .. `\nFunction defined at: {source}:{line}` @@ -53,6 +61,7 @@ local function try(callback: (U...) -> T | any, finally: (() -> ())?, . end local success, result = pcall(callback, ...) + finally() assert(success, result) @@ -87,6 +96,34 @@ local function notify(atom: Atom) end end +--[=[ + Returns the result of the function without subscribing to changes. If a + non-function value is provided, it is returned as is. + + @param molecule The atom or molecule to get the state of. + @param args Arguments to pass to the molecule. + @return The current state. +]=] +local function peek(callback: ((U...) -> T) | T, ...: U...): T + if type(callback) ~= "function" then + return callback + end + + if capturing.index == 0 then + return callback(...) + end + + capturing.index += 1 + capturing.stack[capturing.index] = {} + + local result = try(callback, function() + capturing.stack[capturing.index] = nil + capturing.index -= 1 + end, ...) + + return result +end + --[=[ Captures all atoms that are read during the function call and returns them along with the result of the function. Useful for tracking dependencies. @@ -96,14 +133,17 @@ end ]=] local function capture(molecule: () -> T): (Set>, T) if listeners[molecule] then - return { [molecule] = true }, molecule() + return { [molecule] = true }, peek(molecule) end local dependencies: Set> = {} - capturing[dependencies] = true + + capturing.index += 1 + capturing.stack[capturing.index] = dependencies local result = try(molecule, function() - capturing[dependencies] = nil + capturing.stack[capturing.index] = nil + capturing.index -= 1 end) return dependencies, result @@ -133,36 +173,6 @@ local function batch(callback: () -> ()) table.clear(batched) end ---[=[ - Returns the result of the function without subscribing to changes. If a - non-function value is provided, it is returned as is. - - @param molecule The atom or molecule to get the state of. - @param args Arguments to pass to the molecule. - @return The current state. -]=] -local function peek(callback: ((U...) -> T) | T, ...: U...): T - if type(callback) ~= "function" then - return callback - end - - if not next(capturing) then - return callback(...) - end - - local snapshot = table.clone(capturing) - - table.clear(capturing) - - local result = try(callback, function() - for set in next, snapshot do - capturing[set] = true - end - end, ...) - - return result -end - --[=[ Subscribes the listener to the changes of the given atoms.