From aa52c98cd97d819e93293e55d3b2eb84e1aeaba2 Mon Sep 17 00:00:00 2001 From: Richard Date: Sun, 9 Jun 2024 20:10:32 -0700 Subject: [PATCH] feat: allow conditional dependencies --- src/__tests__/benchmarks/atom.bench.luau | 19 +-- src/__tests__/capture.spec.luau | 21 +++- src/__tests__/computed.spec.luau | 31 +++++ src/__tests__/effect.spec.luau | 29 +++++ src/__tests__/peek.spec.luau | 8 +- src/__tests__/subscribe.spec.luau | 28 +++++ src/atom.luau | 4 +- src/computed.luau | 9 +- src/effect.luau | 12 +- src/index.d.ts | 2 +- src/store.luau | 144 +++++++++++++++++------ src/subscribe.luau | 12 +- 12 files changed, 235 insertions(+), 84 deletions(-) diff --git a/src/__tests__/benchmarks/atom.bench.luau b/src/__tests__/benchmarks/atom.bench.luau index e7819f9..6663caa 100644 --- a/src/__tests__/benchmarks/atom.bench.luau +++ b/src/__tests__/benchmarks/atom.bench.luau @@ -7,34 +7,23 @@ local atom = require(Package.atom) local computed = require(Package.computed) local function parameterGenerator() - return atom({ count = 0 }) + return atom(0) end local benchmarks = { atom = function(_, source: any) for _ = 1, 500 do - source({ count = math.random(1, 2) }) + source(math.random(1, 2)) end end, computed = function(_, source: any) computed(function() - return { count = 2 * source().count } + return 2 * source() end) - for _ = 1, 500 do - source({ count = math.random(1, 2) }) - end - end, - computed_nested = function(_, source: any) - local double = computed(function() - return { count = 2 * source().count } - end) - computed(function() - return { count = 2 * double().count } - end) for _ = 1, 500 do - source({ count = math.random(1, 2) }) + source(math.random(1, 2)) end end, } diff --git a/src/__tests__/capture.spec.luau b/src/__tests__/capture.spec.luau index 59caaa4..37b37ff 100644 --- a/src/__tests__/capture.spec.luau +++ b/src/__tests__/capture.spec.luau @@ -2,21 +2,30 @@ local atom = require(script.Parent.Parent.atom) local store = require(script.Parent.Parent.store) return function() + it("gracefully handles yielding", function() + expect(function() + store.capture(function() + task.wait(10) + return 1 + end) + end).to.throw() + end) + it("captures atoms from a function", function() local a, b, c = atom(1), atom(1), atom(2) - local captured, result = store.capture(function() + local deps, result = store.capture(function() return a() + b() end) - expect(captured[a]).to.be.ok() - expect(captured[b]).to.be.ok() - expect(captured[c]).to.never.be.ok() + expect(deps[a]).to.be.ok() + expect(deps[b]).to.be.ok() + expect(deps[c]).to.never.be.ok() expect(result).to.equal(2) end) it("captures an atom by itself", function() local a = atom(1) - local captured, result = store.capture(a) - expect(captured[a]).to.be.ok() + local deps, result = store.capture(a) + expect(deps[a]).to.be.ok() expect(result).to.equal(1) end) end diff --git a/src/__tests__/computed.spec.luau b/src/__tests__/computed.spec.luau index b9b3a19..87d0d6b 100644 --- a/src/__tests__/computed.spec.luau +++ b/src/__tests__/computed.spec.luau @@ -68,4 +68,35 @@ return function() assert(count(store.listeners[source]) == 0, "nested computed atom did not gc") end) + + it("tracks conditional dependencies", function() + local condition = atom(false) + local a = atom(0) + local b = atom(1) + local reruns = 0 + local current + + -- assign variable to avoid gc + local _ = computed(function() + reruns += 1 + current = if condition() then a() else b() + return current + end) + + expect(reruns).to.equal(1) + expect(current).to.equal(1) + + condition(true) + expect(reruns).to.equal(2) + expect(current).to.equal(0) + + a(1) + b(2) + expect(reruns).to.equal(3) + expect(current).to.equal(1) + + condition(false) + expect(reruns).to.equal(4) + expect(current).to.equal(2) + end) end diff --git a/src/__tests__/effect.spec.luau b/src/__tests__/effect.spec.luau index a6bb66f..fd077c9 100644 --- a/src/__tests__/effect.spec.luau +++ b/src/__tests__/effect.spec.luau @@ -43,4 +43,33 @@ return function() unsubscribe() expect(cleanups).to.equal(2) end) + + it("tracks conditional dependencies", function() + local condition = atom(false) + local a = atom(0) + local b = atom(1) + local reruns = 0 + local current + + effect(function() + reruns += 1 + current = if condition() then a() else b() + end) + + expect(reruns).to.equal(1) + expect(current).to.equal(1) + + condition(true) + expect(reruns).to.equal(2) + expect(current).to.equal(0) + + a(1) + b(2) + expect(reruns).to.equal(3) + expect(current).to.equal(1) + + condition(false) + expect(reruns).to.equal(4) + expect(current).to.equal(2) + end) end diff --git a/src/__tests__/peek.spec.luau b/src/__tests__/peek.spec.luau index bde5def..a861517 100644 --- a/src/__tests__/peek.spec.luau +++ b/src/__tests__/peek.spec.luau @@ -4,12 +4,12 @@ local store = require(script.Parent.Parent.store) return function() it("calls an atom without capturing it", function() local a, b, c = atom(1), atom(1), atom(1) - local captured, result = store.capture(function() + local deps, result = store.capture(function() return a() + store.peek(b) + c() end) - expect(captured[a]).to.be.ok() - expect(captured[b]).to.never.be.ok() - expect(captured[c]).to.be.ok() + expect(deps[a]).to.be.ok() + expect(deps[b]).to.never.be.ok() + expect(deps[c]).to.be.ok() expect(result).to.equal(3) end) diff --git a/src/__tests__/subscribe.spec.luau b/src/__tests__/subscribe.spec.luau index 5097213..d5256e5 100644 --- a/src/__tests__/subscribe.spec.luau +++ b/src/__tests__/subscribe.spec.luau @@ -67,4 +67,32 @@ return function() assert(count(store.listeners[source]) == 0, "computed atom did not collect") end) + + it("tracks conditional dependencies", function() + local condition = atom(false) + local a = atom(0) + local b = atom(1) + local reruns = 0 + local current + + subscribe(function() + return if condition() then a() else b() + end, function(state) + reruns += 1 + current = state + end) + + condition(true) + expect(reruns).to.equal(1) + expect(current).to.equal(0) + + a(1) + b(2) + expect(reruns).to.equal(2) + expect(current).to.equal(1) + + condition(false) + expect(reruns).to.equal(3) + expect(current).to.equal(2) + end) end diff --git a/src/atom.luau b/src/atom.luau index 52abf93..b17b9b0 100644 --- a/src/atom.luau +++ b/src/atom.luau @@ -15,8 +15,8 @@ local function atom(state: T, options: AtomOptions?): Atom local function atom(...) if select("#", ...) == 0 then - if store.isCapturing() then - store.captured[atom] = true + if store.capturing.current then + store.capturing.current[atom] = true end return state diff --git a/src/computed.luau b/src/computed.luau index 0b0a910..cdfa072 100644 --- a/src/computed.luau +++ b/src/computed.luau @@ -15,19 +15,18 @@ type Molecule = types.Molecule @return A new read-only atom. ]=] local function computed(molecule: Molecule, options: AtomOptions?): Molecule - local captured, state = store.capture(molecule) + local dependencies, state = store.capture(molecule) local computedAtom = atom(state, options) local computedRef = setmetatable({ current = computedAtom }, { __mode = "v" }) local function listener() if computedRef.current then - computedRef.current(molecule()) + dependencies, state = store.reconnect(molecule, dependencies, listener, computedRef.current) + computedRef.current(state) end end - for atom in captured do - store.listeners[atom][listener] = computedAtom - end + store.connect(dependencies, listener, computedAtom) return computedAtom end diff --git a/src/effect.luau b/src/effect.luau index 4e27531..8f5450e 100644 --- a/src/effect.luau +++ b/src/effect.luau @@ -13,24 +13,20 @@ type Cleanup = () -> () @return A function that unsubscribes the callback. ]=] local function effect(callback: () -> any): () -> () - local captured, cleanup = store.capture(callback) + local dependencies, cleanup = store.capture(callback) local function listener() if cleanup then cleanup() end - cleanup = callback() + dependencies, cleanup = store.reconnect(callback, dependencies, listener) end - for atom in captured do - store.listeners[atom][listener] = true - end + store.connect(dependencies, listener) return function() - for atom in captured do - store.listeners[atom][listener] = nil - end + store.disconnect(dependencies, listener) if cleanup then cleanup() diff --git a/src/index.d.ts b/src/index.d.ts index edcbd90..8c4f0e2 100644 --- a/src/index.d.ts +++ b/src/index.d.ts @@ -123,7 +123,7 @@ declare namespace Charm { * @param molecule The function to run. * @returns A tuple containing the captured atoms and the result of the function. */ - function capture(molecule: Molecule): LuaTuple<[captured: Set>, state: State]>; + function capture(molecule: Molecule): LuaTuple<[dependencies: Set>, state: State]>; /** * Notifies all subscribers of the given atom that the state has changed. diff --git a/src/store.luau b/src/store.luau index 1847611..7cfe094 100644 --- a/src/store.luau +++ b/src/store.luau @@ -4,19 +4,55 @@ type WeakMap = types.WeakMap local listeners: WeakMap, WeakMap<() -> (), unknown>> = setmetatable({}, { __mode = "k" }) -local captured: { [Atom]: true } = {} -local capturing = false - +local capturing: { current: { [Atom]: true }? } = {} local batched: { [() -> ()]: true } = {} local batching = false +--[=[ + 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 + the outcome, the `finally` function is called to clean up any resources. + + @param callback The function to run. + @param finally Cleanup logic to run before error handling. + @param ... Arguments to pass to the callback. + @return The result of the callback. +]=] +local function try(callback: (U...) -> T | any, finally: (() -> ())?, ...: U...): T + if _G.__DEV__ then + local thread = coroutine.create(callback) + local success, result = coroutine.resume(thread, ...) + + if finally then + finally() + end + + if coroutine.status(thread) == "suspended" then + coroutine.close(thread) + error("Yielding is not allowed in atom functions") + elseif not success then + error("An error occurred while running an atom function") + end + + return result + end + + local result = callback(...) + + if finally then + finally() + end + + return result +end + --[=[ Returns whether the given value is an atom. @param value The value to check. @return `true` if the value is an atom, otherwise `false`. ]=] -local function isAtom(value: any) +local function isAtom(value: any): boolean return not not (value and listeners[value]) end @@ -34,7 +70,7 @@ local function notify(atom: Atom) end for listener in next, table.clone(listeners[atom]) do - listener() + try(listener) end end @@ -45,17 +81,19 @@ end @param molecule The function to run. @return A tuple containing the captured atoms and the result of the function. ]=] -local function capture(callback: () -> T): ({ [Atom]: true }, T) - if listeners[callback] then - return { [callback] = true }, callback() +local function capture(molecule: () -> T): ({ [Atom]: true }, T) + if listeners[molecule] then + return { [molecule] = true }, molecule() end - capturing = true - local result = callback() - capturing = false + local dependencies = {} + local previous = capturing.current - local dependencies = table.clone(captured) - table.clear(captured) + capturing.current = dependencies + + local result = try(molecule, function() + capturing.current = previous + end) return dependencies, result end @@ -68,16 +106,17 @@ end ]=] local function batch(callback: () -> ()) if batching then - callback() - return + return callback() end batching = true - callback() - batching = false + + try(callback, function() + batching = false + end) for listener in batched do - listener() + try(listener) end table.clear(batched) @@ -96,44 +135,77 @@ local function peek(callback: ((U...) -> T) | T, ...: U...): T return callback end - if not capturing then + if not capturing.current then return callback(...) end - capturing = false - local result = callback(...) - capturing = true + local previous = capturing.current + + capturing.current = nil + + local result = try(callback, function() + capturing.current = previous + end, ...) return result end --[=[ - Sets whether the store is capturing atoms. If `true`, atoms that are read - will be tracked and returned by `capture`. - - @param value Whether to capture atoms. + Subscribes the listener to the changes of the given atoms. + + @param dependencies The atoms to listen to. + @param listener The function to call when the atoms change. + @param ref Optionally bind the lifetime of the listener to a value. ]=] -local function setCapturing(value: boolean) - capturing = value +local function connect(dependencies: { [Atom]: true }, listener: () -> (), ref: unknown?) + for atom in dependencies do + listeners[atom][listener] = ref or true + end end --[=[ - Returns whether the store is currently capturing atoms. - - @return `true` if the store is capturing atoms, otherwise `false`. + Unsubscribes the listener from the changes of the given atoms. + + @param dependencies The atoms to stop listening to. + @param listener The function to stop calling when the atoms change. ]=] -local function isCapturing() - return capturing +local function disconnect(dependencies: { [Atom]: true }, listener: () -> ()) + for atom in dependencies do + listeners[atom][listener] = nil + end +end + +--[=[ + Disconnects the listener and evaluates the molecule again to capture any + new dependencies. Useful for tracking conditional dependencies. + + @param molecule The atom or molecule to evaluate. + @param dependencies The current dependencies. + @param listener The listener to disconnect. + @param ref Optionally bind the lifetime of the listener to a value. + @return The new dependencies and the state of the molecule. +]=] +local function reconnect( + molecule: Atom, + dependencies: { [Atom]: true }, + listener: () -> (), + ref: unknown? +): ({ [Atom]: true }, T) + disconnect(dependencies, listener) + local nextDependencies, state = capture(molecule) + connect(nextDependencies, listener, ref) + return nextDependencies, state end return { listeners = listeners, - captured = captured, + capturing = capturing, isAtom = isAtom, notify = notify, capture = capture, batch = batch, peek = peek, - setCapturing = setCapturing, - isCapturing = isCapturing, + connect = connect, + disconnect = disconnect, + reconnect = reconnect, } diff --git a/src/subscribe.luau b/src/subscribe.luau index f285c90..1083d99 100644 --- a/src/subscribe.luau +++ b/src/subscribe.luau @@ -16,20 +16,18 @@ local function subscribe(molecule: Molecule, callback: (state: T, prev: T) local function listener() local prevState = state - state = molecule() + + captured, state = store.reconnect(molecule, captured, listener) + if state ~= prevState then callback(state, prevState) end end - for atom in captured do - store.listeners[atom][listener] = true - end + store.connect(captured, listener) return function() - for atom in captured do - store.listeners[atom][listener] = nil - end + store.disconnect(captured, listener) end end