[@xstate/store] Add store.with(…) for store extensions (#5427)

* Add store.with(…) extension, deprecate inline undoRedo

* Changeset

* Use event payload map more consistently, fix go-to-def for undo/redo

Author: davidkpiano

.changeset/odd-wombats-thank.md

@@ -0,0 +1,24 @@
+---
+'@xstate/store': minor
+---
+
+Add `.with()` method for store extensions.
+
+```ts
+import { createStore } from '@xstate/store';
+import { undoRedo } from '@xstate/store/undo';
+
+const store = createStore({
+  context: { count: 0 },
+  on: {
+    inc: (ctx) => ({ count: ctx.count + 1 }),
+    dec: (ctx) => ({ count: ctx.count - 1 })
+  }
+}).with(undoRedo());
+
+store.trigger.inc(); // count = 1
+
+// Added from the undoRedo extension
+store.trigger.undo(); // count = 0
+store.trigger.redo(); // count = 1
+```

packages/xstate-store/src/index.ts

@@ -37,5 +37,6 @@ export type {
   ReadonlyAtom,
   EventFromStoreConfig,
   EmitsFromStoreConfig,
-  ContextFromStoreConfig
+  ContextFromStoreConfig,
+  StoreExtension
 } from './types';

packages/xstate-store/src/store.ts

@@ -186,6 +186,10 @@ function createStoreCore<
       return createAtom(() => selector(store.get().context), {
         compare: equalityFn
       });
+    },
+    with(extension) {
+      const extendedLogic = extension(logic as any);
+      return createStoreCore(extendedLogic) as any;
     }
   };
 

packages/xstate-store/src/types.ts

@@ -146,6 +146,29 @@ export interface Store<
     ExtractEvents<TEventPayloadMap>,
     TEmitted
   >;
+  /**
+   * Extends the store with additional functionality via a store extension.
+   *
+   * @example
+   *
+   * ```ts
+   * const store = createStore({
+   *   context: { count: 0 },
+   *   on: { inc: (ctx) => ({ count: ctx.count + 1 }) }
+   * }).with(undoRedo());
+   *
+   * store.trigger.inc();
+   * store.trigger.undo(); // undoes the increment
+   * ```
+   */
+  with<TNewEventPayloadMap extends EventPayloadMap>(
+    extension: StoreExtension<
+      TContext,
+      TEventPayloadMap,
+      TNewEventPayloadMap,
+      TEmitted
+    >
+  ): Store<TContext, TEventPayloadMap & TNewEventPayloadMap, TEmitted>;
 }
 
 export type StoreTransition<
@@ -443,6 +466,35 @@ export type StoreLogic<
 };
 export type AnyStoreLogic = StoreLogic<any, any, any>;
 
+/**
+ * A store extension that transforms store logic, optionally adding new events.
+ *
+ * @example
+ *
+ * ```ts
+ * const store = createStore({
+ *   context: { count: 0 },
+ *   on: { inc: (ctx) => ({ count: ctx.count + 1 }) }
+ * }).with(undoRedo());
+ * ```
+ */
+export type StoreExtension<
+  TContext extends StoreContext,
+  TEventPayloadMap extends EventPayloadMap,
+  TNewEventPayloadMap extends EventPayloadMap,
+  TEmitted extends EventObject
+> = (
+  logic: StoreLogic<
+    StoreSnapshot<TContext>,
+    ExtractEvents<TEventPayloadMap>,
+    TEmitted
+  >
+) => StoreLogic<
+  StoreSnapshot<TContext>,
+  ExtractEvents<TEventPayloadMap> | ExtractEvents<TNewEventPayloadMap>,
+  TEmitted
+>;
+
 export type AnyStoreConfig = StoreConfig<any, any, any>;
 export type EventFromStoreConfig<TStore extends AnyStoreConfig> =
   TStore extends StoreConfig<any, infer TEventPayloadMap, any>