refactor: Add mailbox types and differentiate use of maps from map-like stores

packages/SwingSet/src/devices/mailbox/mailbox.js

@@ -71,17 +71,39 @@ import { Nat } from '@endo/nat';
 // replace it with one that tracks which parts of the state have been
 // modified, to build more efficient Merkle proofs.
 
+/**
+ * @typedef {object} Mailbox
+ * @property {bigint} ack
+ * @property {Map<bigint, unknown>} outbox
+ */
+
+/**
+ * @typedef {object} MailboxExport
+ * @property {number} ack
+ * @property {Array<[number, unknown]>} outbox
+ */
+
+/**
+ * @param {MailboxExport} data
+ * @param {Partial<Mailbox>} [inout]
+ * @returns {Mailbox}
+ */
 export function importMailbox(data, inout = {}) {
   const outbox = new Map();
   for (const m of data.outbox) {
     outbox.set(Nat(m[0]), m[1]);
   }
   inout.ack = Nat(data.ack);
   inout.outbox = outbox;
-  return inout;
+  return /** @type {Mailbox} */ (inout);
 }
 
+/**
+ * @param {Mailbox} inout
+ * @returns {MailboxExport}
+ */
 export function exportMailbox(inout) {
+  /** @type {MailboxExport['outbox']} */
   const messages = [];
   for (const [msgnum, body] of inout.outbox) {
     messages.push([Number(msgnum), body]);
@@ -93,64 +115,75 @@ export function exportMailbox(inout) {
   };
 }
 
-export function buildMailboxStateMap(state = harden(new Map())) {
-  function getOrCreatePeer(peer) {
-    if (!state.has(peer)) {
-      const inout = {
-        outbox: harden(new Map()),
-        ack: 0n,
-      };
-      state.set(peer, inout);
-    }
-    return state.get(peer);
+/**
+ * @param {Map<string, Mailbox>} state
+ */
+export function exportMailboxData(state) {
+  /** @type {Record<string, {inboundAck: MailboxExport['ack'], outbox: MailboxExport['outbox']}>} */
+  const data = {};
+  for (const [peer, inout] of state.entries()) {
+    const exported = exportMailbox(inout);
+    data[peer] = {
+      inboundAck: exported.ack,
+      outbox: exported.outbox,
+    };
   }
+  return harden(data);
+}
 
-  function add(peer, msgnum, body) {
-    getOrCreatePeer(`${peer}`).outbox.set(Nat(msgnum), `${body}`);
+function getOrCreatePeer(state, peer) {
+  if (!state.has(peer)) {
+    const inout = {
+      outbox: harden(new Map()),
+      ack: 0n,
+    };
+    state.set(peer, inout);
   }
+  return state.get(peer);
+}
 
-  function remove(peer, msgnum) {
-    const messages = getOrCreatePeer(`${peer}`).outbox;
-    messages.delete(Nat(msgnum));
+/**
+ * @param {ReturnType<exportMailboxData>} data
+ * @returns {Map<string, Mailbox>}
+ */
+export function makeEphemeralMailboxStorage(data) {
+  const state = harden(new Map());
+  for (const peer of Object.getOwnPropertyNames(data)) {
+    const inout = getOrCreatePeer(state, peer);
+    const d = data[peer];
+    importMailbox(
+      {
+        ack: d.inboundAck,
+        outbox: d.outbox,
+      },
+      inout,
+    );
   }
+  return state;
+}
 
-  function setAcknum(peer, msgnum) {
-    getOrCreatePeer(`${peer}`).ack = Nat(msgnum);
+/**
+ * @template [T=unknown]
+ * @param {Pick<Map<string, T>, 'has' | 'get'> & {set: (key: string, value: T) => void}} [state]
+ */
+export function buildMailboxStateMap(state = harden(new Map())) {
+  function add(peer, msgnum, body) {
+    getOrCreatePeer(state, `${peer}`).outbox.set(Nat(msgnum), `${body}`);
   }
 
-  function exportToData() {
-    const data = {};
-    for (const [peer, inout] of state.entries()) {
-      const exported = exportMailbox(inout);
-      data[peer] = {
-        inboundAck: exported.ack,
-        outbox: exported.outbox,
-      };
-    }
-    return harden(data);
+  function remove(peer, msgnum) {
+    const messages = getOrCreatePeer(state, `${peer}`).outbox;
+    messages.delete(Nat(msgnum));
   }
 
-  function populateFromData(data) {
-    !state.size || Fail`cannot populateFromData: outbox is not empty`;
-    for (const peer of Object.getOwnPropertyNames(data)) {
-      const inout = getOrCreatePeer(peer);
-      const d = data[peer];
-      importMailbox(
-        {
-          ack: d.inboundAck,
-          outbox: d.outbox,
-        },
-        inout,
-      );
-    }
+  function setAcknum(peer, msgnum) {
+    getOrCreatePeer(state, `${peer}`).ack = Nat(msgnum);
   }
 
   return harden({
     add,
     remove,
     setAcknum,
-    exportToData,
-    populateFromData,
   });
 }
 

packages/SwingSet/src/index.js

@@ -13,6 +13,8 @@ export { upgradeSwingset } from './controller/upgradeSwingset.js';
 export {
   buildMailboxStateMap,
   buildMailbox,
+  exportMailboxData,
+  makeEphemeralMailboxStorage,
 } from './devices/mailbox/mailbox.js';
 export { buildTimer } from './devices/timer/timer.js';
 export { buildBridge } from './devices/bridge/bridge.js';

packages/SwingSet/src/types-external.js

@@ -283,6 +283,13 @@ export {};
  * @property { boolean } [restartWorkerOnSnapshot]     Reload worker immediately upon snapshot creation
  */
 
+/**
+ * @typedef { import('./devices/mailbox/mailbox.js').Mailbox } Mailbox
+ */
+/**
+ * @typedef { import('./devices/mailbox/mailbox.js').MailboxExport } MailboxExport
+ */
+
 /**
  * Vat Creation and Management
  *

packages/SwingSet/test/device-mailbox/device-mailbox.test.js

@@ -13,6 +13,8 @@ import {
 import {
   buildMailboxStateMap,
   buildMailbox,
+  exportMailboxData,
+  makeEphemeralMailboxStorage,
 } from '../../src/devices/mailbox/mailbox.js';
 import { bundleOpts } from '../util.js';
 
@@ -24,7 +26,8 @@ test.before(async t => {
 });
 
 test('mailbox outbound', async t => {
-  const s = buildMailboxStateMap();
+  const mailboxStorage = harden(new Map());
+  const s = buildMailboxStateMap(mailboxStorage);
   const mb = buildMailbox(s);
   const config = {
     bootstrap: 'bootstrap',
@@ -49,8 +52,8 @@ test('mailbox outbound', async t => {
   const c = await makeSwingsetController(kernelStorage, devEndows, runtimeOpts);
   t.teardown(c.shutdown);
   await c.run();
-  // exportToData() provides plain Numbers to the host that needs to convey the messages
-  t.deepEqual(s.exportToData(), {
+  // exportMailboxData() provides plain Numbers to the host that needs to convey the messages
+  t.deepEqual(exportMailboxData(mailboxStorage), {
     peer1: {
       inboundAck: 13,
       outbox: [
@@ -68,9 +71,9 @@ test('mailbox outbound', async t => {
     },
   });
 
-  const s2 = buildMailboxStateMap();
-  s2.populateFromData(s.exportToData());
-  t.deepEqual(s.exportToData(), s2.exportToData());
+  const mailboxDataExport = exportMailboxData(mailboxStorage);
+  const mailboxStorageCopy = makeEphemeralMailboxStorage(mailboxDataExport);
+  t.deepEqual(exportMailboxData(mailboxStorageCopy), mailboxDataExport);
 });
 
 test('mailbox inbound', async t => {

packages/SwingSet/test/vattp.test.js

@@ -7,6 +7,7 @@ import { initializeSwingset, makeSwingsetController } from '../src/index.js';
 import {
   buildMailboxStateMap,
   buildMailbox,
+  exportMailboxData,
 } from '../src/devices/mailbox/mailbox.js';
 
 async function restartVatTP(controller) {
@@ -19,7 +20,8 @@ async function restartVatTP(controller) {
 }
 
 test.serial('vattp', async t => {
-  const s = buildMailboxStateMap();
+  const mailboxStorage = harden(new Map());
+  const s = buildMailboxStateMap(mailboxStorage);
   const mb = buildMailbox(s);
   const config = {
     bootstrap: 'bootstrap',
@@ -46,7 +48,7 @@ test.serial('vattp', async t => {
   const c = await makeSwingsetController(kernelStorage, deviceEndowments);
   t.teardown(c.shutdown);
   await c.run();
-  t.deepEqual(s.exportToData(), {});
+  t.deepEqual(exportMailboxData(mailboxStorage), {});
 
   const m1 = [1, 'msg1'];
   const m2 = [2, 'msg2'];
@@ -57,15 +59,20 @@ test.serial('vattp', async t => {
     'ch.receive msg1',
     'ch.receive msg2',
   ]);
-  t.deepEqual(s.exportToData(), { remote1: { outbox: [], inboundAck: 2 } });
+  t.deepEqual(exportMailboxData(mailboxStorage), {
+    remote1: { outbox: [], inboundAck: 2 },
+  });
 
   t.is(mb.deliverInbound('remote1', [m1, m2], 0), true);
   await c.run();
-  t.deepEqual(s.exportToData(), { remote1: { outbox: [], inboundAck: 2 } });
+  t.deepEqual(exportMailboxData(mailboxStorage), {
+    remote1: { outbox: [], inboundAck: 2 },
+  });
 });
 
 test.serial('vattp 2', async t => {
-  const s = buildMailboxStateMap();
+  const mailboxStorage = harden(new Map());
+  const s = buildMailboxStateMap(mailboxStorage);
   const mb = buildMailbox(s);
   const config = {
     bootstrap: 'bootstrap',
@@ -93,19 +100,23 @@ test.serial('vattp 2', async t => {
   t.teardown(c.shutdown);
   c.pinVatRoot('bootstrap');
   await c.run();
-  t.deepEqual(s.exportToData(), {
+  t.deepEqual(exportMailboxData(mailboxStorage), {
     remote1: { outbox: [[1, 'out1']], inboundAck: 0 },
   });
 
   t.is(mb.deliverInbound('remote1', [], 1), true);
   await c.run();
   t.deepEqual(c.dump().log, []);
-  t.deepEqual(s.exportToData(), { remote1: { outbox: [], inboundAck: 0 } });
+  t.deepEqual(exportMailboxData(mailboxStorage), {
+    remote1: { outbox: [], inboundAck: 0 },
+  });
 
   t.is(mb.deliverInbound('remote1', [[1, 'msg1']], 1), true);
   await c.run();
   t.deepEqual(c.dump().log, ['ch.receive msg1']);
-  t.deepEqual(s.exportToData(), { remote1: { outbox: [], inboundAck: 1 } });
+  t.deepEqual(exportMailboxData(mailboxStorage), {
+    remote1: { outbox: [], inboundAck: 1 },
+  });
 
   t.is(mb.deliverInbound('remote1', [[1, 'msg1']], 1), true);
 
@@ -114,7 +125,9 @@ test.serial('vattp 2', async t => {
   t.is(mb.deliverInbound('remote1', [m1, m2], 1), true);
   await c.run();
   t.deepEqual(c.dump().log, ['ch.receive msg1', 'ch.receive msg2']);
-  t.deepEqual(s.exportToData(), { remote1: { outbox: [], inboundAck: 2 } });
+  t.deepEqual(exportMailboxData(mailboxStorage), {
+    remote1: { outbox: [], inboundAck: 2 },
+  });
 
   // now upgrade vattp, and see if the state is retained
   await restartVatTP(c);
@@ -123,7 +136,9 @@ test.serial('vattp 2', async t => {
   t.is(mb.deliverInbound('remote1', [m2], 1), true);
   t.deepEqual(c.dump().log, ['ch.receive msg1', 'ch.receive msg2']);
   // vattp remembers the inboundAck number
-  t.deepEqual(s.exportToData(), { remote1: { outbox: [], inboundAck: 2 } });
+  t.deepEqual(exportMailboxData(mailboxStorage), {
+    remote1: { outbox: [], inboundAck: 2 },
+  });
 
   // vattp remembers the receiver: new inbound msgs are still delivered
   const m3 = [3, 'msg3'];
@@ -134,13 +149,15 @@ test.serial('vattp 2', async t => {
     'ch.receive msg2',
     'ch.receive msg3',
   ]);
-  t.deepEqual(s.exportToData(), { remote1: { outbox: [], inboundAck: 3 } });
+  t.deepEqual(exportMailboxData(mailboxStorage), {
+    remote1: { outbox: [], inboundAck: 3 },
+  });
 
   // vattp still supports the transmitter object
   c.queueToVatRoot('bootstrap', 'transmit', ['out2']);
   await c.run();
   // vattp remembers the outbound seqnum
-  t.deepEqual(s.exportToData(), {
+  t.deepEqual(exportMailboxData(mailboxStorage), {
     remote1: { outbox: [[2, 'out2']], inboundAck: 3 },
   });
 });

packages/boot/tools/supports.ts

@@ -511,8 +511,10 @@ export const makeSwingsetTestKit = async (
       },
     });
   }
+  const mailboxStorage = new Map();
   const { controller, timer, bridgeInbound } = await buildSwingset(
-    new Map(),
+    // @ts-expect-error missing method 'getNextKey'
+    mailboxStorage,
     bridgeOutbound,
     kernelStorage,
     configPath,

packages/cosmic-swingset/src/launch-chain.js

@@ -52,7 +52,8 @@ import { makeQueue, makeQueueStorageMock } from './helpers/make-queue.js';
 import { exportStorage } from './export-storage.js';
 import { parseLocatedJson } from './helpers/json.js';
 
-/** @import {RunPolicy} from '@agoric/swingset-vat' */
+/** @import { Mailbox, RunPolicy, SwingSetConfig } from '@agoric/swingset-vat' */
+/** @import { KVStore } from './helpers/bufferedStorage.js' */
 
 const console = anylogger('launch-chain');
 const blockManagerConsole = anylogger('block-manager');
@@ -74,8 +75,6 @@ const parseUpgradePlanInfo = (upgradePlan, prefix = '') => {
   return harden(upgradePlanInfo || {});
 };
 
-/** @import {SwingSetConfig} from '@agoric/swingset-vat' */
-
 /**
  * @typedef {object} CosmicSwingsetConfig
  * @property {import('@agoric/deploy-script-support/src/extract-proposal.js').ConfigProposal[]} [coreProposals]
@@ -96,7 +95,7 @@ const parseUpgradePlanInfo = (upgradePlan, prefix = '') => {
 const getHostKey = path => `host.${path}`;
 
 /**
- * @param {Map<*, *>} mailboxStorage
+ * @param {KVStore<Mailbox>} mailboxStorage
  * @param {((dstID: string, obj: any) => any)} bridgeOutbound
  * @param {SwingStoreKernelStorage} kernelStorage
  * @param {string | (() => string | Promise<string>)} vatconfig absolute path or thunk