From a8cb227036e8f640cf2405a8f8ddb9fdfbb05fb3 Mon Sep 17 00:00:00 2001
From: Rainer Hahnekamp <rainer.hahnekamp@gmail.com>
Date: Tue, 8 Aug 2023 11:02:47 +0200
Subject: [PATCH 1/5] freeze state in signalState

---
 src/lib/signal-state.spec.ts   |  92 --------------------
 src/lib/signal-state.ts        |  12 +--
 src/lib/utils.ts               |  39 +++++++++
 src/tests/signal-state.spec.ts | 152 +++++++++++++++++++++++++++++++++
 src/tests/test2.ts             |   1 -
 src/tests/with-effect.ts       |  26 ++++++
 6 files changed, 224 insertions(+), 98 deletions(-)
 delete mode 100644 src/lib/signal-state.spec.ts
 create mode 100644 src/lib/utils.ts
 create mode 100644 src/tests/signal-state.spec.ts
 create mode 100644 src/tests/with-effect.ts

diff --git a/src/lib/signal-state.spec.ts b/src/lib/signal-state.spec.ts
deleted file mode 100644
index 15c7730..0000000
--- a/src/lib/signal-state.spec.ts
+++ /dev/null
@@ -1,92 +0,0 @@
-import { signalState } from './signal-state';
-import { ApplicationRef, Component, effect } from '@angular/core';
-import { fakeAsync, TestBed, waitForAsync } from '@angular/core/testing';
-import { state } from '@angular/animations';
-
-describe('Signal State', () => {
-  const initialState = {
-    user: {
-      firstName: 'John',
-      lastName: 'Smith',
-    },
-    foo: 'bar',
-    numbers: [1, 2, 3],
-  };
-
-  const setup = () => signalState(initialState);
-
-  it('should support nested signals', () => {
-    const state = setup();
-
-    expect(state()).toBe(initialState);
-    expect(state.user()).toBe(initialState.user);
-    expect(state.user.firstName()).toBe(initialState.user.firstName);
-  });
-
-  it('should allow updates', () => {
-    const state = setup();
-    state.$update((state) => ({
-      ...state,
-      user: { firstName: 'Johannes', lastName: 'Schmidt' },
-    }));
-    expect(state()).toEqual({
-      ...initialState,
-      user: { firstName: 'Johannes', lastName: 'Schmidt' },
-    });
-  });
-
-  it('should update immutably', () => {
-    const state = setup();
-    state.$update((state) => ({
-      ...state,
-      foo: 'bar',
-      numbers: [3, 2, 1],
-    }));
-    expect(state.user()).toBe(initialState.user);
-    expect(state.foo()).toBe(initialState.foo);
-    expect(state.numbers()).not.toBe(initialState.numbers);
-  });
-
-  // TODO: Find a better way to execute effects
-  describe('Effects in Signals', () => {
-    @Component({
-      template: '',
-    })
-    class TestComponent {
-      userFired = 0;
-      numbersFired = 0;
-      state = setup();
-
-      constructor() {
-        effect(() => {
-          this.userFired++;
-          this.state.user();
-        });
-
-        const numbersListener = effect(() => {
-          this.numbersFired++;
-          this.state.numbers();
-        });
-      }
-    }
-
-    it('should not fire all signals on update', fakeAsync(() => {
-      const fixture = TestBed.createComponent(TestComponent);
-      const component = fixture.componentInstance;
-
-      expect(component.numbersFired).toBe(0);
-      expect(component.userFired).toBe(0);
-
-      fixture.detectChanges();
-
-      expect(component.numbersFired).toBe(1);
-      expect(component.userFired).toBe(1);
-
-      component.state.$update((state) => ({ ...state, numbers: [4, 5, 6] }));
-      fixture.detectChanges();
-
-      expect(component.numbersFired).toBe(2);
-      expect(component.userFired).toBe(1);
-    }));
-  });
-});
diff --git a/src/lib/signal-state.ts b/src/lib/signal-state.ts
index 0e6061d..5869399 100644
--- a/src/lib/signal-state.ts
+++ b/src/lib/signal-state.ts
@@ -1,6 +1,7 @@
 import { signal, WritableSignal } from '@angular/core';
 import { DeepSignal, toDeepSignal } from './deep-signal';
 import { defaultEqualityFn } from './select-signal';
+import { deepFreeze, isFunction, isObjectLike } from './utils';
 
 export type SignalState<State extends Record<string, unknown>> =
   DeepSignal<State> & SignalStateUpdate<State>;
