feat(store): allow extending State({...}) instead of decorating

packages/store/src/decorators/state.ts

@@ -1,27 +1,44 @@
+import { inject } from '@angular/core';
 import {
-  ɵStateClass,
   ɵMETA_KEY,
   ɵMETA_OPTIONS_KEY,
   ɵMetaDataModel,
+  ɵStateClass,
   ɵStateClassInternal,
   ɵStoreOptions,
   ɵensureStoreMetadata
 } from '@ngxs/store/internals';
+import { Observable } from 'rxjs';
 
+import { StateOperator } from '../symbols';
 import { ensureStateNameIsValid } from '../utils/store-validators';
+import { StateContextFactory } from '../internal/state-context-factory';
 
 interface MutateMetaOptions<T> {
   meta: ɵMetaDataModel;
   inheritedStateClass: ɵStateClassInternal;
   optionsWithInheritance: ɵStoreOptions<T>;
 }
 
+interface BaseState<T> {
+  getState(): T;
+  setState(value: T): void;
+  patchState(value: Partial<T>): void;
+  dispatch(actions: any | any[]): Observable<void>;
+}
+
+interface StateDecorator<T> {
+  (target: ɵStateClass): void;
+  new (): BaseState<T>;
+}
+
 /**
  * Decorates a class with ngxs state information.
  */
-export function State<T>(options: ɵStoreOptions<T>) {
-  return (target: ɵStateClass): void => {
-    const stateClass: ɵStateClassInternal = target;
+export function State<T>(options: ɵStoreOptions<T>): StateDecorator<T> {
+  return function (this: BaseState<T>) {
+    // eslint-disable-next-line prefer-rest-params
+    const stateClass = (new.target || arguments[0]) as unknown as ɵStateClassInternal;
     const meta: ɵMetaDataModel = ɵensureStoreMetadata(stateClass);
     const inheritedStateClass: ɵStateClassInternal = Object.getPrototypeOf(stateClass);
     const optionsWithInheritance: ɵStoreOptions<T> = getStateOptions(
@@ -30,7 +47,32 @@ export function State<T>(options: ɵStoreOptions<T>) {
     );
     mutateMetaData<T>({ meta, inheritedStateClass, optionsWithInheritance });
     stateClass[ɵMETA_OPTIONS_KEY] = optionsWithInheritance;
-  };
+
+    // If this function is being called as a constructor function.
+    if (new.target) {
+      const stateContextFactory = inject(StateContextFactory);
+
+      this.getState = () => {
+        const ctx = stateContextFactory.createStateContext<T>(meta.path!);
+        return ctx.getState();
+      };
+
+      this.setState = (value: T | StateOperator<T>) => {
+        const ctx = stateContextFactory.createStateContext<T>(meta.path!);
+        ctx.setState(value);
+      };
+
+      this.patchState = (value: Partial<T>) => {
+        const ctx = stateContextFactory.createStateContext<T>(meta.path!);
+        ctx.patchState(value);
+      };
+
+      this.dispatch = actions => {
+        const ctx = stateContextFactory.createStateContext<T>(meta.path!);
+        return ctx.dispatch(actions);
+      };
+    }
+  } as unknown as StateDecorator<T>;
 }
 
 function getStateOptions<T>(

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

@@ -182,6 +182,14 @@ export class StateFactory implements OnDestroy {
       ensureStatesAreDecorated(stateClasses);
     }
 
+    // Just playing around with this because the metadata is not
+    // set until the constructor function is called (if `State` is used
+    // as a factory function rather than as a decorator, which returns
+    // a base state).
+    for (const stateClass of stateClasses) {
+      this._injector.get(stateClass);
+    }
+
     const { newStates } = this.addToStatesMap(stateClasses);
     if (!newStates.length) return [];
 
@@ -295,7 +303,7 @@ export class StateFactory implements OnDestroy {
 
       if (actionMetas) {
         for (const actionMeta of actionMetas) {
-          const stateContext = this._stateContextFactory.createStateContext(metadata);
+          const stateContext = this._stateContextFactory.createStateContext(metadata.path);
           try {
             let result = metadata.instance[actionMeta.fn](stateContext, action);
 

packages/store/tests/state-extends.spec.ts

@@ -0,0 +1,40 @@
+import { Component, Injectable } from '@angular/core';
+import { bootstrapApplication } from '@angular/platform-browser';
+import { Action, InitState, NgxsOnInit, provideStore, State, Store } from '@ngxs/store';
+import { freshPlatform } from '@ngxs/store/internals/testing';
+
+describe('State decorator returns a constructor function', () => {
+  @Component({ selector: 'app-root', template: '', standalone: true })
+  class TestComponent {}
+
+  it(
+    'should call an InitState action handler before the ngxsOnInit method on root module initialisation',
+    freshPlatform(async () => {
+      // Arrange
+      @Injectable()
+      class FooState
+        extends State<string[]>({ name: 'foo', defaults: [] })
+        implements NgxsOnInit
+      {
+        ngxsOnInit() {
+          this.setState([...this.getState(), 'onInit']);
+        }
+
+        @Action(InitState)
+        initState() {
+          this.setState([...this.getState(), 'initState']);
+        }
+      }
+
+      // Act
+      const { injector } = await bootstrapApplication(TestComponent, {
+        providers: [provideStore([FooState])]
+      });
+
+      const store = injector.get(Store);
+
+      // Assert
+      expect(store.snapshot().foo).toEqual(['initState', 'onInit']);
+    })
+  );
+});