feat(signals): enhance withComputed to accept computation functions (#4822)

Closes #4782

Co-authored-by: Marko Stanimirović <markostanimirovic95@gmail.com>
Co-authored-by: Tim Deschryver <28659384+timdeschryver@users.noreply.github.com>
Co-authored-by: michael-small <33669563+michael-small@users.noreply.github.com>

Author: 4 people

modules/signals/spec/types/with-computed.types.spec.ts

@@ -0,0 +1,90 @@
+import { expecter } from 'ts-snippet';
+import { compilerOptions } from './helpers';
+import { signalStore, withComputed } from 'modules/signals/src';
+import { TestBed } from '@angular/core/testing';
+
+describe('withComputed', () => {
+  const expectSnippet = expecter(
+    (code) => `
+        import {
+          deepComputed,
+          signalStore,
+          withComputed,
+        } from '@ngrx/signals';
+        import { TestBed } from '@angular/core/testing';
+        import { signal } from '@angular/core';
+        
+        ${code}
+      `,
+    compilerOptions()
+  );
+
+  it('creates a Signal automatically', () => {
+    const snippet = `
+      const Store = signalStore(
+        withComputed(() => ({
+          user: () => ({ firstName: 'John', lastName: 'Doe' })
+        }))
+      );
+
+      const store = TestBed.inject(Store);
+      const user = store.user;
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer(
+      'user',
+      'Signal<{ firstName: string; lastName: string; }>'
+    );
+  });
+
+  it('keeps a WritableSignal intact, if passed', () => {
+    const snippet = `
+      const user = signal({ firstName: 'John', lastName: 'Doe' });
+
+      const Store = signalStore(
+        withComputed(() => ({
+          user,
+        }))
+      );
+
+      const store = TestBed.inject(Store);
+      const userSignal = store.user;
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer(
+      'userSignal',
+      'WritableSignal<{ firstName: string; lastName: string; }>'
+    );
+  });
+
+  it('keeps a DeepSignal intact, if passed', () => {
+    const snippet = `
+    const user = deepComputed(
+      signal({
+        name: 'John Doe',
+        address: {
+          street: '123 Main St',
+          city: 'Anytown',
+        },
+      })
+    );
+
+    const Store = signalStore(
+      withComputed(() => ({
+        user,
+      }))
+    );
+
+    const store = TestBed.inject(Store);
+    const userSignal = store.user;
+  `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer(
+      'userSignal',
+      'DeepSignal<{ name: string; address: { street: string; city: string; }; }>'
+    );
+  });
+});

modules/signals/spec/with-computed.spec.ts

@@ -1,6 +1,13 @@
-import { signal } from '@angular/core';
-import { withComputed, withMethods, withState } from '../src';
-import { getInitialInnerStore } from '../src/signal-store';
+import { computed, signal } from '@angular/core';
+import {
+  deepComputed,
+  signalStoreFeature,
+  withComputed,
+  withMethods,
+  withState,
+} from '../src';
+import { getInitialInnerStore, signalStore } from '../src/signal-store';
+import { TestBed } from '@angular/core/testing';
 
 describe('withComputed', () => {
   it('adds computed signals to the store immutably', () => {
@@ -54,4 +61,87 @@ describe('withComputed', () => {
       'p1, s2, m1, Symbol(computed_secret)'
     );
   });
+
+  it('adds computed automatically if the value is a function', () => {
+    const initialStore = getInitialInnerStore();
+
+    const store = signalStoreFeature(
+      withState({ a: 2, b: 3 }),
+      withComputed(({ a, b }) => ({
+        sum: () => a() + b(),
+        product: () => a() * b(),
+      }))
+    )(initialStore);
+
+    expect(store.props.sum()).toBe(5);
+    expect(store.props.product()).toBe(6);
+  });
+
+  it('allows to mix user-provided computeds and automatically computed ones', () => {
+    const initialStore = getInitialInnerStore();
+
+    const store = signalStoreFeature(
+      withState({ a: 2, b: 3 }),
+      withComputed(({ a, b }) => ({
+        sum: () => a() + b(),
+        product: computed(() => a() * b()),
+      }))
+    )(initialStore);
+
+    expect(store.props.sum()).toBe(5);
+    expect(store.props.product()).toBe(6);
+  });
+
+  it('does not change a WritableSignal', () => {
+    const user = signal({ firstName: 'John', lastName: 'Doe' });
+
+    const Store = signalStore(
+      { providedIn: 'root' },
+      withComputed(() => ({
+        user,
+      }))
+    );
+
+    const store = TestBed.inject(Store);
+
+    expect(store.user).toBe(user);
+  });
+
+  it('does not change a DeepSignal', () => {
+    const user = deepComputed(
+      signal({
+        name: 'John Doe',
+        address: {
+          street: '123 Main St',
+          city: 'Anytown',
+        },
+      })
+    );
+
+    const Store = signalStore(
+      { providedIn: 'root' },
+      withComputed(() => ({
+        user,
+      }))
+    );
+
+    const store = TestBed.inject(Store);
+
+    expect(store.user).toBe(user);
+  });
+
+  it('does not change a Signal', () => {
+    const user = computed(() => ({ firstName: 'John', lastName: 'Doe' }));
+
+    const Store = signalStore(
+      { providedIn: 'root' },
+      withComputed(() => ({
+        user,
+      }))
+    );
+
+    const store = TestBed.inject(Store);
+
+    expect(store.user).toBe(user);
+  });
 });

modules/signals/src/with-computed.ts

@@ -1,22 +1,51 @@
+import { computed, isSignal, Signal } from '@angular/core';
 import {
-  SignalsDictionary,
   SignalStoreFeature,
   SignalStoreFeatureResult,
   StateSignals,
 } from './signal-store-models';
 import { Prettify } from './ts-helpers';
 import { withProps } from './with-props';
 
+type ComputedResult<
+  ComputedDictionary extends Record<
+    string | symbol,
+    Signal<unknown> | (() => unknown)
+  >
+> = {
+  [P in keyof ComputedDictionary]: ComputedDictionary[P] extends Signal<unknown>
+    ? ComputedDictionary[P]
+    : ComputedDictionary[P] extends () => infer V
+    ? Signal<V>
+    : never;
+};
+
 export function withComputed<
   Input extends SignalStoreFeatureResult,
-  ComputedSignals extends SignalsDictionary
+  ComputedDictionary extends Record<
+    string | symbol,
+    Signal<unknown> | (() => unknown)
+  >
 >(
-  signalsFactory: (
+  computedFactory: (
     store: Prettify<StateSignals<Input['state']> & Input['props']>
-  ) => ComputedSignals
+  ) => ComputedDictionary
 ): SignalStoreFeature<
   Input,
-  { state: {}; props: ComputedSignals; methods: {} }
+  { state: {}; props: ComputedResult<ComputedDictionary>; methods: {} }
 > {
-  return withProps(signalsFactory);
+  return withProps((store) => {
+    const computedResult = computedFactory(store);
+    const computedResultKeys = Reflect.ownKeys(computedResult);
+
+    return computedResultKeys.reduce((prev, key) => {
+      const signalOrComputation = computedResult[key];
+      return {
+        ...prev,
+        [key]: isSignal(signalOrComputation)
+          ? signalOrComputation
+          : computed(signalOrComputation),
+      };
+    }, {} as ComputedResult<ComputedDictionary>);
+  });
 }

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

@@ -145,7 +145,7 @@ export class BookSearch {
 
 Computed signals can be added to the store using the `withComputed` feature.
 This feature accepts a factory function as an input argument, which is executed within the injection context.
-The factory should return a dictionary of computed signals, utilizing previously defined state signals and properties that are accessible through its input argument.
+The factory should return a dictionary containing either computed signals or functions that return values (which are automatically wrapped in computed signals), utilizing previously defined state signals and properties that are accessible through its input argument.
 
 <code-example header="book-search-store.ts">
 
@@ -162,13 +162,14 @@ export const BookSearchStore = signalStore(
   // 👇 Accessing previously defined state signals and properties.
   withComputed(({ books, filter }) => ({
     booksCount: computed(() => books().length),
-    sortedBooks: computed(() => {
+    // 👇 Adds computed automatically
+    sortedBooks: () => {
       const direction = filter.order() === 'asc' ? 1 : -1;
 
       return books().toSorted((a, b) =>
         direction * a.title.localeCompare(b.title)
       );
-    }),
+    },
   }))
 );