From fc3cd6d9ad998be5f39dd7234b8045a14749943c Mon Sep 17 00:00:00 2001
From: Richard <richerd.rdc@gmail.com>
Date: Sun, 3 Mar 2024 10:24:54 -0800
Subject: [PATCH] fix: add unmount behavior to split atoms

---
 src/__tests__/split.spec.luau |  8 ++++++++
 src/atom.luau                 | 25 ++++++++++---------------
 src/derive.luau               |  2 +-
 src/split.luau                |  6 ++++++
 4 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/src/__tests__/split.spec.luau b/src/__tests__/split.spec.luau
index 829ee5f..1504054 100644
--- a/src/__tests__/split.spec.luau
+++ b/src/__tests__/split.spec.luau
@@ -125,4 +125,12 @@ return function()
 		expect(updated).to.equal(1)
 		expect(deleted).to.equal(3)
 	end)
+
+	it("unmounts from parent atom", function()
+		local itemsAtom = atom({})
+		local itemAtoms = split(itemsAtom)
+		expect((next(itemsAtom.listeners))).to.be.ok()
+		itemAtoms:unmount()
+		expect((next(itemsAtom.listeners))).to.never.be.ok()
+	end)
 end
diff --git a/src/atom.luau b/src/atom.luau
index 2e76e77..211bac2 100644
--- a/src/atom.luau
+++ b/src/atom.luau
@@ -33,16 +33,11 @@ local atomCount = 0
 local function atom<State>(state: State): Atom<State>
 	atomCount += 1
 
-	local listeners: { [(State) -> ()]: true } = {}
-	local label = `atom{atomCount}`
-	local traceback = debug.traceback(nil, 2)
-	local owners: { [any]: true } = {}
-
 	local atom = {
-		listeners = listeners,
-		debugTraceback = traceback,
-		debugLabel = label,
-		owners = owners,
+		listeners = {},
+		debugTraceback = debug.traceback(nil, 2),
+		debugLabel = `atom{atomCount}`,
+		owners = {},
 	} :: Atom<State>
 
 	local equals: ((State, State) -> boolean)?
@@ -68,7 +63,7 @@ local function atom<State>(state: State): Atom<State>
 		if not (if equals then equals(state, nextState) else state == nextState) then
 			state = nextState
 
-			for listener in table.clone(listeners) do
+			for listener in table.clone(self.listeners) do
 				listener(state)
 			end
 		end
@@ -80,7 +75,7 @@ local function atom<State>(state: State): Atom<State>
 		if not (if equals then equals(state, nextState) else state == nextState) then
 			state = nextState
 
-			for listener in table.clone(listeners) do
+			for listener in table.clone(self.listeners) do
 				listener(state)
 			end
 		end
@@ -99,20 +94,20 @@ local function atom<State>(state: State): Atom<State>
 	function atom:readonly(): ReadonlyAtom<State>
 		function setter(_, _)
 			if _G.__DEV__ then
-				error(`Attempt to set read-only atom '{label}' created at: {traceback}`, 3)
+				error(`Attempt to set read-only atom '{self.debugLabel}' created at: {self.debugTraceback}`, 3)
 			end
 		end
 		return self
 	end
 
 	function atom:named(name: string): Atom<State>
-		label = name
+		self.debugLabel = name
 		return self
 	end
 
 	function atom:unmount()
-		table.clear(listeners)
-		table.clear(owners)
+		table.clear(self.listeners)
+		table.clear(self.owners)
 		reference.release(self)
 	end
 
diff --git a/src/derive.luau b/src/derive.luau
index 614d2d1..9d07458 100644
--- a/src/derive.luau
+++ b/src/derive.luau
@@ -29,7 +29,7 @@ local function derive<Result>(...): ReadonlyAtom<Result>
 	local getDerivedAtom = reference.weak(derivedAtom)
 
 	function derivedAtom:unmount()
-		derivedAtomUnmount(self)
+		derivedAtomUnmount(derivedAtom)
 
 		for listener in listeners do
 			for _, atom in atoms do
diff --git a/src/split.luau b/src/split.luau
index 6152c3a..92a4477 100644
--- a/src/split.luau
+++ b/src/split.luau
@@ -13,6 +13,7 @@ local function split<Item, Key>(
 	keyExtractor: ((Item) -> Key)?
 ): ReadonlyAtom<{ [Key]: Atom<Item> }>
 	local splitAtom: ReadonlyAtom<{ [Key]: Atom<Item> }> = atom({}):readonly()
+	local splitAtomUnmount = splitAtom.unmount
 	local getSplitAtom = reference.weak(splitAtom)
 
 	local previousAtoms: { [Key]: Atom<Item> } = {}
@@ -67,6 +68,11 @@ local function split<Item, Key>(
 		splitAtom:setForced(atoms)
 	end
 
+	function splitAtom:unmount()
+		splitAtomUnmount(splitAtom)
+		cleanup()
+	end
+
 	itemsAtom.listeners[listener] = true
 	listener(itemsAtom:get())