envio 3.5.1 → 3.6.1

package/src/EntityTables.res

@@ -0,0 +1,32 @@
+// One in-memory table per entity, for a single chain scope. The indexer holds
+// the cross-chain partition; each ChainState holds the per-chain one.
+
+type t = dict<InMemoryTable.Entity.t>
+
+exception UndefinedEntity({entityName: string})
+
+let make = (entities: array<Internal.entityConfig>): t => {
+  let init = Dict.make()
+  entities->Array.forEach(entityConfig => {
+    init->Dict.set((entityConfig.name :> string), InMemoryTable.Entity.make())
+  })
+  init
+}
+
+let get = (self: t, ~entityName: string) => {
+  switch self->Utils.Dict.dangerouslyGetNonOption(entityName) {
+  | Some(table) => table
+  | None =>
+    UndefinedEntity({entityName: entityName})->ErrorHandling.mkLogAndRaise(
+      ~msg="Unexpected, entity InMemoryTable is undefined",
+    )
+  }
+}
+
+// Entities whose rows are shared by every chain, so their tables live on the
+// indexer rather than on a ChainState.
+let crossChain = (entities: array<Internal.entityConfig>) =>
+  entities->Array.filter((entityConfig: Internal.entityConfig) => entityConfig.crossChain)
+
+let perChain = (entities: array<Internal.entityConfig>) =>
+  entities->Array.filter((entityConfig: Internal.entityConfig) => !entityConfig.crossChain)

package/src/EntityTables.res.mjs

@@ -0,0 +1,44 @@
+// Generated by ReScript, PLEASE EDIT WITH CARE
+
+import * as ErrorHandling from "./ErrorHandling.res.mjs";
+import * as InMemoryTable from "./InMemoryTable.res.mjs";
+import * as Primitive_exceptions from "@rescript/runtime/lib/es6/Primitive_exceptions.js";
+
+let UndefinedEntity = /* @__PURE__ */Primitive_exceptions.create("EntityTables.UndefinedEntity");
+
+function make(entities) {
+  let init = {};
+  entities.forEach(entityConfig => {
+    init[entityConfig.name] = InMemoryTable.Entity.make();
+  });
+  return init;
+}
+
+function get(self, entityName) {
+  let table = self[entityName];
+  if (table !== undefined) {
+    return table;
+  } else {
+    return ErrorHandling.mkLogAndRaise(undefined, "Unexpected, entity InMemoryTable is undefined", {
+      RE_EXN_ID: UndefinedEntity,
+      entityName: entityName
+    });
+  }
+}
+
+function crossChain(entities) {
+  return entities.filter(entityConfig => entityConfig.crossChain);
+}
+
+function perChain(entities) {
+  return entities.filter(entityConfig => !entityConfig.crossChain);
+}
+
+export {
+  UndefinedEntity,
+  make,
+  get,
+  crossChain,
+  perChain,
+}
+/* ErrorHandling Not a pure module */

package/src/Envio.res

