feat: add setAtom bound to the state manager

example/src/modules/x/childState/childState.js

@@ -2,14 +2,14 @@ import { defineState } from '@lwc/state';
 import parentStateFactory from 'x/parentState';
 import anotherParentStateFactory from 'x/anotherParentState';
 
-export default defineState((atom, _computed, update, fromContext) => (initialName = 'bar') => {
+export default defineState((atom, computed, fromContext, setAtom) => (initialName = 'bar') => {
   const name = atom(initialName);
   const parentState = fromContext(parentStateFactory);
   const anotherParentState = fromContext(anotherParentStateFactory);
 
-  const updateName = update({ name }, (_, newName) => ({
-    name: newName,
-  }));
+  const updateName = (newName) => {
+    setAtom(name, newName);
+  };
 
   return {
     name,

example/src/modules/x/grandChildState/grandChildState.js

@@ -1,7 +1,7 @@
 import { defineState } from '@lwc/state';
 import childStateFactory from 'x/childState';
 
-export default defineState((_atom, _computed, _update, fromContext) => () => {
+export default defineState((_atom, _computed, fromContext) => () => {
   const childState = fromContext(childStateFactory);
 
   return {

example/src/modules/x/parentState/parentState.js

@@ -1,11 +1,11 @@
 import { defineState } from '@lwc/state';
 
-export default defineState((atom, computed, update, fromContext) => (initialName = 'foo') => {
+export default defineState((atom, computed, fromContext, setAtom) => (initialName = 'foo') => {
   const name = atom(initialName);
 
-  const updateName = update({ name }, (_, newName) => ({
-    name: newName,
-  }));
+  const updateName = (newName) => {
+    setAtom(name, newName);
+  };
 
   return {
     name,

packages/@lwc/state/README.md

@@ -28,8 +28,10 @@ The main function for creating state definitions. It provides four utilities:
 
 - `atom<T>(initialValue: T)`: Creates a reactive atomic value
 - `computed(signals, computeFn)`: Creates derived state based on provided signals map
-- `update(signals, updateFn)`: Creates state mutation functions
 - `fromContext(stateDefinition)`: Consumes context from parent components
+- `setAtom<T>(signal: Signal<T>, newValue: T)`: Directly sets the value of an atom signal. Can only be used to modify atoms that were created within the same state manager instance. Attempting to modify atoms from other state managers will:
+  - In development: Throw a ReferenceError
+  - In production: Silently fail without modifying the atom
 
 ### ContextfulLightningElement
 
@@ -46,20 +48,27 @@ Create a state definition using `defineState`:
 import { defineState } from '@lwc/state';
 
 const useCounter = defineState(
-  (atom, computed, update) =>
+  (atom, computed, _fromContext, setAtom) =>
     (initialValue = 0) => {
       // Create reactive atom
       const count = atom(initialValue);
       // Create computed value
       const doubleCount = computed({ count }, ({ count }) => count * 2);
       // Create update function
-      const increment = update({ count }, ({ count }) => ({
-        count: count + 1,
-      }));
+      const increment = () => {
+        setAtom(count, count.value + 1)
+      };
+
+      // Create a function that directly sets the count atom
+      const setCount = (newValue: number) => {
+        setAtom(count, newValue);
+      };
+
       return {
         count,
         doubleCount,
         increment,
+        setCount,
       };
     }
 );
@@ -106,7 +115,7 @@ const context = defineState(
 import contextFactory from '<parentState>';
 
 const useTheme = defineState(
-  (_atom, _computed, _update, fromContext) =>
+  (_atom, _computed, fromContext) =>
     () => {
       const theme = fromContext(contextFactory);
 

packages/@lwc/state/src/__tests__/index.test.ts

@@ -1,13 +1,23 @@
 import { afterEach, describe, expect, test, vi } from 'vitest';
 import { defineState } from '../index.js';
+
 // biome-ignore lint: test only
 let doubleCountNotifySpy: any;
 // biome-ignore lint: test only
 let fruitNameAndCountNotifySpy: any;
 // biome-ignore lint: test only
 let fruitNameAndCountComputeValueSpy: any;
 
-const state = defineState((atom, computed, update, _fromContext) => (...args) => {
+const anotherSM = defineState((atom) => (...args) => {
+  const exposedAtom = atom(1);
+
+  return {
+    exposedAtom,
+    getExposedAtom: () => exposedAtom,
+  };
+})();
+
+const state = defineState((atom, computed, fromContext, setAtom) => (...args) => {
   const countArg = args[0] as number;
   const fruitArg = args[1] as string;
 
@@ -27,15 +37,26 @@ const state = defineState((atom, computed, update, _fromContext) => (...args) =>
   // @ts-ignore
   fruitNameAndCountComputeValueSpy = vi.spyOn(fruitNameAndCount, 'computeValue');
 
-  const increment = update({ count }, ({ count: countValue }) => ({
-    count: countValue + 1,
-  }));
-  const incrementBy = update({ count }, ({ count: countValue }, amount: number) => ({
-    count: countValue + amount,
-  }));
-  const changeFruit = update({ fruit }, ({ fruit }, newFruit: string) => ({
-    fruit: newFruit,
-  }));
+  const increment = () => {
+    setAtom(count, count.value + 1);
+  };
+
+  const incrementBy = (amount: number) => {
+    setAtom(count, count.value + amount);
+  };
+
+  const changeFruit = (newFruit: string) => {
+    setAtom(fruit, newFruit);
+  };
+
+  const multiplyBy = (num) => {
+    setAtom(count, count.value * num);
+  };
+
+  const tryModifyAtomsFromAnotherSM = () => {
+    const atomFromOtherSM = anotherSM.value.getExposedAtom();
+    setAtom(atomFromOtherSM, 3421);
+  };
 
   return {
     count,
@@ -44,6 +65,8 @@ const state = defineState((atom, computed, update, _fromContext) => (...args) =>
     incrementBy,
     fruitNameAndCount,
     changeFruit,
+    multiplyBy,
+    tryModifyAtomsFromAnotherSM,
   };
 });
 
@@ -89,6 +112,36 @@ describe('state manager', () => {
     expect(s.value.doubleCount).toBe(8);
   });
 
+  test('increment updates count and doubleCount', () => {
+    const s = state(1);
+    s.value.multiplyBy(3);
+
+    expect(s.value.count).toBe(3);
+    expect(s.value.doubleCount).toBe(6);
+  });
+
+  test('modifying atoms from another SM does not work in production', () => {
+    const beforeNodeEnvValue = process.env.NODE_ENV;
+    process.env.NODE_ENV = 'production';
+
+    const s = state(1);
+
+    expect(anotherSM.value.exposedAtom).toBe(1);
+    s.value.tryModifyAtomsFromAnotherSM();
+    expect(anotherSM.value.exposedAtom).toBe(1);
+
+    process.env.NODE_ENV = beforeNodeEnvValue;
+  });
+
+  test('modifying atoms from another SM does not work in non-production environments', () => {
+    const s = state(1);
+
+    expect(anotherSM.value.exposedAtom).toBe(1);
+    expect(() => s.value.tryModifyAtomsFromAnotherSM()).toThrow(
+      'The atom being set is not defined by this state manager.',
+    );
+  });
+
   test('multiple updates', () => {
     const s = state(0);
 

packages/@lwc/state/src/index.ts

@@ -8,7 +8,6 @@ import type {
   MakeAtom,
   MakeComputed,
   MakeContextHook,
-  MakeUpdate,
   UnwrapSignal,
 } from './types.ts';
 export { setTrustedSignalSet } from '@lwc/signals';
@@ -95,32 +94,27 @@ const atom: MakeAtom = <T,>(initialValue: T) => new AtomSignal<T>(initialValue);
 const computed: MakeComputed = (inputSignalsObj, computer) =>
   new ComputedSignal(inputSignalsObj, computer);
 
-const update: MakeUpdate = <
-  SignalSubType extends Signal<unknown>,
-  AdditionalArguments extends unknown[],
-  Updater extends (signalValues: ValuesObj, ...args: AdditionalArguments) => ValuesObj,
-  SignalsObj extends Record<string, SignalSubType>,
-  ValuesObj extends {
-    [SignalName in keyof SignalsObj]?: UnwrapSignal<SignalsObj[SignalName]>;
-  },
->(
-  signalsToUpdate: SignalsObj,
-  userProvidedUpdaterFn: Updater,
-) => {
-  return (...uniqueArgs: AdditionalArguments) => {
-    const signalValues = {} as ValuesObj;
-
-    for (const [signalName, signal] of Object.entries(signalsToUpdate)) {
-      signalValues[signalName as keyof ValuesObj] = signal.value as ValuesObj[keyof ValuesObj];
-    }
-
-    const newValues = userProvidedUpdaterFn(signalValues, ...uniqueArgs);
-
-    for (const [atomName, newValue] of Object.entries(newValues)) {
-      signalsToUpdate[atomName][atomSetter](newValue);
-    }
+function createStateManagerAtomsClosure() {
+  const smAtoms = new WeakSet();
+
+  return {
+    atom: <T,>(v: T) => {
+      const a = atom(v);
+      smAtoms.add(a);
+
+      return a;
+    },
+    setAtom: <T = unknown>(a: AtomSignal<T>, v: T) => {
+      if (smAtoms.has(a)) {
+        a[atomSetter](v);
+      } else {
+        if (process.env.NODE_ENV !== 'production') {
+          throw new ReferenceError('The atom being set is not defined by this state manager.');
+        }
+      }
+    },
   };
-};
+}
 
 export const defineState: DefineState = <
   InnerStateShape extends Record<string, Signal<unknown> | ExposedUpdater>,
@@ -133,8 +127,8 @@ export const defineState: DefineState = <
   defineStateCallback: (
     atom: MakeAtom,
     computed: MakeComputed,
-    update: MakeUpdate,
     fromContext: MakeContextHook<ContextShape>,
+    setAtom: <T>(signal: AtomSignal<T>, newValue: T) => void,
   ) => (...args: Args) => InnerStateShape,
 ) => {
   const stateDefinition = (...args: Args) => {
@@ -194,7 +188,13 @@ export const defineState: DefineState = <
           return localContextSignal;
         };
 
-        this.internalStateShape = defineStateCallback(atom, computed, update, fromContext)(...args);
+        const { atom, setAtom } = createStateManagerAtomsClosure();
+        this.internalStateShape = defineStateCallback(
+          atom,
+          computed,
+          fromContext,
+          setAtom,
+        )(...args);
 
         for (const signalOrUpdater of Object.values(this.internalStateShape)) {
           if (signalOrUpdater && !isUpdater(signalOrUpdater)) {

packages/@lwc/state/src/types.ts

@@ -20,17 +20,6 @@ export type MakeComputed = <
 
 export type MutatorArgs = Array<unknown>;
 
-export type MakeUpdate = <
-  SignalSubType extends Signal<unknown>,
-  SignalsToMutate extends Record<string, SignalSubType>,
-  Values extends {
-    [SignalName in keyof SignalsToMutate]?: UnwrapSignal<SignalsToMutate[SignalName]>;
-  },
->(
-  signalsToMutate: SignalsToMutate,
-  mutator: (signalValues: Values, ...mutatorArgs: MutatorArgs) => Values,
-) => (...mutatorArgs: MutatorArgs) => void;
-
 export type MakeContextHook<T> = <StateDef extends () => Signal<T>>(
   stateDef: StateDef,
 ) => Signal<T>;
@@ -50,7 +39,7 @@ export type DefineState = <
   defineStateCb: (
     atom: MakeAtom,
     computed: MakeComputed,
-    update: MakeUpdate,
     fromContext: MakeContextHook<ContextShape>,
+    setAtom: <T>(a: Signal<T>, newValue: T) => void,
   ) => (...args: Args) => InnerStateShape,
 ) => (...args: Args) => Signal<OuterStateShape>;