feat(account-tree-controller): add payload/snapshot support

packages/account-tree-controller/src/state/id-map.test.ts

@@ -0,0 +1,104 @@
+import { IdMap } from './id-map.js';
+
+describe('IdMap', () => {
+  describe('constructor', () => {
+    it('creates an empty map when called with no arguments', () => {
+      const map = new IdMap();
+      expect(map.getPayloadId('entropy:wallet-1')).toBeUndefined();
+      expect(map.getLocalId('wallet:entropy-source-1')).toBeUndefined();
+    });
+
+    it('pre-populates the map from the provided entries', () => {
+      const map = new IdMap([
+        ['entropy:wallet-1', 'wallet:entropy-source-1'],
+        ['keyring:wallet-2', 'wallet:private-key'],
+      ]);
+      expect(map.getPayloadId('entropy:wallet-1')).toBe(
+        'wallet:entropy-source-1',
+      );
+      expect(map.getPayloadId('keyring:wallet-2')).toBe('wallet:private-key');
+      expect(map.getLocalId('wallet:entropy-source-1')).toBe(
+        'entropy:wallet-1',
+      );
+      expect(map.getLocalId('wallet:private-key')).toBe('keyring:wallet-2');
+    });
+  });
+
+  describe('add', () => {
+    it('registers a local-to-payload pair and its reverse', () => {
+      const map = new IdMap();
+      map.add('entropy:wallet-1', 'wallet:entropy-source-1');
+      expect(map.getPayloadId('entropy:wallet-1')).toBe(
+        'wallet:entropy-source-1',
+      );
+      expect(map.getLocalId('wallet:entropy-source-1')).toBe(
+        'entropy:wallet-1',
+      );
+    });
+
+    it('throws if the same local ID is registered twice', () => {
+      const map = new IdMap();
+      map.add('entropy:wallet-1', 'wallet:entropy-source-1');
+      expect(() =>
+        map.add('entropy:wallet-1', 'wallet:entropy-source-2'),
+      ).toThrow('Local ID already registered: entropy:wallet-1');
+    });
+
+    it('throws if the same payload ID is registered twice', () => {
+      const map = new IdMap();
+      map.add('entropy:wallet-1', 'wallet:entropy-source-1');
+      expect(() =>
+        map.add('entropy:wallet-2', 'wallet:entropy-source-1'),
+      ).toThrow('Payload ID already registered: wallet:entropy-source-1');
+    });
+
+    it('handles wallet and group IDs in the same map', () => {
+      const map = new IdMap();
+      map.add('entropy:wallet-1', 'wallet:entropy-source-1');
+      map.add('entropy:wallet-1/0', 'wallet:entropy-source-1/0');
+      expect(map.getPayloadId('entropy:wallet-1')).toBe(
+        'wallet:entropy-source-1',
+      );
+      expect(map.getPayloadId('entropy:wallet-1/0')).toBe(
+        'wallet:entropy-source-1/0',
+      );
+    });
+  });
+
+  describe('getPayloadId', () => {
+    it('returns the payload ID for a known local wallet ID', () => {
+      const map = new IdMap([['entropy:wallet-1', 'wallet:entropy-source-1']]);
+      expect(map.getPayloadId('entropy:wallet-1')).toBe(
+        'wallet:entropy-source-1',
+      );
+    });
+
+    it('returns undefined for an unknown local ID', () => {
+      const map = new IdMap();
+      expect(map.getPayloadId('entropy:wallet-unknown')).toBeUndefined();
+    });
+  });
+
+  describe('getLocalId', () => {
+    it('returns the local ID for a known payload wallet ID', () => {
+      const map = new IdMap([['entropy:wallet-1', 'wallet:entropy-source-1']]);
+      expect(map.getLocalId('wallet:entropy-source-1')).toBe(
+        'entropy:wallet-1',
+      );
+    });
+
+    it('returns undefined for an unknown payload ID', () => {
+      const map = new IdMap();
+      expect(map.getLocalId('wallet:entropy-source-unknown')).toBeUndefined();
+    });
+
+    it('returns the local ID for a group payload ID', () => {
+      const map = new IdMap([
+        ['entropy:wallet-1/0', 'wallet:entropy-source-1/0'],
+      ]);
+      expect(map.getLocalId('wallet:entropy-source-1/0')).toBe(
+        'entropy:wallet-1/0',
+      );
+    });
+  });
+});

packages/account-tree-controller/src/state/id-map.ts

@@ -0,0 +1,74 @@
+import type { AccountGroupId, AccountWalletId } from '@metamask/account-api';
+import { assert } from '@metamask/utils';
+
+import type {
+  AccountGroupPayloadId,
+  AccountWalletPayloadId,
+} from './payload.js';
+
+type LocalId = AccountWalletId | AccountGroupId;
+type PayloadId = AccountWalletPayloadId | AccountGroupPayloadId;
+
+/**
+ * Bidirectional map between local controller IDs and portable payload IDs.
+ *
+ * Populated during {@link exportState} so that callers can bridge between
+ * device-local wallet/group IDs and the stable cross-device IDs that appear
+ * in a serialized {@link AccountTreePayload}.
+ */
+export class IdMap {
+  readonly #localToPayload: Map<LocalId, PayloadId> = new Map();
+
+  readonly #payloadToLocal: Map<PayloadId, LocalId> = new Map();
+
+  /**
+   * @param entries - Optional seed pairs `[localId, payloadId]` to pre-populate the map.
+   */
+  constructor(entries: [localId: LocalId, payloadId: PayloadId][] = []) {
+    for (const [localId, payloadId] of entries) {
+      this.add(localId, payloadId);
+    }
+  }
+
+  /**
+   * Registers a local↔payload ID pair.
+   *
+   * @param localId - Local controller wallet or group ID.
+   * @param payloadId - Corresponding portable payload ID.
+   */
+  add(localId: LocalId, payloadId: PayloadId): void {
+    // Each local ID maps to exactly one payload ID and vice versa: wallet and
+    // group IDs are unique by construction, so a duplicate signals a bug in the
+    // caller (e.g. exportState visiting the same node twice).
+    assert(
+      !this.#localToPayload.has(localId),
+      `Local ID already registered: ${localId}`,
+    );
+    assert(
+      !this.#payloadToLocal.has(payloadId),
+      `Payload ID already registered: ${payloadId}`,
+    );
+    this.#localToPayload.set(localId, payloadId);
+    this.#payloadToLocal.set(payloadId, localId);
+  }
+
+  /**
+   * Returns the portable payload ID for a given local ID.
+   *
+   * @param localId - Local controller wallet or group ID.
+   * @returns The payload ID, or `undefined` if not registered.
+   */
+  getPayloadId(localId: LocalId): PayloadId | undefined {
+    return this.#localToPayload.get(localId);
+  }
+
+  /**
+   * Returns the local controller ID for a given payload ID.
+   *
+   * @param payloadId - Portable payload ID.
+   * @returns The local wallet or group ID, or `undefined` if not registered.
+   */
+  getLocalId(payloadId: PayloadId): LocalId | undefined {
+    return this.#payloadToLocal.get(payloadId);
+  }
+}