@@ -186,7 +186,8 @@ and effectOptions<'input, 'output> = {
   rateLimit: rateLimit,
   /** Whether the effect should be cached. */
   cache?: bool,
-  /** Whether the effect's cache is shared across all chains. Defaults to `true`.
+  /** Whether the effect's cache is shared across all chains. Defaults to `true`,
+   or to `false` when config.yaml sets `disable_default_cross_chain: true`.
    Set to `false` to isolate the cache (and rate limiting) per chain and enable
    `context.chain.id` inside the handler. */
   crossChain?: bool,
@@ -196,8 +197,8 @@ and effectContext = {
   log: logger,
   effect: 'input 'output. (effect<'input, 'output>, 'input) => promise<'output>,
   mutable cache: bool,
-  /** The chain the effect was called on. Only available on effects created with
-   `crossChain: false`; accessing it on a cross-chain effect throws. */
+  /** The chain the effect was called on. Only available on chain-scoped
+   effects; accessing it on a cross-chain effect throws. */
   chain: effectChain,
 }
 and effectArgs<'input> = {
@@ -259,10 +260,9 @@ let createEffect = (
     | Some(true) => true
     | _ => false
     },
-    crossChain: switch options.crossChain {
-    | Some(false) => false
-    | _ => true
-    },
+    // Left unresolved: the config's `defaultCrossChain` fills it in when the
+    // effect didn't state one, and the config isn't available here.
+    crossChain: options.crossChain,
     rateLimit: switch options.rateLimit {
     | Disable => None
     | Enable({calls, per}) =>

package/src/Envio.res.mjs

@@ -30,12 +30,11 @@ function createEffect(options, handler) {
     output: s.m(outputSchema)
   }));
   let match = options.cache;
-  let match$1 = options.crossChain;
-  let match$2 = options.rateLimit;
+  let match$1 = options.rateLimit;
   let tmp;
-  tmp = match$2 === false ? undefined : ({
-      callsPerDuration: match$2.calls,
-      durationMs: durationToMs(match$2.per)
+  tmp = match$1 === false ? undefined : ({
+      callsPerDuration: match$1.calls,
+      durationMs: durationToMs(match$1.per)
     });
   return {
     name: options.name,
@@ -45,7 +44,7 @@ function createEffect(options, handler) {
       outputSchema: outputSchema
     },
     defaultShouldCache: match !== undefined ? match : false,
-    crossChain: match$1 !== undefined ? match$1 : true,
+    crossChain: options.crossChain,
     output: outputSchema,
     input: S$RescriptSchema.schema(param => options.input),
     rateLimit: tmp

package/src/Hasura.res

@@ -281,6 +281,21 @@ let createSelectPermission = async (
   )
 }
 
+// The column_mapping references columns by their db names. A reference between
+// two per-chain entities is only meaningful within one chain, so the chain-id
+// column joins alongside the id — without it the relationship would resolve to
+// another chain's row with the same id.
+let makeColumnMapping = (~relationalKey, ~isDerivedFrom, ~chainIdColumn) => {
+  let pairs = [
+    isDerivedFrom ? `"id": "${relationalKey}"` : `"${relationalKey}": "id"`,
+  ]
+  switch chainIdColumn {
+  | Some(column) => pairs->Array.push(`"${column}": "${column}"`)->ignore
+  | None => ()
+  }
+  `{${pairs->Array.joinUnsafe(", ")}}`
+}
+
 let createEntityRelationship = async (
   ~endpoint,
   ~auth,
@@ -291,10 +306,9 @@ let createEntityRelationship = async (
   ~objectName: string,
   ~mappedEntity: string,
   ~isDerivedFrom: bool,
+  ~chainIdColumn: option<string>,
   ~comment: option<string>=?,
 ) => {
-  // The column_mapping references columns by their db names
-  let derivedFromTo = isDerivedFrom ? `"id": "${relationalKey}"` : `"${relationalKey}" : "id"`
 
   let tableJson = {
     "schema": pgSchema,
@@ -306,7 +320,9 @@ let createEntityRelationship = async (
         "schema": pgSchema,
         "name": mappedEntity,
       },
-      "column_mapping": JSON.parseOrThrow(`{${derivedFromTo}}`),
+      "column_mapping": JSON.parseOrThrow(
+        makeColumnMapping(~relationalKey, ~isDerivedFrom, ~chainIdColumn),
+      ),
     },
   }->(Utils.magic: {..} => JSON.t)
 
@@ -388,9 +404,25 @@ let trackDatabase = async (
     )
   }
 
+  // Both sides of a relationship must be per-chain for the chain to be part of
+  // the join; a per-chain entity referencing a cross-chain one resolves by id
+  // alone. The reverse (cross-chain referencing per-chain) is rejected at
+  // codegen, so it can't reach here.
+  let chainIdColumnOf = (entityName: string) =>
+    userEntities
+    ->Array.find((e: Internal.entityConfig) => e.name === entityName)
+    ->Option.flatMap(e => e.table->Table.getChainIdField)
+    ->Option.map(Table.getPgDbFieldName)
+
   for i in 0 to userEntities->Array.length - 1 {
     let entityConfig = userEntities->Array.getUnsafe(i)
     let {tableName} = entityConfig.table
+    let ownChainIdColumn = entityConfig.table->Table.getChainIdField->Option.map(Table.getPgDbFieldName)
+    let sharedChainIdColumn = mappedEntity =>
+      switch (ownChainIdColumn, chainIdColumnOf(mappedEntity)) {
+      | (Some(column), Some(_)) => Some(column)
+      | _ => None
+      }
 
     //Set array relationships
     let derivedFromFields = entityConfig.table->Table.getDerivedFromFields
@@ -410,6 +442,7 @@ let trackDatabase = async (
         ~objectName=derivedFromField.fieldName,
         ~relationalKey=relationalFieldName,
         ~mappedEntity=derivedFromField.derivedFromEntity,
+        ~chainIdColumn=sharedChainIdColumn(derivedFromField.derivedFromEntity),
         ~comment=?derivedFromField.description,
       )
     }
@@ -428,6 +461,7 @@ let trackDatabase = async (
         ~objectName=field.fieldName,
         ~relationalKey=field->Table.getPgDbFieldName,
         ~mappedEntity=linkedEntityName,
+        ~chainIdColumn=sharedChainIdColumn(linkedEntityName),
         ~comment=?field.description,
       )
     }

package/src/Hasura.res.mjs

@@ -7,6 +7,7 @@ import * as Utils from "./Utils.res.mjs";
 import * as Schema from "./db/Schema.res.mjs";
 import * as Logging from "./Logging.res.mjs";
 import * as InternalTable from "./db/InternalTable.res.mjs";
+import * as Stdlib_Option from "@rescript/runtime/lib/es6/Stdlib_Option.js";
 import * as S$RescriptSchema from "rescript-schema/src/S.res.mjs";
 import * as Primitive_exceptions from "@rescript/runtime/lib/es6/Primitive_exceptions.js";
 
@@ -247,8 +248,15 @@ async function createSelectPermission(endpoint, auth, tableName, pgSchema, respo
   });
 }
 
