feat(signals): add watchState function (#4442)

Closes #4416

Author: markostanimirovic

modules/signals/spec/state-source.spec.ts

@@ -1,14 +1,22 @@
-import { effect } from '@angular/core';
+import {
+  createEnvironmentInjector,
+  effect,
+  EnvironmentInjector,
+  Injectable,
+} from '@angular/core';
 import { TestBed } from '@angular/core/testing';
 import {
   getState,
   patchState,
   signalState,
   signalStore,
+  watchState,
+  withHooks,
   withMethods,
   withState,
 } from '../src';
 import { STATE_SOURCE } from '../src/state-source';
+import { createLocalService } from './helpers';
 
 describe('StateSource', () => {
   const initialState = {
@@ -153,4 +161,155 @@ describe('StateSource', () => {
       });
     });
   });
+
+  describe('watchState', () => {
+    describe('with signalState', () => {
+      it('watches state changes', () => {
+        const state = signalState({ count: 0 });
+        const stateHistory: number[] = [];
+
+        TestBed.runInInjectionContext(() => {
+          watchState(state, (state) => stateHistory.push(state.count));
+        });
+
+        patchState(state, { count: 1 });
+        patchState(state, { count: 2 });
+        patchState(state, { count: 3 });
+
+        expect(stateHistory).toEqual([0, 1, 2, 3]);
+      });
+
+      it('stops watching on injector destroy', () => {
+        const stateHistory: number[] = [];
+        const state = signalState({ count: 0 });
+
+        @Injectable()
+        class TestService {
+          constructor() {
+            watchState(state, (state) => stateHistory.push(state.count));
+          }
+        }
+
+        const { destroy } = createLocalService(TestService);
+
+        patchState(state, { count: 1 });
+
+        destroy();
+
+        patchState(state, { count: 2 });
+        patchState(state, { count: 3 });
+
+        expect(stateHistory).toEqual([0, 1]);
+      });
+
+      it('stops watching on manual destroy', () => {
+        const state = signalState({ count: 0 });
+        const stateHistory: number[] = [];
+
+        const { destroy } = TestBed.runInInjectionContext(() =>
+          watchState(state, (state) => stateHistory.push(state.count))
+        );
+
+        patchState(state, { count: 1 });
+        patchState(state, { count: 2 });
+
+        destroy();
+
+        patchState(state, { count: 3 });
+
+        expect(stateHistory).toEqual([0, 1, 2]);
+      });
+
+      it('stops watching on provided injector destroy', () => {
+        const injector1 = createEnvironmentInjector(
+          [],
+          TestBed.inject(EnvironmentInjector)
+        );
+        const injector2 = createEnvironmentInjector(
+          [],
+          TestBed.inject(EnvironmentInjector)
+        );
+        const state = signalState({ count: 0 });
+        const stateHistory1: number[] = [];
+        const stateHistory2: number[] = [];
+
+        watchState(state, (state) => stateHistory1.push(state.count), {
+          injector: injector1,
+        });
+        watchState(state, (state) => stateHistory2.push(state.count), {
+          injector: injector2,
+        });
+
+        patchState(state, { count: 1 });
+        patchState(state, { count: 2 });
+
+        injector1.destroy();
+
+        patchState(state, { count: 3 });
+
+        injector2.destroy();
+
+        patchState(state, { count: 4 });
+
+        expect(stateHistory1).toEqual([0, 1, 2]);
+        expect(stateHistory2).toEqual([0, 1, 2, 3]);
+      });
+
+      it('throws an error when called out of injection context', () => {
+        expect(() => watchState(signalState({}), () => {})).toThrow(
+          /NG0203: watchState\(\) can only be used within an injection context/
+        );
+      });
+    });
+
+    describe('with signalStore', () => {
+      it('watches state changes when used within the store', () => {
+        const stateHistory: number[] = [];
+        const CounterStore = signalStore(
+          withState({ count: 0 }),
+          withHooks({
+            onInit(store) {
+              patchState(store, { count: 1 });
+
+              watchState(store, (state) => stateHistory.push(state.count));
+
+              patchState(store, { count: 2 });
+              patchState(store, { count: 3 });
+            },
+          })
+        );
+
+        TestBed.configureTestingModule({ providers: [CounterStore] });
+        TestBed.inject(CounterStore);
+
+        expect(stateHistory).toEqual([1, 2, 3]);
+      });
+
+      it('watches state changes when used outside of store', () => {
+        const stateHistory: number[] = [];
+        const CounterStore = signalStore(
+          withState({ count: 0 }),
+          withMethods((store) => ({
+            increment(): void {
+              patchState(store, (state) => ({ count: state.count + 1 }));
+            },
+          }))
+        );
+
+        TestBed.configureTestingModule({ providers: [CounterStore] });
+        const store = TestBed.inject(CounterStore);
+        const injector = TestBed.inject(EnvironmentInjector);
+
+        watchState(store, (state) => stateHistory.push(state.count), {
+          injector,
+        });
+
+        store.increment();
+        store.increment();
+        store.increment();
+
+        expect(stateHistory).toEqual([0, 1, 2, 3]);
+      });
+    });
+  });
 });

