stash

Author: MajorLift

packages/composable-controller/CHANGELOG.md

@@ -15,7 +15,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - **BREAKING:** Passing a non-controller into `controllers` constructor option now throws an error ([#3904](https://github.com/MetaMask/core/pull/3904))
 - **BREAKING:** The `AllowedActions` parameter of the `ComposableControllerMessenger` type is narrowed from `string` to `never`, as `ComposableController` does not use any external controller actions. ([#3904](https://github.com/MetaMask/core/pull/3904))
-- Subscribe to the `stateChange` events of `BaseControllerV1` controllers that have a `messagingSystem` class field with an assigned instance of the `RestrictedControllerMessenger` class. ([#3964](https://github.com/MetaMask/core/pull/3964))
+- **BREAKING:** The `ComposableController` class is now a generic class that expects one generic argument `ChildControllers` in the form of a union of the controllers passed into the `controllers` array constructor option ([#3941](https://github.com/MetaMask/core/pull/3941))
+  - Child controllers that extend `BaseControllerV1` must have an overridden `name` property that is defined using the `as const` assertion for the `ComposableController` class to be typed correctly.
+  - The controller state is now constrained by a `ComposedControllerState` generic argument, which is automaticallly derived from the `ChildControllers` argument without needing to be passed in by the user.
+  - The `messenger` constructor option is constrained by a `ComposedControllerMessenger` generic argument which is also automatically derived, and is required to contain the `stateChange` events for all child controllers in its `Events` and `AllowedEvents` parameters, as well as the composable controller instance in its `Events` parameter.
 
 ### Removed
 

packages/composable-controller/src/ComposableController.test.ts

@@ -60,6 +60,50 @@ class FooController extends BaseController<
   }
 }
 
+type QuzControllerState = {
+  quz: string;
+};
+type QuzControllerEvent = {
+  type: `QuzController:stateChange`;
+  payload: [QuzControllerState, Patch[]];
+};
+
+type QuzMessenger = RestrictedControllerMessenger<
+  'QuzController',
+  never,
+  QuzControllerEvent,
+  never,
+  never
+>;
+
+const quzControllerStateMetadata = {
+  quz: {
+    persist: true,
+    anonymous: true,
+  },
+};
+
+class QuzController extends BaseController<
+  'QuzController',
+  QuzControllerState,
+  QuzMessenger
+> {
+  constructor(messagingSystem: QuzMessenger) {
+    super({
+      messenger: messagingSystem,
+      metadata: quzControllerStateMetadata,
+      name: 'QuzController',
+      state: { quz: 'quz' },
+    });
+  }
+
+  updateQuz(quz: string) {
+    super.update((state) => {
+      state.quz = quz;
+    });
+  }
+}
+
 // Mock BaseControllerV1 classes
 
 type BarControllerState = BaseState & {
@@ -71,7 +115,7 @@ class BarController extends BaseControllerV1<never, BarControllerState> {
     bar: 'bar',
   };
 
-  override name = 'BarController';
+  override name = 'BarController' as const;
 
   constructor() {
     super();
@@ -92,7 +136,7 @@ class BazController extends BaseControllerV1<never, BazControllerState> {
     baz: 'baz',
   };
 
-  override name = 'BazController';
+  override name = 'BazController' as const;
 
   constructor() {
     super();
@@ -157,23 +201,39 @@ describe('ComposableController', () => {
     it('should compose controller state', () => {
       const controllerMessenger = new ControllerMessenger<
         never,
-        FooControllerEvent
+        FooControllerEvent | QuzControllerEvent
       >();
-      const fooControllerMessenger = controllerMessenger.getRestricted({
+      const fooMessenger = controllerMessenger.getRestricted<
+        'FooController',
+        never,
+        never
+      >({
         name: 'FooController',
       });
-      const fooController = new FooController(fooControllerMessenger);
+      const quzMessenger = controllerMessenger.getRestricted<
+        'QuzController',
+        never,
+        never
+      >({
+        name: 'QuzController',
+      });
+      const fooController = new FooController(fooMessenger);
+      const quzController = new QuzController(quzMessenger);
 
       const composableControllerMessenger = controllerMessenger.getRestricted({
         name: 'ComposableController',
-        allowedEvents: ['FooController:stateChange'],
+        allowedEvents: [
+          'FooController:stateChange',
+          'QuzController:stateChange',
+        ],
       });
       const composableController = new ComposableController({
-        controllers: [fooController],
+        controllers: [fooController, quzController],
         messenger: composableControllerMessenger,
       });
       expect(composableController.state).toStrictEqual({
         FooController: { foo: 'foo' },
+        QuzController: { quz: 'quz' },
       });
     });
 

packages/composable-controller/src/ComposableController.ts

@@ -3,8 +3,10 @@ import type {
   BaseState,
   BaseConfig,
   RestrictedControllerMessenger,
+  EventConstraint,
   StateMetadata,
   StateConstraint,
+  ControllerStateChangeEvent,
 } from '@metamask/base-controller';
 import type { Patch } from 'immer';
 
@@ -113,13 +115,44 @@ export type ComposableControllerMessenger = RestrictedControllerMessenger<
   AllowedEvents['type']
 >;
 
+type GetStateChangeEvents<Controller extends ControllerInstance> =
+  Controller extends ControllerInstance
+    ? Controller extends BaseControllerV1Instance
+      ? {
+          type: `${Controller['name']}:stateChange`;
+          payload: [Controller['state'], Patch[]];
+        }
+      : ControllerStateChangeEvent<Controller['name'], Controller['state']>
+    : never;
+
 /**
  * Controller that can be used to compose multiple controllers together.
  */
-export class ComposableController extends BaseController<
+export class ComposableController<
+  ChildControllers extends ControllerInstance,
+  ComposedControllerState extends ComposableControllerState = {
+    [P in ChildControllers as P['name']]: P['state'];
+  },
+  ComposedControllerStateChangeEvent extends EventConstraint & {
+    type: `${typeof controllerName}:stateChange`;
+  } = ControllerStateChangeEvent<
+    typeof controllerName,
+    ComposedControllerState
+  >,
+  ChildControllersStateChangeEvents extends EventConstraint & {
+    type: `${string}:stateChange`;
+  } = GetStateChangeEvents<ChildControllers>,
+  ComposedControllerMessenger extends ComposableControllerMessenger = RestrictedControllerMessenger<
+    typeof controllerName,
+    never,
+    ComposedControllerStateChangeEvent | ChildControllersStateChangeEvents,
+    never,
+    ChildControllersStateChangeEvents['type']
+  >,
+> extends BaseController<
   typeof controllerName,
-  ComposableControllerState,
-  ComposableControllerMessenger
+  ComposedControllerState,
+  ComposedControllerMessenger
 > {
   /**
    * Creates a ComposableController instance.
@@ -133,29 +166,29 @@ export class ComposableController extends BaseController<
     controllers,
     messenger,
   }: {
-    controllers: ControllerInstance[];
-    messenger: ComposableControllerMessenger;
+    controllers: ChildControllers[];
+    messenger: ComposedControllerMessenger;
   }) {
     if (messenger === undefined) {
       throw new Error(`Messaging system is required`);
     }
 
     super({
       name: controllerName,
-      metadata: controllers.reduce<StateMetadata<ComposableControllerState>>(
+      metadata: controllers.reduce<StateMetadata<ComposedControllerState>>(
         (metadata, controller) => ({
           ...metadata,
           [controller.name]: isBaseController(controller)
             ? controller.metadata
             : { persist: true, anonymous: true },
         }),
-        {},
+        {} as never,
       ),
-      state: controllers.reduce<ComposableControllerState>(
+      state: controllers.reduce<ComposedControllerState>(
         (state, controller) => {
           return { ...state, [controller.name]: controller.state };
         },
-        {},
+        {} as never,
       ),
       messenger,
     });
@@ -179,10 +212,11 @@ export class ComposableController extends BaseController<
     const { name } = controller;
     if (isBaseControllerV1(controller)) {
       controller.subscribe((childState) => {
-        this.update((state) => ({
-          ...state,
-          [name]: childState,
-        }));
+        this.update((state) => {
+          Object.assign(state, {
+            [name]: childState,
+          });
+        });
       });
     }
     if (
@@ -191,11 +225,12 @@ export class ComposableController extends BaseController<
     ) {
       this.messagingSystem.subscribe(
         `${name}:stateChange`,
-        (childState: StateConstraint) => {
-          this.update((state) => ({
-            ...state,
-            [name]: childState,
-          }));
+        (childState: Record<string, unknown>) => {
+          this.update((state) => {
+            Object.assign(state, {
+              [name]: childState,
+            });
+          });
         },
       );
     }