-async function createEntityRelationship(endpoint, auth, pgSchema, tableName, relationshipType, relationalKey, objectName, mappedEntity, isDerivedFrom, comment) {
-  let derivedFromTo = isDerivedFrom ? `"id": "` + relationalKey + `"` : `"` + relationalKey + `" : "id"`;
+function makeColumnMapping(relationalKey, isDerivedFrom, chainIdColumn) {
+  let pairs = [isDerivedFrom ? `"id": "` + relationalKey + `"` : `"` + relationalKey + `": "id"`];
+  if (chainIdColumn !== undefined) {
+    pairs.push(`"` + chainIdColumn + `": "` + chainIdColumn + `"`);
+  }
+  return `{` + pairs.join(", ") + `}`;
+}
+
+async function createEntityRelationship(endpoint, auth, pgSchema, tableName, relationshipType, relationalKey, objectName, mappedEntity, isDerivedFrom, chainIdColumn, comment) {
   let tableJson = {
     schema: pgSchema,
     name: tableName
@@ -259,7 +267,7 @@ async function createEntityRelationship(endpoint, auth, pgSchema, tableName, rel
         schema: pgSchema,
         name: mappedEntity
       },
-      column_mapping: JSON.parse(`{` + derivedFromTo + `}`)
+      column_mapping: JSON.parse(makeColumnMapping(relationalKey, isDerivedFrom, chainIdColumn))
     }
   };
   let args = {
@@ -313,21 +321,30 @@ async function trackDatabase(endpoint, auth, pgSchema, userEntities, aggregateEn
     let tableName = tableNames[i];
     await createSelectPermission(endpoint, auth, tableName, pgSchema, responseLimit, aggregateEntities);
   }
+  let chainIdColumnOf = entityName => Stdlib_Option.map(Stdlib_Option.flatMap(userEntities.find(e => e.name === entityName), e => Table.getChainIdField(e.table)), Table.getPgDbFieldName);
   for (let i$1 = 0, i_finish$1 = userEntities.length; i$1 < i_finish$1; ++i$1) {
     let entityConfig = userEntities[i$1];
     let match = entityConfig.table;
     let tableName$1 = match.tableName;
+    let ownChainIdColumn = Stdlib_Option.map(Table.getChainIdField(entityConfig.table), Table.getPgDbFieldName);
+    let sharedChainIdColumn = mappedEntity => {
+      let match = chainIdColumnOf(mappedEntity);
+      if (ownChainIdColumn !== undefined && match !== undefined) {
+        return ownChainIdColumn;
+      }
+    };
     let derivedFromFields = Table.getDerivedFromFields(entityConfig.table);
     for (let j = 0, j_finish = derivedFromFields.length; j < j_finish; ++j) {
       let derivedFromField = derivedFromFields[j];
       let relationalFieldName = Utils.unwrapResultExn(Schema.getDerivedFromPgFieldName(schema, derivedFromField));
-      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "array", relationalFieldName, derivedFromField.fieldName, derivedFromField.derivedFromEntity, true, derivedFromField.description);
+      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "array", relationalFieldName, derivedFromField.fieldName, derivedFromField.derivedFromEntity, true, sharedChainIdColumn(derivedFromField.derivedFromEntity), derivedFromField.description);
     }
     let linkedEntityFields = Table.getLinkedEntityFields(entityConfig.table);
     for (let j$1 = 0, j_finish$1 = linkedEntityFields.length; j$1 < j_finish$1; ++j$1) {
       let match$1 = linkedEntityFields[j$1];
+      let linkedEntityName = match$1[1];
       let field = match$1[0];
-      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "object", Table.getPgDbFieldName(field), field.fieldName, match$1[1], false, field.description);
+      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "object", Table.getPgDbFieldName(field), field.fieldName, linkedEntityName, false, sharedChainIdColumn(linkedEntityName), field.description);
     }
   }
   return Logging.info("Hasura configuration completed");
