From 051443b8155137b7101fe138efa185a200ea4538 Mon Sep 17 00:00:00 2001
From: Josh Morony <joshua.morony@gmail.com>
Date: Tue, 14 Nov 2023 10:16:04 +1030
Subject: [PATCH] feat(signal-slice): allow supplying source as a function that
 accepts state signal (#146)

* feat(signal-slice): test supplying state to source

* feat(signal-slice): allow supplying function that accepts state signal as a source

* docs(signal-slice): add docs for supplying source as a function
---
 .../docs/utilities/Signals/signal-slice.md    | 19 ++++++++++++++++
 .../signal-slice/src/signal-slice.spec.ts     | 22 +++++++++++++++++++
 .../signal-slice/src/signal-slice.ts          | 19 +++++++++++-----
 3 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/docs/src/content/docs/utilities/Signals/signal-slice.md b/docs/src/content/docs/utilities/Signals/signal-slice.md
index bbad4ff6..be87f706 100644
--- a/docs/src/content/docs/utilities/Signals/signal-slice.md
+++ b/docs/src/content/docs/utilities/Signals/signal-slice.md
@@ -56,6 +56,25 @@ state = signalSlice({
 
 The `source` should be mapped to a partial of the `initialState`. In the example above, when the source emits it will update both the `checklists` and the `loaded` properties in the state signal.
 
+If you need to utilise the current state in a source, instead of supplying the
+observable directly as a source you can supply a function that accepts the state
+signal and returns the source:
+
+```ts
+state = signalSlice({
+	initialState: this.initialState,
+	sources: [
+		this.loadChecklists$,
+		(state) =>
+			this.newMessage$.pipe(
+				map((newMessage) => ({
+					messages: [...state().messages, newMessage],
+				}))
+			),
+	],
+});
+```
+
 ## Reducers and Actions
 
 Another way to update the state is through `reducers` and `actions`. This is good for situations where you need to manually/imperatively trigger some action, and then use the current state in some way in order to calculate the new state.
diff --git a/libs/ngxtension/signal-slice/src/signal-slice.spec.ts b/libs/ngxtension/signal-slice/src/signal-slice.spec.ts
index bdf83d18..c6243ba1 100644
--- a/libs/ngxtension/signal-slice/src/signal-slice.spec.ts
+++ b/libs/ngxtension/signal-slice/src/signal-slice.spec.ts
@@ -74,6 +74,28 @@ describe(signalSlice.name, () => {
 			expect(state().user.firstName).toEqual(testUpdate.user.firstName);
 			expect(state().age).toEqual(testUpdate2.age);
 		});
+
+		it('should allow supplying function that takes state signal', () => {
+			const ageSource$ = new Subject<number>();
+
+			TestBed.runInInjectionContext(() => {
+				state = signalSlice({
+					initialState,
+					sources: [
+						testSource$,
+						(state) =>
+							ageSource$.pipe(
+								map((incrementAge) => ({ age: state().age + incrementAge }))
+							),
+					],
+				});
+			});
+
+			const incrementAge = 5;
+			ageSource$.next(incrementAge);
+
+			expect(state().age).toEqual(initialState.age + incrementAge);
+		});
 	});
 
 	describe('reducers', () => {
diff --git a/libs/ngxtension/signal-slice/src/signal-slice.ts b/libs/ngxtension/signal-slice/src/signal-slice.ts
index f322cca1..b5207000 100644
--- a/libs/ngxtension/signal-slice/src/signal-slice.ts
+++ b/libs/ngxtension/signal-slice/src/signal-slice.ts
@@ -111,6 +111,8 @@ type ActionStreams<
 		: never;
 };
 
+export type Source<TSignalValue> = Observable<PartialOrValue<TSignalValue>>;
+
 export type SignalSlice<
 	TSignalValue,
 	TReducers extends NamedReducers<TSignalValue>,
@@ -132,7 +134,10 @@ export function signalSlice<
 	TEffects extends NamedEffects
 >(config: {
 	initialState: TSignalValue;
-	sources?: Array<Observable<PartialOrValue<TSignalValue>>>;
+	sources?: Array<
+		| Source<TSignalValue>
+		| ((state: Signal<TSignalValue>) => Source<TSignalValue>)
+	>;
 	reducers?: TReducers;
 	asyncReducers?: TAsyncReducers;
 	selectors?: (state: Signal<TSignalValue>) => TSelectors;
@@ -159,13 +164,17 @@ export function signalSlice<
 
 	const state = signal(initialState);
 
-	for (const source of sources) {
-		connect(state, source);
-	}
-
 	const readonlyState = state.asReadonly();
 	const subs: Subject<unknown>[] = [];
 
+	for (const source of sources) {
+		if (isObservable(source)) {
+			connect(state, source);
+		} else {
+			connect(state, source(readonlyState));
+		}
+	}
+
 	for (const [key, reducer] of Object.entries(reducers as TReducers)) {
 		const subject = new Subject();