@@ -16,6 +17,7 @@ export type SignalStateUpdate<State extends Record<string, unknown>> = {
 export function signalState<State extends Record<string, unknown>>(
   initialState: State
 ): SignalState<State> {
+  deepFreeze(initialState);
   const stateSignal = signal(initialState, { equal: defaultEqualityFn });
   const deepSignal = toDeepSignal(stateSignal.asReadonly());
   (deepSignal as SignalState<State>).$update =
@@ -29,12 +31,12 @@ export function signalStateUpdateFactory<State extends Record<string, unknown>>(
 ): SignalStateUpdate<State>['$update'] {
   return (...updaters) =>
     stateSignal.update((state) =>
-      updaters.reduce(
-        (currentState: State, updater) => ({
+      updaters.reduce((currentState: State, updater) => {
+        deepFreeze(currentState);
+        return {
           ...currentState,
           ...(typeof updater === 'function' ? updater(currentState) : updater),
-        }),
-        state
-      )
+        };
+      }, state)
     );
 }
diff --git a/src/lib/utils.ts b/src/lib/utils.ts
new file mode 100644
index 0000000..2da10e5
--- /dev/null
+++ b/src/lib/utils.ts
@@ -0,0 +1,39 @@
+export function deepFreeze(target: any) {
+  Object.freeze(target);
+
+  const targetIsFunction = isFunction(target);
+
+  for (const prop of Object.getOwnPropertyNames(target)) {
+    // Do not freeze Ivy properties (e.g. router state)
+    // https://github.com/ngrx/platform/issues/2109#issuecomment-582689060
+    if (prop.startsWith('ɵ')) {
+      continue;
+    }
+
+    if (
+      Object.hasOwn(target, prop) &&
+      (targetIsFunction
+        ? prop !== 'caller' && prop !== 'callee' && prop !== 'arguments'
+        : true)
+    ) {
+      const propValue: unknown = target[prop];
+
+      if (
+        (isObjectLike(propValue) || isFunction(propValue)) &&
+        !Object.isFrozen(propValue)
+      ) {
+        deepFreeze(propValue);
+      }
+    }
+  }
+
+  return target;
+}
+
+export function isFunction(target: unknown): target is () => unknown {
+  return typeof target === 'function';
+}
+
+export function isObjectLike(target: unknown): target is object {
+  return typeof target === 'object' && target !== null;
+}
diff --git a/src/tests/signal-state.spec.ts b/src/tests/signal-state.spec.ts
new file mode 100644
index 0000000..dd3db5f
--- /dev/null
+++ b/src/tests/signal-state.spec.ts
@@ -0,0 +1,152 @@
+import { effect } from '@angular/core';
+import { withEffect } from './with-effect';
+import { signalState } from '../lib/signal-state';
+import { selectSignal } from '../lib/select-signal';
+
+describe('Signal State', () => {
+  const initialState = {
+    user: {
+      firstName: 'John',
+      lastName: 'Smith',
+    },
+    foo: 'bar',
+    numbers: [1, 2, 3],
+  };
+
+  const setup = () => signalState(initialState);
+
+  it('should support nested signals', () => {
+    const state = setup();
+
+    expect(state()).toBe(initialState);
+    expect(state.user()).toBe(initialState.user);
+    expect(state.user.firstName()).toBe(initialState.user.firstName);
+  });
+
+  it('should allow updates', () => {
+    const state = setup();
+    state.$update((state) => ({
+      ...state,
+      user: { firstName: 'Johannes', lastName: 'Schmidt' },
+    }));
+    expect(state()).toEqual({
+      ...initialState,
+      user: { firstName: 'Johannes', lastName: 'Schmidt' },
+    });
+  });
+
+  it('should update immutably', () => {
+    const state = setup();
+    state.$update((state) => ({
+      ...state,
+      foo: 'bar',
+      numbers: [3, 2, 1],
+    }));
+    expect(state.user()).toBe(initialState.user);
+    expect(state.foo()).toBe(initialState.foo);
+    expect(state.numbers()).not.toBe(initialState.numbers);
+  });
+
+  describe('equal checks', () => {
+    it(
+      'should not fire unchanged signals on update',
+      withEffect((detectChanges) => {
+        const state = setup();
+
+        const numberEffect = jest.fn(() => state.numbers());
+        effect(numberEffect);
+
+        const userEffect = jest.fn(() => state.user());
+        effect(userEffect);
+
+        expect(numberEffect).toHaveBeenCalledTimes(0);
+        expect(userEffect).toHaveBeenCalledTimes(0);
+
+        // run effects for the first time
+        detectChanges();
+        expect(numberEffect).toHaveBeenCalledTimes(1);
+        expect(userEffect).toHaveBeenCalledTimes(1);
+
+        // update state with effect run
+        state.$update((state) => ({ ...state, numbers: [4, 5, 6] }));
+        detectChanges();
+        expect(numberEffect).toHaveBeenCalledTimes(2);
+        expect(userEffect).toHaveBeenCalledTimes(1);
+      })
+    );
+
+    it(
+      'should not fire for unchanged derived signals',
+      withEffect((detectChanges) => {
+        const state = setup();
+
+        const numberCount = selectSignal(
+          state,
+          (state) => state.numbers.length
+        );
+
+        const numberEffect = jest.fn(() => numberCount());
+        effect(numberEffect);
+
+        // run effects for the first time
+        detectChanges();
+        expect(numberEffect).toHaveBeenCalledTimes(1);
+
+        // update user
+        state.$update({
+          user: { firstName: 'Susanne', lastName: 'Taylor' },
+        });
+        detectChanges();
+        expect(numberEffect).toHaveBeenCalledTimes(1);
+
+        // update numbers
+        state.$update({ numbers: [1] });
+        detectChanges();
+        expect(numberEffect).toHaveBeenCalledTimes(2);
+        expect(numberCount()).toBe(1);
+      })
+    );
+  });
+
+  describe('immutability', () => {
+    it(
+      'should throw on mutable changes',
+      withEffect((detectChanges) => {
+        let numberCounter = 0;
+        const state = setup();
+        const effectFn = jest.fn(() => state.numbers());
+        effect(effectFn);
+        detectChanges();
+        expect(effectFn).toHaveBeenCalledTimes(1);
+
+        expect(() =>
+          state.$update((state) => {
+            const { numbers } = state;
+            numbers.push(4);
+            return { ...state };
+          })
+        ).toThrow();
+      })
+    );
+
+    it(
+      'should throw on a single mutable change',
+      withEffect((detectChanges) => {
+        let numberCounter = 0;
+        const state = setup();
+        const effectFn = jest.fn(() => state.numbers());
+        effect(effectFn);
+        detectChanges();
+        expect(effectFn).toHaveBeenCalledTimes(1);
+
+        expect(() =>
+          state.$update({ foo: 'foobar' }, (state) => {
+            const { numbers } = state;
+            numbers.push(4);
+            return { ...state };
+          })
+        ).toThrow();
+      })
+    );
+  });
+});
diff --git a/src/tests/test2.ts b/src/tests/test2.ts
index ac61cc6..b2a57c4 100644
--- a/src/tests/test2.ts
+++ b/src/tests/test2.ts
@@ -1,5 +1,4 @@
 import {
-  selectSignal,
   signalStore,
   signalStoreFeature,
   type,
diff --git a/src/tests/with-effect.ts b/src/tests/with-effect.ts
new file mode 100644
index 0000000..c07d9ec
--- /dev/null
+++ b/src/tests/with-effect.ts
@@ -0,0 +1,26 @@
+import { Component } from '@angular/core';
+import { TestBed } from '@angular/core/testing';
+
+/**
+ * Testing function which executes inside of an injectionContext and provides Change Detection trigger.
+ *
+ * It looks like, we need to have at least one component to execute effects.
+ * AppRef::tick() does not work. Only ComponentFixture::detectChanges() seems to do the job.
+ * withEffects renders a TestComponent and executes the tests inside an injectionContext.
+ * This allows us to generate effects very easily.
+ */
+export const withEffect =
+  (testFn: (detectChanges: () => void) => void): (() => void) =>
+  () => {
+    const fixture = TestBed.configureTestingModule({
+      imports: [TestComponent],
+    }).createComponent(TestComponent);
+
+    TestBed.runInInjectionContext(() => testFn(() => fixture.detectChanges()));
+  };
+
+@Component({
+  template: '',
+  standalone: true,
+})
+class TestComponent {}

From 603d0a54b427202f6f8449137bef9d4e48cf0969 Mon Sep 17 00:00:00 2001
From: Rainer Hahnekamp <rainer.hahnekamp@gmail.com>
Date: Tue, 8 Aug 2023 11:39:49 +0200
Subject: [PATCH 2/5] only freeze state in signal state, if the state has been
 updated with a function, i.e. state is exposed.

@markostanimirovic: I am not happy with the `@ts-ignore` in `deepFreeze`. Maybe you can find a better way.
---
 src/lib/signal-state.ts | 14 ++++++++------
 src/lib/utils.ts        |  5 +++--
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/lib/signal-state.ts b/src/lib/signal-state.ts
index 5869399..887118f 100644
--- a/src/lib/signal-state.ts
+++ b/src/lib/signal-state.ts
@@ -31,12 +31,14 @@ export function signalStateUpdateFactory<State extends Record<string, unknown>>(
 ): SignalStateUpdate<State>['$update'] {
   return (...updaters) =>
     stateSignal.update((state) =>
-      updaters.reduce((currentState: State, updater) => {
-        deepFreeze(currentState);
-        return {
+      updaters.reduce(
+        (currentState: State, updater) => ({
           ...currentState,
-          ...(typeof updater === 'function' ? updater(currentState) : updater),
-        };
-      }, state)
+          ...(typeof updater === 'function'
+            ? updater(deepFreeze(currentState))
+            : updater),
+        }),
+        state
+      )
     );
 }
diff --git a/src/lib/utils.ts b/src/lib/utils.ts
index 2da10e5..f4e8599 100644
--- a/src/lib/utils.ts
+++ b/src/lib/utils.ts
@@ -1,4 +1,4 @@
-export function deepFreeze(target: any) {
+export function deepFreeze<T extends object>(target: T): T {
   Object.freeze(target);
 
   const targetIsFunction = isFunction(target);
@@ -16,7 +16,8 @@ export function deepFreeze(target: any) {
         ? prop !== 'caller' && prop !== 'callee' && prop !== 'arguments'
         : true)
     ) {
-      const propValue: unknown = target[prop];
+      // @ts-ignore
+      const propValue = target[prop];
 
       if (
         (isObjectLike(propValue) || isFunction(propValue)) &&

From dfa5e7e78ec163025c9a1c8f5bf10228b9f74b26 Mon Sep 17 00:00:00 2001
From: Rainer Hahnekamp <rainer.hahnekamp@gmail.com>
Date: Wed, 9 Aug 2023 21:22:33 +0200
Subject: [PATCH 3/5] - rename util to deep-freeze - rename withEffect to
 testEffects - type assertion in freezing function - move signal state spec
 back to original location

---
 src/lib/{utils.ts => deep-freeze.ts}          | 7 +++----
 src/{tests => lib}/signal-state.spec.ts       | 0
 src/tests/{with-effect.ts => test-effects.ts} | 2 +-
 3 files changed, 4 insertions(+), 5 deletions(-)
 rename src/lib/{utils.ts => deep-freeze.ts} (81%)
 rename src/{tests => lib}/signal-state.spec.ts (100%)
 rename src/tests/{with-effect.ts => test-effects.ts} (97%)

diff --git a/src/lib/utils.ts b/src/lib/deep-freeze.ts
similarity index 81%
rename from src/lib/utils.ts
rename to src/lib/deep-freeze.ts
index f4e8599..c27da95 100644
--- a/src/lib/utils.ts
+++ b/src/lib/deep-freeze.ts
@@ -16,8 +16,7 @@ export function deepFreeze<T extends object>(target: T): T {
         ? prop !== 'caller' && prop !== 'callee' && prop !== 'arguments'
         : true)
     ) {
-      // @ts-ignore
-      const propValue = target[prop];
+      const propValue = (target as Record<string, unknown>)[prop];
 
       if (
         (isObjectLike(propValue) || isFunction(propValue)) &&
@@ -31,10 +30,10 @@ export function deepFreeze<T extends object>(target: T): T {
   return target;
 }
 
-export function isFunction(target: unknown): target is () => unknown {
+function isFunction(target: unknown): target is () => unknown {
   return typeof target === 'function';
 }
 
-export function isObjectLike(target: unknown): target is object {
+function isObjectLike(target: unknown): target is object {
   return typeof target === 'object' && target !== null;
 }
diff --git a/src/tests/signal-state.spec.ts b/src/lib/signal-state.spec.ts
similarity index 100%
rename from src/tests/signal-state.spec.ts
rename to src/lib/signal-state.spec.ts
diff --git a/src/tests/with-effect.ts b/src/tests/test-effects.ts
similarity index 97%
rename from src/tests/with-effect.ts
rename to src/tests/test-effects.ts
index c07d9ec..5176fc0 100644
--- a/src/tests/with-effect.ts
+++ b/src/tests/test-effects.ts
@@ -9,7 +9,7 @@ import { TestBed } from '@angular/core/testing';
  * withEffects renders a TestComponent and executes the tests inside an injectionContext.
  * This allows us to generate effects very easily.
  */
-export const withEffect =
+export const testEffects =
   (testFn: (detectChanges: () => void) => void): (() => void) =>
   () => {
     const fixture = TestBed.configureTestingModule({

From baa6b6057ccdb003c4b12b75082aa23a6e405769 Mon Sep 17 00:00:00 2001
From: Rainer Hahnekamp <rainer.hahnekamp@gmail.com>
Date: Wed, 9 Aug 2023 22:32:22 +0200
Subject: [PATCH 4/5] fix tests

---
 src/lib/signal-state.spec.ts | 14 +++++++-------
 src/lib/signal-state.ts      |  2 +-
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/lib/signal-state.spec.ts b/src/lib/signal-state.spec.ts
index dd3db5f..4dafbcb 100644
--- a/src/lib/signal-state.spec.ts
+++ b/src/lib/signal-state.spec.ts
@@ -1,7 +1,7 @@
 import { effect } from '@angular/core';
-import { withEffect } from './with-effect';
-import { signalState } from '../lib/signal-state';
-import { selectSignal } from '../lib/select-signal';
+import { testEffects } from '../tests/test-effects';
+import { signalState } from './signal-state';
+import { selectSignal } from './select-signal';
 
 describe('Signal State', () => {
   const initialState = {
@@ -50,7 +50,7 @@ describe('Signal State', () => {
   describe('equal checks', () => {
     it(
       'should not fire unchanged signals on update',
-      withEffect((detectChanges) => {
+      testEffects((detectChanges) => {
         const state = setup();
 
         const numberEffect = jest.fn(() => state.numbers());
@@ -77,7 +77,7 @@ describe('Signal State', () => {
 
     it(
       'should not fire for unchanged derived signals',
-      withEffect((detectChanges) => {
+      testEffects((detectChanges) => {
         const state = setup();
 
         const numberCount = selectSignal(
@@ -111,7 +111,7 @@ describe('Signal State', () => {
   describe('immutability', () => {
     it(
       'should throw on mutable changes',
-      withEffect((detectChanges) => {
+      testEffects((detectChanges) => {
         let numberCounter = 0;
         const state = setup();
         const effectFn = jest.fn(() => state.numbers());
@@ -131,7 +131,7 @@ describe('Signal State', () => {
 
     it(
       'should throw on a single mutable change',
-      withEffect((detectChanges) => {
+      testEffects((detectChanges) => {
         let numberCounter = 0;
         const state = setup();
         const effectFn = jest.fn(() => state.numbers());
diff --git a/src/lib/signal-state.ts b/src/lib/signal-state.ts
index 887118f..004dbba 100644
--- a/src/lib/signal-state.ts
+++ b/src/lib/signal-state.ts
@@ -1,7 +1,7 @@
 import { signal, WritableSignal } from '@angular/core';
 import { DeepSignal, toDeepSignal } from './deep-signal';
 import { defaultEqualityFn } from './select-signal';
-import { deepFreeze, isFunction, isObjectLike } from './utils';
+import { deepFreeze } from './deep-freeze';
 
 export type SignalState<State extends Record<string, unknown>> =
   DeepSignal<State> & SignalStateUpdate<State>;

From 491f21c33ef0bd3cb9438ea9d57b21a79cb715f6 Mon Sep 17 00:00:00 2001
From: Rainer Hahnekamp <rainer.hahnekamp@gmail.com>
Date: Fri, 11 Aug 2023 16:19:13 +0200
Subject: [PATCH 5/5] freeze signalState optionally (only in dev mode). $update
 function provides state as readonly

- refactored tests
- added further tests for mutable and immutable changes
---
 src/lib/deep-freeze.ts            |   6 +-
 src/lib/deep-readonly.ts          |   5 +
 src/lib/signal-state.prod.spec.ts |  24 +++++
 src/lib/signal-state.spec.ts      | 173 ++++++++++++++++++++----------
 src/lib/signal-state.ts           |  47 +++++---
 src/tests/test-effects.ts         |   2 +-
 6 files changed, 184 insertions(+), 73 deletions(-)
 create mode 100644 src/lib/deep-readonly.ts
 create mode 100644 src/lib/signal-state.prod.spec.ts

diff --git a/src/lib/deep-freeze.ts b/src/lib/deep-freeze.ts
index c27da95..bd0fd42 100644
--- a/src/lib/deep-freeze.ts
+++ b/src/lib/deep-freeze.ts
@@ -1,4 +1,6 @@
-export function deepFreeze<T extends object>(target: T): T {
+import { DeepReadonly } from './deep-readonly';
+
+export function deepFreeze<T extends object>(target: T): DeepReadonly<T> {
   Object.freeze(target);
 
   const targetIsFunction = isFunction(target);
@@ -27,7 +29,7 @@ export function deepFreeze<T extends object>(target: T): T {
     }
   }
 
-  return target;
+  return target as DeepReadonly<T>;
 }
 
 function isFunction(target: unknown): target is () => unknown {
diff --git a/src/lib/deep-readonly.ts b/src/lib/deep-readonly.ts
new file mode 100644
index 0000000..2f58038
--- /dev/null
+++ b/src/lib/deep-readonly.ts
@@ -0,0 +1,5 @@
+export type DeepReadonly<T> = T extends Array<infer AT>
+  ? ReadonlyArray<AT>
+  : T extends Record<string, unknown>
+  ? Readonly<{ [Key in keyof T]: DeepReadonly<T[Key]> }>
+  : Readonly<T>;
diff --git a/src/lib/signal-state.prod.spec.ts b/src/lib/signal-state.prod.spec.ts
new file mode 100644
index 0000000..b9f1046
--- /dev/null
+++ b/src/lib/signal-state.prod.spec.ts
@@ -0,0 +1,24 @@
+import { testEffects } from '../tests/test-effects';
+import { enableProdMode } from '@angular/core';
+import { signalState } from './signal-state';
+
+// needs a separate file because we can't revert the prod mode.
+it(
+  'should not freeze state with immutability check in prod mode',
+  testEffects(() => {
+    enableProdMode();
+    const state = signalState(
+      {
+        person: { firstName: 'John', lastName: 'Smith' },
+      },
+      { immutabilityCheck: true }
+    );
+
+    expect(() =>
+      state.$update((state) => {
+        (state as any).person.firstName = 'Johannes';
+        return state;
+      })
+    ).not.toThrow();
+  })
+);
diff --git a/src/lib/signal-state.spec.ts b/src/lib/signal-state.spec.ts
index 4dafbcb..7567a5d 100644
--- a/src/lib/signal-state.spec.ts
+++ b/src/lib/signal-state.spec.ts
@@ -1,56 +1,68 @@
-import { effect } from '@angular/core';
+import { effect, isDevMode } from '@angular/core';
+import * as angularCore from '@angular/core';
 import { testEffects } from '../tests/test-effects';
 import { signalState } from './signal-state';
 import { selectSignal } from './select-signal';
 
 describe('Signal State', () => {
-  const initialState = {
+  interface TestState {
+    user: { firstName: string; lastName: string };
+    foo: string;
+    numbers: number[];
+  }
+
+  const getInitialState = () => ({
     user: {
       firstName: 'John',
       lastName: 'Smith',
     },
     foo: 'bar',
     numbers: [1, 2, 3],
-  };
-
-  const setup = () => signalState(initialState);
+  });
 
-  it('should support nested signals', () => {
-    const state = setup();
+  describe('Basics', () => {
+    it('should support nested signals', () => {
+      const initialState = getInitialState();
+      const state = signalState(initialState);
 
-    expect(state()).toBe(initialState);
-    expect(state.user()).toBe(initialState.user);
-    expect(state.user.firstName()).toBe(initialState.user.firstName);
-  });
+      expect(state()).toBe(initialState);
+      expect(state.user()).toBe(initialState.user);
+      expect(state.user.firstName()).toBe('John');
+    });
 
-  it('should allow updates', () => {
-    const state = setup();
-    state.$update((state) => ({
-      ...state,
-      user: { firstName: 'Johannes', lastName: 'Schmidt' },
-    }));
-    expect(state()).toEqual({
-      ...initialState,
-      user: { firstName: 'Johannes', lastName: 'Schmidt' },
+    it('should allow updates', () => {
+      const initialState = getInitialState();
+      const state = signalState(initialState);
+      state.$update((state) => ({
+        ...state,
+        user: { firstName: 'Johannes', lastName: 'Schmidt' },
+      }));
+      expect(state()).toEqual({
+        ...initialState,
+        user: { firstName: 'Johannes', lastName: 'Schmidt' },
+      });
     });
-  });
 
-  it('should update immutably', () => {
-    const state = setup();
-    state.$update((state) => ({
-      ...state,
-      foo: 'bar',
-      numbers: [3, 2, 1],
-    }));
-    expect(state.user()).toBe(initialState.user);
-    expect(state.foo()).toBe(initialState.foo);
-    expect(state.numbers()).not.toBe(initialState.numbers);
+    it('should update immutably', () => {
+      const initialState = getInitialState();
+      const state = signalState(initialState);
+      state.$update((state) => ({
+        ...state,
+        foo: 'bar',
+        numbers: [3, 2, 1],
+      }));
+      expect(state.user()).toBe(initialState.user);
+      expect(state.foo()).toBe(initialState.foo);
+      expect(state.numbers()).not.toBe(initialState.numbers);
+    });
   });
 
   describe('equal checks', () => {
+    const setup = () => signalState(getInitialState());
+
     it(
       'should not fire unchanged signals on update',
-      testEffects((detectChanges) => {
+      testEffects((runEffects) => {
         const state = setup();
 
         const numberEffect = jest.fn(() => state.numbers());
@@ -63,13 +75,13 @@ describe('Signal State', () => {
         expect(userEffect).toHaveBeenCalledTimes(0);
 
         // run effects for the first time
-        detectChanges();
+        runEffects();
         expect(numberEffect).toHaveBeenCalledTimes(1);
         expect(userEffect).toHaveBeenCalledTimes(1);
 
         // update state with effect run
         state.$update((state) => ({ ...state, numbers: [4, 5, 6] }));
-        detectChanges();
+        runEffects();
         expect(numberEffect).toHaveBeenCalledTimes(2);
         expect(userEffect).toHaveBeenCalledTimes(1);
       })
@@ -77,7 +89,7 @@ describe('Signal State', () => {
 
     it(
       'should not fire for unchanged derived signals',
-      testEffects((detectChanges) => {
+      testEffects((runEffects) => {
         const state = setup();
 
         const numberCount = selectSignal(
@@ -89,19 +101,19 @@ describe('Signal State', () => {
         effect(numberEffect);
 
         // run effects for the first time
-        detectChanges();
+        runEffects();
         expect(numberEffect).toHaveBeenCalledTimes(1);
 
         // update user
         state.$update({
           user: { firstName: 'Susanne', lastName: 'Taylor' },
         });
-        detectChanges();
+        runEffects();
         expect(numberEffect).toHaveBeenCalledTimes(1);
 
         // update numbers
         state.$update({ numbers: [1] });
-        detectChanges();
+        runEffects();
         expect(numberEffect).toHaveBeenCalledTimes(2);
         expect(numberCount()).toBe(1);
       })
@@ -109,41 +121,88 @@ describe('Signal State', () => {
   });
 
   describe('immutability', () => {
+    const setup = (immutabilityCheck = false) =>
+      signalState(getInitialState(), { immutabilityCheck });
+
+    it(
+      'should run nested-based effects on mutable updates',
+      testEffects((runEffects) => {
+        let numberCounter = 0;
+        const state = setup();
+        const effectFn = jest.fn(() => state.user());
+        effect(effectFn);
+        runEffects();
+
+        // mutable update
+        state.$update((state) => {
+          (state as any).user = { firstName: 'John', lastName: 'Smith' };
+          return state;
+        });
+
+        runEffects();
+        expect(effectFn).toHaveBeenCalledTimes(2);
+      })
+    );
+
     it(
-      'should throw on mutable changes',
-      testEffects((detectChanges) => {
+      'should run selectSignal-based effects on mutable updates',
+      testEffects((runEffects) => {
         let numberCounter = 0;
         const state = setup();
-        const effectFn = jest.fn(() => state.numbers());
+        const userSignal = selectSignal(state, (state) => state.user);
+        const effectFn = jest.fn(() => userSignal());
         effect(effectFn);
-        detectChanges();
-        expect(effectFn).toHaveBeenCalledTimes(1);
+        runEffects();
+
+        // mutable update
+        state.$update((state) => {
+          (state as any).user = { firstName: 'John', lastName: 'Smith' };
+          return state;
+        });
+        runEffects();
+        expect(effectFn).toHaveBeenCalledTimes(2);
+      })
+    );
+
+    it(
+      'should not freeze on mutable updates',
+      testEffects(() => {
+        const state = setup();
 
         expect(() =>
           state.$update((state) => {
-            const { numbers } = state;
-            numbers.push(4);
+            (state as any).foo = 'bar';
             return { ...state };
           })
+        ).not.toThrow();
+      })
+    );
+
+    it(
+      'should freeze on mutable updates with immutability check',
+      testEffects(() => {
+        const state = setup(true);
+
+        expect(() =>
+          state.$update((state) => {
+            (state as any).foo = 'bar';
+            return state;
+          })
         ).toThrow();
       })
     );
 
     it(
-      'should throw on a single mutable change',
-      testEffects((detectChanges) => {
-        let numberCounter = 0;
-        const state = setup();
-        const effectFn = jest.fn(() => state.numbers());
-        effect(effectFn);
-        detectChanges();
-        expect(effectFn).toHaveBeenCalledTimes(1);
+      'should freeze on consecutive mutable updates with immutability check',
+      testEffects(() => {
+        const state = setup(true);
+
+        state.$update((state) => ({ ...state, foo: 'foobar' }));
 
         expect(() =>
-          state.$update({ foo: 'foobar' }, (state) => {
-            const { numbers } = state;
-            numbers.push(4);
-            return { ...state };
+          state.$update((state) => {
+            (state as any).foo = 'barfoo';
+            return state;
           })
         ).toThrow();
       })
diff --git a/src/lib/signal-state.ts b/src/lib/signal-state.ts
index 004dbba..afc7230 100644
--- a/src/lib/signal-state.ts
+++ b/src/lib/signal-state.ts
@@ -1,44 +1,65 @@
-import { signal, WritableSignal } from '@angular/core';
+import { isDevMode, signal, WritableSignal } from '@angular/core';
 import { DeepSignal, toDeepSignal } from './deep-signal';
 import { defaultEqualityFn } from './select-signal';
+import { DeepReadonly } from './deep-readonly';
 import { deepFreeze } from './deep-freeze';
 
 export type SignalState<State extends Record<string, unknown>> =
   DeepSignal<State> & SignalStateUpdate<State>;
 
+export type UpdatedState<T> = T extends ReadonlyArray<infer A>
+  ? Array<A> | ReadonlyArray<A>
+  : T extends Record<string, unknown>
+  ? { [Key in keyof T]: UpdatedState<T[Key]> }
+  : T;
+
 export type SignalStateUpdater<State extends Record<string, unknown>> =
   | Partial<State>
-  | ((state: State) => Partial<State>);
+  | ((state: DeepReadonly<State>) => UpdatedState<State>);
 
 export type SignalStateUpdate<State extends Record<string, unknown>> = {
   $update: (...updaters: SignalStateUpdater<State>[]) => void;
 };
 
 export function signalState<State extends Record<string, unknown>>(
-  initialState: State
+  initialState: State,
+  options: { immutabilityCheck: boolean } = { immutabilityCheck: false }
 ): SignalState<State> {
-  deepFreeze(initialState);
+  const freezeState = isDevMode() && options.immutabilityCheck;
+  if (freezeState) {
+    deepFreeze(initialState);
+  }
   const stateSignal = signal(initialState, { equal: defaultEqualityFn });
   const deepSignal = toDeepSignal(stateSignal.asReadonly());
-  (deepSignal as SignalState<State>).$update =
-    signalStateUpdateFactory(stateSignal);
+  (deepSignal as SignalState<State>).$update = signalStateUpdateFactory(
+    stateSignal,
+    freezeState
+  );
 
   return deepSignal as SignalState<State>;
 }
 
 export function signalStateUpdateFactory<State extends Record<string, unknown>>(
-  stateSignal: WritableSignal<State>
+  stateSignal: WritableSignal<State>,
+  freezeState: boolean
 ): SignalStateUpdate<State>['$update'] {
-  return (...updaters) =>
-    stateSignal.update((state) =>
-      updaters.reduce(
+  return (...updaters) => {
+    stateSignal.update((state) => {
+      const updatedState = updaters.reduce(
         (currentState: State, updater) => ({
           ...currentState,
           ...(typeof updater === 'function'
-            ? updater(deepFreeze(currentState))
+            ? updater(currentState as DeepReadonly<State>)
             : updater),
         }),
         state
-      )
-    );
+      );
+
+      if (freezeState) {
+        deepFreeze(updatedState);
+      }
+
+      return updatedState;
+    });
+  };
 }
diff --git a/src/tests/test-effects.ts b/src/tests/test-effects.ts
index 5176fc0..d8912c2 100644
--- a/src/tests/test-effects.ts
+++ b/src/tests/test-effects.ts
@@ -10,7 +10,7 @@ import { TestBed } from '@angular/core/testing';
  * This allows us to generate effects very easily.
  */
 export const testEffects =
-  (testFn: (detectChanges: () => void) => void): (() => void) =>
+  (testFn: (runEffects: () => void) => void): (() => void) =>
   () => {
     const fixture = TestBed.configureTestingModule({
       imports: [TestComponent],