@@ -346,6 +363,7 @@ export {
   makeColumnConfigs,
   trackTables,
   createSelectPermission,
+  makeColumnMapping,
   createEntityRelationship,
   trackDatabase,
 }

package/src/InMemoryStore.res

@@ -2,11 +2,45 @@
 // mutations route through IndexerState's domain operations; the write loop and
 // capacity/flush coordination live in Writing.
 
+// The scope must match the entity's own: a cross-chain entity has one table on
+// the indexer, a per-chain entity one table per ChainState. Taking the scope
+// rather than an optional chain id keeps a dummy chain id unrepresentable.
 let getInMemTable = (
   state: IndexerState.t,
   ~entityConfig: Internal.entityConfig,
+  ~scope: Internal.chainScope,
 ): InMemoryTable.Entity.t =>
-  state->IndexerState.entities->IndexerState.EntityTables.get(~entityName=entityConfig.name)
+  switch scope {
+  | CrossChain => state->IndexerState.entities
+  | Chain(chainId) => state->IndexerState.getChainState(~chainId)->ChainState.entities
+  }->EntityTables.get(~entityName=entityConfig.name)
+
+// The scope a given entity's rows live in when reached from a handler running
+// on `chainId`.
+let entityScope = (entityConfig: Internal.entityConfig, ~chainId): Internal.chainScope =>
+  entityConfig.crossChain ? CrossChain : Chain(chainId)
+
+// The chain a row loaded from storage belongs to, taken off the row so what's
+// left matches the entity schema the handlers see. A per-chain entity whose row
+// carries no chain id would silently land in the wrong partition, so it throws.
+let takeRowScope = (
+  entity: Internal.entity,
+  ~entityConfig: Internal.entityConfig,
+): Internal.chainScope =>
+  switch entityConfig.table->Table.getChainIdField {
+  | None => CrossChain
+  | Some(field) =>
+    let row = entity->(Utils.magic: Internal.entity => dict<ChainId.t>)
+    switch row->Utils.Dict.dangerouslyGetNonOption(field.fieldName) {
+    | Some(chainId) =>
+      row->Utils.Dict.deleteInPlace(field.fieldName)
+      Chain(chainId)
+    | None =>
+      JsError.throwWithMessage(
+        `Rollback row for the per-chain entity "${entityConfig.name}" with id "${entity.id}" is missing its "${field.fieldName}" value.`,
+      )
+    }
+  }
 
 let getEffectInMemTable = (
   state: IndexerState.t,
@@ -114,16 +148,16 @@ let prepareRollbackDiff = async (
   let _ = await persistence.allEntities
   ->Array.filter(entityConfig => entityConfig.storage.postgres)
   ->Array.map(async entityConfig => {
-    let entityTable = state->getInMemTable(~entityConfig)
-
-    let (removedIds, restoredEntitiesResult) = await persistence.storage.getRollbackData(
+    let (removals, restoredEntitiesResult) = await persistence.storage.getRollbackData(
       ~entityConfig,
       ~rollbackTargetCheckpointId,
     )
 
-    removedIds->Array.forEach(entityId => {
+    removals->Array.forEach(({entityId, scope}: Persistence.rollbackRemoval) => {
       deletedEntities->Utils.Dict.push(entityConfig.name, entityId)
-      entityTable->InMemoryTable.Entity.set(
+      state
+      ->getInMemTable(~entityConfig, ~scope)
+      ->InMemoryTable.Entity.set(
         ~committedCheckpointId,
         Delete({
           entityId,
@@ -138,8 +172,11 @@ let prepareRollbackDiff = async (
       ->(Utils.magic: array<unknown> => array<Internal.entity>)
 
     restoredEntities->Array.forEach((entity: Internal.entity) => {
+      let scope = entity->takeRowScope(~entityConfig)
       setEntities->Utils.Dict.push(entityConfig.name, entity.id)
-      entityTable->InMemoryTable.Entity.set(
+      state
+      ->getInMemTable(~entityConfig, ~scope)
+      ->InMemoryTable.Entity.set(
         ~committedCheckpointId,
         Set({
           entityId: entity.id->EntityId.unsafeOfString,
@@ -162,7 +199,10 @@ let prepareRollbackDiff = async (
 // registered them: the store is where they already live, and it knows which
 // ones the database hasn't seen yet.
 let setBatchDcs = (state: IndexerState.t, ~batch: Batch.t) => {
-  let inMemTable = state->getInMemTable(~entityConfig=InternalTable.EnvioAddresses.entityConfig)
+  let inMemTable = state->getInMemTable(
+    ~entityConfig=InternalTable.EnvioAddresses.entityConfig,
+    ~scope=CrossChain,
+  )
   let committedCheckpointId = state->IndexerState.committedCheckpointId
 
   batch.progressedChainsById->Utils.Dict.forEach(progressedChain => {

package/src/InMemoryStore.res.mjs

@@ -6,13 +6,40 @@ import * as Config from "./Config.res.mjs";
 import * as Internal from "./Internal.res.mjs";
 import * as ChainState from "./ChainState.res.mjs";
 import * as EffectState from "./EffectState.res.mjs";
+import * as EntityTables from "./EntityTables.res.mjs";
 import * as IndexerState from "./IndexerState.res.mjs";
 import * as InMemoryTable from "./InMemoryTable.res.mjs";
 import * as Stdlib_Option from "@rescript/runtime/lib/es6/Stdlib_Option.js";
+import * as Stdlib_JsError from "@rescript/runtime/lib/es6/Stdlib_JsError.js";
+import * as Primitive_option from "@rescript/runtime/lib/es6/Primitive_option.js";
 import * as S$RescriptSchema from "rescript-schema/src/S.res.mjs";
 
-function getInMemTable(state, entityConfig) {
-  return IndexerState.EntityTables.get(IndexerState.entities(state), entityConfig.name);
+function getInMemTable(state, entityConfig, scope) {
+  let tmp;
+  tmp = scope === "crossChain" ? IndexerState.entities(state) : ChainState.entities(IndexerState.getChainState(state, scope));
+  return EntityTables.get(tmp, entityConfig.name);
+}
+
+function entityScope(entityConfig, chainId) {
+  if (entityConfig.crossChain) {
+    return "crossChain";
+  } else {
+    return chainId;
+  }
+}
+
+function takeRowScope(entity, entityConfig) {
+  let field = Table.getChainIdField(entityConfig.table);
+  if (field === undefined) {
+    return "crossChain";
+  }
+  let chainId = entity[field.fieldName];
+  if (chainId !== undefined) {
+    Utils.Dict.deleteInPlace(entity, field.fieldName);
+    return Primitive_option.valFromOption(chainId);
+  } else {
+    return Stdlib_JsError.throwWithMessage(`Rollback row for the per-chain entity "` + entityConfig.name + `" with id "` + entity.id + `" is missing its "` + field.fieldName + `" value.`);
+  }
 }
 
 function getEffectInMemTable(state, effect, scope) {
@@ -88,11 +115,11 @@ async function prepareRollbackDiff(state, rollbackTargetCheckpointId, rollbackDi
   let deletedEntities = {};
   let setEntities = {};
   await Promise.all(persistence.allEntities.filter(entityConfig => entityConfig.storage.postgres).map(async entityConfig => {
-    let entityTable = getInMemTable(state, entityConfig);
     let match = await persistence.storage.getRollbackData(entityConfig, rollbackTargetCheckpointId);
-    match[0].forEach(entityId => {
+    match[0].forEach(param => {
+      let entityId = param.entityId;
       Utils.Dict.push(deletedEntities, entityConfig.name, entityId);
-      InMemoryTable.Entity.set(entityTable, committedCheckpointId, {
+      InMemoryTable.Entity.set(getInMemTable(state, entityConfig, param.scope), committedCheckpointId, {
         type: "DELETE",
         entityId: entityId,
         checkpointId: rollbackDiffCheckpointId
@@ -100,8 +127,9 @@ async function prepareRollbackDiff(state, rollbackTargetCheckpointId, rollbackDi
     });
     let restoredEntities = S$RescriptSchema.parseOrThrow(match[1], Table.pgRowsSchema(entityConfig.table));
     restoredEntities.forEach(entity => {
+      let scope = takeRowScope(entity, entityConfig);
       Utils.Dict.push(setEntities, entityConfig.name, entity.id);
-      InMemoryTable.Entity.set(entityTable, committedCheckpointId, {
+      InMemoryTable.Entity.set(getInMemTable(state, entityConfig, scope), committedCheckpointId, {
         type: "SET",
         entityId: entity.id,
         entity: entity,
@@ -116,7 +144,7 @@ async function prepareRollbackDiff(state, rollbackTargetCheckpointId, rollbackDi
 }
 
 function setBatchDcs(state, batch) {
-  let inMemTable = getInMemTable(state, Config.EnvioAddresses.entityConfig);
+  let inMemTable = getInMemTable(state, Config.EnvioAddresses.entityConfig, "crossChain");
   let committedCheckpointId = IndexerState.committedCheckpointId(state);
   Utils.Dict.forEach(batch.progressedChainsById, progressedChain => {
     let chainId = progressedChain.fetchState.chainId;
@@ -155,6 +183,8 @@ function setBatchDcs(state, batch) {
 
 export {
   getInMemTable,
+  entityScope,
+  takeRowScope,
   getEffectInMemTable,
   hasEffectOutput,
   getEffectOutputUnsafe,

package/src/IndexerState.res

@@ -5,28 +5,6 @@ type rollbackState =
   | FoundReorgDepth({chainId: ChainId.t, rollbackTargetBlockNumber: int})
   | RollbackReady({eventsProcessedDiffByChain: dict<float>})
 
-module EntityTables = {
-  type t = dict<InMemoryTable.Entity.t>
-  exception UndefinedEntity({entityName: string})
-  let make = (entities: array<Internal.entityConfig>): t => {
-    let init = Dict.make()
-    entities->Array.forEach(entityConfig => {
-      init->Dict.set((entityConfig.name :> string), InMemoryTable.Entity.make())
-    })
-    init
-  }
-
-  let get = (self: t, ~entityName: string) => {
-    switch self->Utils.Dict.dangerouslyGetNonOption(entityName) {
-    | Some(table) => table
-    | None =>
-      UndefinedEntity({entityName: entityName})->ErrorHandling.mkLogAndRaise(
-        ~msg="Unexpected, entity InMemoryTable is undefined",
-      )
-    }
-  }
-}
-
 // Per-(contract, event) handler counters rendered into the
 // envio_processing_handler_* and envio_preload_handler_* metrics.
 type handlerStat = {
@@ -75,6 +53,8 @@ type t = {
   persistence: Persistence.t,
   // --- In-memory store: entity/effect tables and the pending-write queue. ---
   allEntities: array<Internal.entityConfig>,
+  // Cross-chain entities only; each ChainState holds its own partition for the
+  // per-chain ones.
   mutable entities: EntityTables.t,
   effectState: EffectState.t,
   mutable rollback: option<Persistence.rollback>,
@@ -186,7 +166,7 @@ let make = (
     config,
     persistence,
     allEntities: persistence.allEntities,
-    entities: EntityTables.make(persistence.allEntities),
+    entities: EntityTables.make(persistence.allEntities->EntityTables.crossChain),
     effectState: EffectState.make(),
     rollback: None,
     committedCheckpointId,
@@ -467,6 +447,30 @@ let config = (state: t) => state.config
 let persistence = (state: t) => state.persistence
 let allEntities = (state: t) => state.allEntities
 let entities = (state: t) => state.entities
+
+// Every in-memory entity table across all scopes, cross-chain first. The size,
+// drop and flush passes walk the whole store this way instead of assuming a
+// single partition.
+let eachEntityTable = (state: t, fn: (~entityConfig: Internal.entityConfig, ~scope: Internal.chainScope, ~table: InMemoryTable.Entity.t) => unit) => {
+  let chainStates = state.crossChainState->CrossChainState.chainStates
+  state.allEntities->Array.forEach(entityConfig =>
+    if entityConfig.crossChain {
+      fn(
+        ~entityConfig,
+        ~scope=Internal.CrossChain,
+        ~table=state.entities->EntityTables.get(~entityName=entityConfig.name),
+      )
+    } else {
+      chainStates->Utils.Dict.forEach(chainState => {
+        fn(
+          ~entityConfig,
+          ~scope=Internal.Chain((chainState->ChainState.chainConfig).id),
+          ~table=chainState->ChainState.entities->EntityTables.get(~entityName=entityConfig.name),
+        )
+      })
+    }
+  )
+}
 let effectState = (state: t) => state.effectState
 let committedCheckpointId = (state: t) => state.committedCheckpointId
 let processedCheckpointId = (state: t) => state.processedCheckpointId
@@ -802,7 +806,9 @@ let beginRollbackDiff = (
   ~diffCheckpointId,
   ~progressBlockNumberByChainId,
 ) => {
-  state.entities = EntityTables.make(state.allEntities)
+  let perChainEntities = state.allEntities->EntityTables.perChain
+  state.entities = EntityTables.make(state.allEntities->EntityTables.crossChain)
+  state->chainStates->Utils.Dict.forEach(cs => cs->ChainState.resetEntities(~perChainEntities))
   state.effectState->EffectState.resetForRollback
   state.rollback = Some({
     targetCheckpointId,