modules/signals/src/index.ts

@@ -13,6 +13,7 @@ export {
   PartialStateUpdater,
   patchState,
   StateSource,
+  watchState,
   WritableStateSource,
 } from './state-source';
 export { Prettify } from './ts-helpers';

modules/signals/src/state-source.ts

@@ -1,6 +1,17 @@
-import { Signal, WritableSignal } from '@angular/core';
+import {
+  assertInInjectionContext,
+  DestroyRef,
+  inject,
+  Injector,
+  Signal,
+  untracked,
+  WritableSignal,
+} from '@angular/core';
+import { SIGNAL } from '@angular/core/primitives/signals';
 import { Prettify } from './ts-helpers';
 
+const STATE_WATCHERS = new WeakMap<object, Array<StateWatcher<any>>>();
+
 export const STATE_SOURCE = Symbol('STATE_SOURCE');
 
 export type WritableStateSource<State extends object> = {
@@ -15,6 +26,10 @@ export type PartialStateUpdater<State extends object> = (
   state: State
 ) => Partial<State>;
 
+export type StateWatcher<State extends object> = (
+  state: NoInfer<State>
+) => void;
+
 export function patchState<State extends object>(
   stateSource: WritableStateSource<State>,
   ...updaters: Array<
@@ -30,10 +45,72 @@ export function patchState<State extends object>(
       currentState
     )
   );
+
+  notifyWatchers(stateSource);
 }
 
 export function getState<State extends object>(
   stateSource: StateSource<State>
 ): State {
   return stateSource[STATE_SOURCE]();
 }
+
+export function watchState<State extends object>(
+  stateSource: StateSource<State>,
+  watcher: StateWatcher<State>,
+  config?: { injector?: Injector }
+): { destroy(): void } {
+  if (!config?.injector) {
+    assertInInjectionContext(watchState);
+  }
+
+  const injector = config?.injector ?? inject(Injector);
+  const destroyRef = injector.get(DestroyRef);
+
+  addWatcher(stateSource, watcher);
+  watcher(getState(stateSource));
+
+  const destroy = () => removeWatcher(stateSource, watcher);
+  destroyRef.onDestroy(destroy);
+
+  return { destroy };
+}
+
+function getWatchers<State extends object>(
+  stateSource: StateSource<State>
+): Array<StateWatcher<State>> {
+  return STATE_WATCHERS.get(stateSource[STATE_SOURCE][SIGNAL] as object) || [];
+}
+
+function notifyWatchers<State extends object>(
+  stateSource: StateSource<State>
+): void {
+  const watchers = getWatchers(stateSource);
+
+  for (const watcher of watchers) {
+    const state = untracked(() => getState(stateSource));
+    watcher(state);
+  }
+}
+
+function addWatcher<State extends object>(
+  stateSource: StateSource<State>,
+  watcher: StateWatcher<State>
+): void {
+  const watchers = getWatchers(stateSource);
+  STATE_WATCHERS.set(stateSource[STATE_SOURCE][SIGNAL] as object, [
+    ...watchers,
+    watcher,
+  ]);
+}
+
+function removeWatcher<State extends object>(
+  stateSource: StateSource<State>,
+  watcher: StateWatcher<State>
+): void {
+  const watchers = getWatchers(stateSource);
+  STATE_WATCHERS.set(
+    stateSource[STATE_SOURCE][SIGNAL] as object,
+    watchers.filter((w) => w !== watcher)
+  );
+}

projects/ngrx.io/content/guide/signals/signal-store/index.md

@@ -141,30 +141,6 @@ export class BooksComponent {
 
 </code-example>
 
-The `@ngrx/signals` package also offers the `getState` function to get the current state value of the SignalStore.
-When used within the reactive context, state changes are automatically tracked.
-
-<code-example header="books.component.ts">
-
-import { Component, effect, inject } from '@angular/core';
-import { getState } from '@ngrx/signals';
-import { BooksStore } from './books.store';
-
-@Component({ /* ... */ })
-export class BooksComponent {
-  readonly store = inject(BooksStore);
-
-  constructor() {
-    effect(() => {
-      // 👇 The effect will be re-executed whenever the state changes.
-      const state = getState(this.store);
-      console.log('books state changed', state);
-    });
-  }
-}
-
-</code-example>
-
 ## Defining Computed Signals
 
 Computed signals can be added to the store using the `withComputed` feature.