perf(store): avoid going over states list every time action is dispat…

…ched (#2219)

In this commit, I'm working on optimizing how we handle action dispatches by avoiding
repeated traversal of the `states` list. Instead, we prepare a map each time a new state
is added, allowing us to perform O(1) lookups by action type in the future. This approach
reduces complexity and improves performance.

I've tested it with benchmark.js, and here are the results:

```
class Increment {
  static readonly type = 'Increment';
}

const states = Array.from({ length: 50 }).map((_, index) => {
  @State({
    name: `counter_${index + 1}`,
    defaults: 0,
  })
  @Injectable()
  class CounterState {
    @action(Increment)
    increment(ctx: StateContext<number>) {
      ctx.setState((counter) => counter + 1);
    }
  }

  return CounterState;
});

store.dispatch() before changes x 3,435 ops/sec ±0.45% (65 runs sampled)
store.dispatch() after changes x 3,942 ops/sec ±1.21% (25 runs sampled)
```

.bundlemonrc.json

@@ -30,7 +30,7 @@
     },
     {
       "path": "./fesm2022/ngxs-store.mjs",
-      "maxSize": "100kB",
+      "maxSize": "101kB",
       "maxPercentIncrease": 0.5
     }
   ],

CHANGELOG.md

@@ -7,6 +7,7 @@ $ npm install @ngxs/store@dev
 ### To become next patch version
 
 - Fix(store): allow selector utils usage within state class [#2210](https://github.com/ngxs/store/pull/2210)
+- Performance(store): Avoid going over states list every time action is dispatched [#2219](https://github.com/ngxs/store/pull/2219)
 
 ### 18.1.1 2024-08-06
 

packages/store/src/internal/lifecycle-state-manager.ts

@@ -99,6 +99,6 @@ export class LifecycleStateManager implements OnDestroy {
   }
 
   private _getStateContext(mappedStore: MappedStore): StateContext<any> {
-    return this._stateContextFactory.createStateContext(mappedStore);
+    return this._stateContextFactory.createStateContext(mappedStore.path);
   }
 }

packages/store/src/internal/state-context-factory.ts

@@ -4,7 +4,7 @@ import { ExistingState, StateOperator, isStateOperator } from '@ngxs/store/opera
 import { Observable } from 'rxjs';
 
 import { StateContext } from '../symbols';
-import { MappedStore, StateOperations } from '../internal/internals';
+import { StateOperations } from '../internal/internals';
 import { InternalStateOperations } from '../internal/state-operations';
 import { simplePatch } from './state-operators';
 
@@ -19,25 +19,25 @@ export class StateContextFactory {
   /**
    * Create the state context
    */
-  createStateContext<T>(mappedStore: MappedStore): StateContext<T> {
+  createStateContext<T>(path: string): StateContext<T> {
     const root = this._internalStateOperations.getRootStateOperations();
 
     return {
       getState(): T {
         const currentAppState = root.getState();
-        return getState(currentAppState, mappedStore.path);
+        return getState(currentAppState, path);
       },
       patchState(val: Partial<T>): void {
         const currentAppState = root.getState();
         const patchOperator = simplePatch<T>(val);
-        setStateFromOperator(root, currentAppState, patchOperator, mappedStore.path);
+        setStateFromOperator(root, currentAppState, patchOperator, path);
       },
       setState(val: T | StateOperator<T>): void {
         const currentAppState = root.getState();
         if (isStateOperator(val)) {
-          setStateFromOperator(root, currentAppState, val, mappedStore.path);
+          setStateFromOperator(root, currentAppState, val, path);
         } else {
-          setStateValue(root, currentAppState, val, mappedStore.path);
+          setStateValue(root, currentAppState, val, path);
         }
       },
       dispatch(actions: any | any[]): Observable<void> {

packages/store/src/internal/state-factory.ts

@@ -17,7 +17,8 @@ import {
   ɵStateClassInternal,
   ɵINITIAL_STATE_TOKEN,
   ɵSharedSelectorOptions,
-  ɵRuntimeSelectorContext
+  ɵRuntimeSelectorContext,
+  ɵActionHandlerMetaData
 } from '@ngxs/store/internals';
 import { getActionTypeFromInstance, getValue, setValue } from '@ngxs/store/plugins';
 import {
@@ -78,6 +79,11 @@ function cloneDefaults(defaults: any): any {
   return value;
 }
 
+interface InvokableActionHandlerMetaData extends ɵActionHandlerMetaData {
+  path: string;
+  instance: any;
+}
+
 /**
  * The `StateFactory` class adds root and feature states to the graph.
  * This extracts state names from state classes, checks if they already
@@ -113,6 +119,19 @@ export class StateFactory implements OnDestroy {
     private _initialState: any
   ) {}
 
+  // Instead of going over the states list every time an action is dispatched,
+  // we are constructing a map of action types to lists of action metadata.
+  // If the `@@Init` action is handled in two different states, the action
+  // metadata list will contain two objects that have the state `instance` and
+  // method names to be used as action handlers (decorated with `@Action(InitState)`).
+  private _actionTypeToMetasMap = new Map<string, InvokableActionHandlerMetaData[]>();
+
+  get actionTypeToMetasMap(): Map<string, InvokableActionHandlerMetaData[]> {
+    return this._parentFactory
+      ? this._parentFactory.actionTypeToMetasMap
+      : this._actionTypeToMetasMap;
+  }
+
   private _states: MappedStore[] = [];
 
   get states(): MappedStore[] {
@@ -223,6 +242,7 @@ export class StateFactory implements OnDestroy {
       }
 
       this.states.push(stateMap);
+      this.hydrateActionMetasMap(stateMap);
     }
 
     return bootstrappedStores;
@@ -290,64 +310,62 @@ export class StateFactory implements OnDestroy {
     // to `true` within the below `for` loop if any `actionMetas` has been found.
     let actionHasBeenHandled = false;
 
-    for (const metadata of this.states) {
-      const actionMetas = metadata.actions[type];
-
-      if (actionMetas) {
-        for (const actionMeta of actionMetas) {
-          const stateContext = this._stateContextFactory.createStateContext(metadata);
-          try {
-            let result = metadata.instance[actionMeta.fn](stateContext, action);
-
-            // We need to use `isPromise` instead of checking whether
-            // `result instanceof Promise`. In zone.js patched environments, `global.Promise`
-            // is the `ZoneAwarePromise`. Some APIs, which are likely not patched by zone.js
-            // for certain reasons, might not work with `instanceof`. For instance, the dynamic
-            // import returns a native promise (not a `ZoneAwarePromise`), causing this check to
-            // be falsy.
-            if (ɵisPromise(result)) {
-              result = from(result);
-            }
+    const actionMetas = this.actionTypeToMetasMap.get(type);
+
+    if (actionMetas) {
+      for (const actionMeta of actionMetas) {
+        const stateContext = this._stateContextFactory.createStateContext(actionMeta.path);
 
-            if (isObservable(result)) {
-              // If this observable has been completed w/o emitting
-              // any value then we wouldn't want to complete the whole chain
-              // of actions. Since if any observable completes then
-              // action will be canceled.
-              // For instance if any action handler would've had such statement:
-              // `handler(ctx) { return EMPTY; }`
-              // then the action will be canceled.
+        let result;
+        try {
+          result = actionMeta.instance[actionMeta.fn](stateContext, action);
+
+          // We need to use `isPromise` instead of checking whether
+          // `result instanceof Promise`. In zone.js patched environments, `global.Promise`
+          // is the `ZoneAwarePromise`. Some APIs, which are likely not patched by zone.js
+          // for certain reasons, might not work with `instanceof`. For instance, the dynamic
+          // import returns a native promise (not a `ZoneAwarePromise`), causing this check to
+          // be falsy.
+          if (ɵisPromise(result)) {
+            result = from(result);
+          }
+
+          if (isObservable(result)) {
+            result = result.pipe(
+              mergeMap((value: any) => {
+                if (ɵisPromise(value)) {
+                  return from(value);
+                }
+                if (isObservable(value)) {
+                  return value;
+                }
+                return of(value);
+              }),
+              // If this observable has completed without emitting any values,
+              // we wouldn't want to complete the entire chain of actions.
+              // If any observable completes, then the action will be canceled.
+              // For instance, if any action handler had a statement like
+              // `handler(ctx) { return EMPTY; }`, then the action would be canceled.
               // See https://github.com/ngxs/store/issues/1568
+              defaultIfEmpty({})
+            );
+
+            if (actionMeta.options.cancelUncompleted) {
+              // todo: ofActionDispatched should be used with action class
               result = result.pipe(
-                mergeMap((value: any) => {
-                  if (ɵisPromise(value)) {
-                    return from(value);
-                  }
-                  if (isObservable(value)) {
-                    return value;
-                  }
-                  return of(value);
-                }),
-                defaultIfEmpty({})
+                takeUntil(dispatched$.pipe(ofActionDispatched(action as any)))
               );
-
-              if (actionMeta.options.cancelUncompleted) {
-                // todo: ofActionDispatched should be used with action class
-                result = result.pipe(
-                  takeUntil(dispatched$.pipe(ofActionDispatched(action as any)))
-                );
-              }
-            } else {
-              result = of({}).pipe(shareReplay());
             }
-
-            results.push(result);
-          } catch (e) {
-            results.push(throwError(e));
+          } else {
+            result = of({}).pipe(shareReplay());
           }
-
-          actionHasBeenHandled = true;
+        } catch (e) {
+          result = throwError(e);
         }
+
+        results.push(result);
+
+        actionHasBeenHandled = true;
       }
     }
 
@@ -358,9 +376,7 @@ export class StateFactory implements OnDestroy {
       // The `NgxsUnhandledActionsLogger` will not be resolved by the injector if the
       // `NgxsDevelopmentModule` is not provided. It's enough to check whether the `injector.get`
       // didn't return `null` so we may ensure the module has been imported.
-      if (unhandledActionsLogger) {
-        unhandledActionsLogger.warn(action);
-      }
+      unhandledActionsLogger?.warn(action);
     }
 
     if (!results.length) {
@@ -406,4 +422,31 @@ export class StateFactory implements OnDestroy {
     // its lifecycle is in 'bootstrapped' state.
     return this.statesByName[name] && valueIsBootstrappedInInitialState;
   }
+
+  private hydrateActionMetasMap({ path, actions, instance }: MappedStore): void {
+    const actionTypeToMetasMap = this.actionTypeToMetasMap;
+
+    for (const actionType of Object.keys(actions)) {
+      // Initialize the map entry if it does not already exist for that
+      // action type. Note that action types may overlap between states,
+      // as the same action can be handled by different states.
+      if (!actionTypeToMetasMap.has(actionType)) {
+        actionTypeToMetasMap.set(actionType, []);
+      }
+
+      const extendedActionMetas = actionTypeToMetasMap.get(actionType)!;
+
+      extendedActionMetas.push(
+        // This involves combining each individual action metadata with
+        // the state instance and the path—essentially everything needed
+        // to invoke an action. This eliminates the need to loop over states
+        // every time an action is dispatched.
+        ...actions[actionType].map(actionMeta => ({
+          ...actionMeta,
+          path,
+          instance
+        }))
+      );
+    }
+  }
 }