envio 3.7.0 → 3.9.0-resolvers.0

package/src/FetchState.res.mjs

@@ -1123,25 +1123,16 @@ function warnRejectedRegistration(verdict, chainId, contractAddress, contractNam
       contractName: contractName
     }, `Skipping contract registration: Not a valid address for this chain's ecosystem.`);
   }
-  switch (verdict.TAG) {
-    case "Added" :
-      return;
-    case "Duplicate" :
-      let existingEffectiveStartBlock = verdict.existingEffectiveStartBlock;
-      let effectiveStartBlock = verdict.effectiveStartBlock;
-      if (existingEffectiveStartBlock > effectiveStartBlock) {
-        return warnAddressRegistration(chainId, contractAddress, {
-          existingBlockNumber: existingEffectiveStartBlock,
-          newBlockNumber: effectiveStartBlock
-        }, `Skipping contract registration: Contract address is already registered at a later block number. Currently registration of the same contract address is not supported by Envio. Reach out to us if it's a problem for you.`);
-      } else {
-        return;
-      }
-    case "Conflict" :
-      return warnAddressRegistration(chainId, contractAddress, {
-        existingContractType: verdict.existingContractName,
-        newContractType: contractName
-      }, `Skipping contract registration: Contract address is already registered for one contract and cannot be registered for another contract.`);
+  if (verdict.TAG === "Added") {
+    return;
+  }
+  let existingEffectiveStartBlock = verdict.existingEffectiveStartBlock;
+  let effectiveStartBlock = verdict.effectiveStartBlock;
+  if (existingEffectiveStartBlock > effectiveStartBlock) {
+    return warnAddressRegistration(chainId, contractAddress, {
+      existingBlockNumber: existingEffectiveStartBlock,
+      newBlockNumber: effectiveStartBlock
+    }, `Skipping same-contract re-registration: the address is already registered for this contract. The start block does not move earlier.`);
   }
 }
 
@@ -1727,7 +1718,7 @@ function getReadyItemsCount(fetchState, targetSize, fromItem) {
   return acc;
 }
 
-function make$1(startBlock, endBlock, onEventRegistrations, addressStore, addresses, maxAddrInPartition, chainId, maxOnBlockBufferSize, knownHeight, progressBlockNumberOpt, onBlockRegistrationsOpt, blockLagOpt, firstEventBlockOpt, clientFilterAddressThresholdOpt, isResumedOpt) {
+function make$1(startBlock, endBlock, onEventRegistrations, addressStore, addressRows, maxAddrInPartition, chainId, maxOnBlockBufferSize, knownHeight, progressBlockNumberOpt, onBlockRegistrationsOpt, blockLagOpt, firstEventBlockOpt, clientFilterAddressThresholdOpt, isResumedOpt) {
   let progressBlockNumber = progressBlockNumberOpt !== undefined ? progressBlockNumberOpt : startBlock - 1 | 0;
   let onBlockRegistrations = onBlockRegistrationsOpt !== undefined ? onBlockRegistrationsOpt : [];
   let blockLag = blockLagOpt !== undefined ? blockLagOpt : 0;
@@ -1762,29 +1753,24 @@ function make$1(startBlock, endBlock, onEventRegistrations, addressStore, addres
     });
   }
   let normalSelection = makeSelection(normalRegistrations, true, undefined);
-  AddressStore.seedBatch(addressStore, addresses.map(contract => ({
-    address: contract.address,
-    contractName: contract.contractName,
-    registrationBlock: contract.registrationBlock
-  }))).forEach((verdict, idx) => {
-    let contract = addresses[idx];
-    warnRejectedRegistration(verdict, chainId, contract.address, contract.contractName);
-  });
+  AddressStore.seedRows(addressStore, addressRows).forEach(rejected => warnAddressRegistration(chainId, rejected.address, {
+    contractName: rejected.contractName,
+    existingBlockNumber: rejected.existingEffectiveStartBlock,
+    newBlockNumber: rejected.effectiveStartBlock
+  }, `Skipping a stored address: it is already registered for this contract.`));
   let dynamicContracts = new Set();
   let clientFilteredContracts = new Set();
   let registeringSetsByContract = {};
-  addresses.forEach(contract => {
-    let contractName = contract.contractName;
+  Array.from(contractNamesWithNormalEvents).forEach(contractName => {
+    if (addressStore.contractCount(contractName) > 0) {
+      registeringSetsByContract[contractName] = AddressStore.makeSet(addressStore, contractName, undefined);
+      return;
+    }
+  });
+  addressStore.dynamicContractNames().forEach(contractName => {
     if (contractNamesWithNormalEvents.has(contractName)) {
-      if (!(contractName in registeringSetsByContract)) {
-        registeringSetsByContract[contractName] = AddressStore.makeSet(addressStore, contractName, undefined);
-      }
-      if (contract.registrationBlock !== -1) {
-        dynamicContracts.add(contractName);
-        return;
-      } else {
-        return;
-      }
+      dynamicContracts.add(contractName);
+      return;
     }
   });
   if (clientFilterAddressThreshold !== undefined) {
@@ -1797,7 +1783,7 @@ function make$1(startBlock, endBlock, onEventRegistrations, addressStore, addres
   }
   let optimizedPartitions = createPartitions(registeringSetsByContract, addressStore, dynamicContracts, clientFilteredContracts, normalSelection, maxAddrInPartition, partitions.length, partitions, progressBlockNumber, knownHeight);
   if (optimizedPartitions.idsInAscOrder.length === 0 && Utils.$$Array.isEmpty(onBlockRegistrations)) {
-    Stdlib_JsError.throwWithMessage(`Invalid configuration: Nothing to fetch on chain ` + ChainId.toString(chainId) + `. ` + (`addresses=` + addresses.length.toString() + `, `) + (`onEventRegistrations=` + onEventRegistrations.length.toString() + `, `) + (`normalRegistrations=` + normalRegistrations.length.toString() + `. `) + `Make sure that you provided at least one contract address to index, or have events with Wildcard mode enabled, or have onBlock handlers.`);
+    Stdlib_JsError.throwWithMessage(`Invalid configuration: Nothing to fetch on chain ` + ChainId.toString(chainId) + `. ` + (`addresses=` + addressRows.addresses.length.toString() + `, `) + (`onEventRegistrations=` + onEventRegistrations.length.toString() + `, `) + (`normalRegistrations=` + normalRegistrations.length.toString() + `. `) + `Make sure that you provided at least one contract address to index, or have events with Wildcard mode enabled, or have onBlock handlers.`);
   }
   let buffer = [];
   let latestOnBlockBlockNumber;
@@ -1861,8 +1847,8 @@ function rollbackPendingQueries(mutPendingQueries, targetBlockNumber) {
   return adjusted;
 }
 
-function rollback(fetchState, addressStore, targetBlockNumber) {
-  addressStore.rollback(targetBlockNumber);
+function rollback(fetchState, rolledBackAddressStore, targetBlockNumber) {
+  let rolledBackAddresses = rolledBackAddressStore.rollback(targetBlockNumber);
   let keptPartitions = [];
   let nextKeptIdRef = 0;
   let registeringSetsByContract = {};
@@ -1921,8 +1907,8 @@ function rollback(fetchState, addressStore, targetBlockNumber) {
       });
     }
   }
-  let optimizedPartitions = createPartitions(registeringSetsByContract, addressStore, fetchState.optimizedPartitions.dynamicContracts, fetchState.optimizedPartitions.clientFilteredContracts, fetchState.normalSelection, fetchState.optimizedPartitions.maxAddrInPartition, nextKeptIdRef, keptPartitions, targetBlockNumber, fetchState.knownHeight);
-  return updateInternal({
+  let optimizedPartitions = createPartitions(registeringSetsByContract, rolledBackAddressStore, fetchState.optimizedPartitions.dynamicContracts, fetchState.optimizedPartitions.clientFilteredContracts, fetchState.normalSelection, fetchState.optimizedPartitions.maxAddrInPartition, nextKeptIdRef, keptPartitions, targetBlockNumber, fetchState.knownHeight);
+  let rolledBack = updateInternal({
     optimizedPartitions: fetchState.optimizedPartitions,
     startBlock: fetchState.startBlock,
     endBlock: fetchState.endBlock,
@@ -1941,6 +1927,10 @@ function rollback(fetchState, addressStore, targetBlockNumber) {
     tmp = item.kind === 0 ? item.blockNumber : item.blockNumber;
     return tmp <= targetBlockNumber;
   }), true, undefined, undefined);
+  return {
+    fetchState: rolledBack,
+    rolledBackAddresses: rolledBackAddresses
+  };
 }
 
 function resetPendingQueries(fetchState) {

package/src/HandlerRegister.res

@@ -154,6 +154,8 @@ let buildOnEventRegistrationWith = (
       ~isWildcard,
       ~handler,
       ~contractRegister,
+      ~where,
+      ~fieldSelection?,
       ~startBlock?,
     ) :> Internal.onEventRegistration)
   | Evm =>
@@ -299,21 +301,29 @@ let addOnEventRegistration = (
   let fieldSelection = switch eventOptions->Option.flatMap(v => v.fields) {
   | None => None
   | Some(fields) =>
-    // Only EVM resolves a registration selection of its own; Fuel and SVM take
-    // theirs from the event config, so an inline one would be silently dropped.
-    if registration.config.ecosystem.name !== Evm {
+    switch registration.config.ecosystem.name {
+    | Evm =>
+      Some(
+        EventConfigBuilder.resolveInlineFieldSelection(
+          fields,
+          ~contractName,
+          ~eventName,
+          ~enableRawEvents=registration.config.enableRawEvents,
+        ),
+      )
+    | Svm =>
+      Some(
+        EventConfigBuilder.resolveSvmInlineFieldSelection(
+          fields,
+          ~contractName,
+          ~eventName,
+        ),
+      )
+    | Fuel =>
       JsError.throwWithMessage(
-        `The fields option of the "${eventName}" event registration on contract "${contractName}" is only supported on EVM. Select the fields in your config instead.`,
+        `The fields option of the "${eventName}" event registration on contract "${contractName}" is not supported on Fuel. Select the fields in your config instead.`,
       )
     }
-    Some(
-      EventConfigBuilder.resolveInlineFieldSelection(
-        fields,
-        ~contractName,
-        ~eventName,
-        ~enableRawEvents=registration.config.enableRawEvents,
-      ),
-    )
   }
   let matched = ref(false)
   registration.config.chainMap
@@ -350,17 +360,18 @@ let addOnEventRegistration = (
   if !matched.contents {
     let (contractNames, eventNames) = registration->describeConfigured(~contractName)
     let listOr = (names, empty) =>
-      names->Utils.Array.isEmpty ? empty : names->Array.joinUnsafe(", ")
+      names->Utils.Array.isEmpty ? empty : names->Utils.Array.quotedJoin
+    let {contractNoun, eventNoun} = registration.config.ecosystem
     if eventNames->Utils.Array.isEmpty {
       JsError.throwWithMessage(
-        `Contract "${contractName}" is not configured on any chain, so its handler for "${eventName}" would never run. Add it to your config, or remove the registration. Configured contracts: ${listOr(
+        `${contractNoun->Utils.String.capitalize} "${contractName}" is not configured on any chain, so its handler for "${eventName}" would never run. Add it to your config, or remove the registration. Configured ${contractNoun}s: ${listOr(
             contractNames,
             "none",
           )}.`,
       )
     } else {
       JsError.throwWithMessage(
-        `Event "${eventName}" is not configured on contract "${contractName}", so its handler would never run. Add it to your config, or remove the registration. Configured events on "${contractName}": ${listOr(
+        `${eventNoun->Utils.String.capitalize} "${eventName}" is not configured on ${contractNoun} "${contractName}", so its handler would never run. Add it to your config, or remove the registration. Configured ${eventNoun}s on "${contractName}": ${listOr(
             eventNames,
             "none",
           )}.`,

package/src/HandlerRegister.res.mjs

@@ -102,7 +102,7 @@ function buildOnEventRegistrationWith(config, chainId, eventConfig, isWildcard,
     case "fuel" :
       return EventConfigBuilder.buildFuelOnEventRegistration(eventConfig, isWildcard, handler, contractRegister, startBlock);
     case "svm" :
-      return EventConfigBuilder.buildSvmOnEventRegistration(eventConfig, isWildcard, handler, contractRegister, startBlock);
+      return EventConfigBuilder.buildSvmOnEventRegistration(eventConfig, isWildcard, handler, contractRegister, where, fieldSelection, startBlock);
   }
 }
 
@@ -200,7 +200,24 @@ function addOnEventRegistration(registration, contractName, eventName, handler,
   let isWildcard = Stdlib_Option.getOr(Stdlib_Option.flatMap(eventOptions, v => v.wildcard), false);
   let where = Stdlib_Option.flatMap(eventOptions, v => v.where);
   let fields = Stdlib_Option.flatMap(eventOptions, v => v.fields);
-  let fieldSelection = fields !== undefined ? (registration.config.ecosystem.name !== "evm" ? Stdlib_JsError.throwWithMessage(`The fields option of the "` + eventName + `" event registration on contract "` + contractName + `" is only supported on EVM. Select the fields in your config instead.`) : undefined, EventConfigBuilder.resolveInlineFieldSelection(fields, contractName, eventName, registration.config.enableRawEvents)) : undefined;
+  let fieldSelection;
+  if (fields !== undefined) {
+    let fields$1 = Primitive_option.valFromOption(fields);
+    let match = registration.config.ecosystem.name;
+    switch (match) {
+      case "evm" :
+        fieldSelection = EventConfigBuilder.resolveInlineFieldSelection(fields$1, contractName, eventName, registration.config.enableRawEvents);
+        break;
+      case "fuel" :
+        fieldSelection = Stdlib_JsError.throwWithMessage(`The fields option of the "` + eventName + `" event registration on contract "` + contractName + `" is not supported on Fuel. Select the fields in your config instead.`);
+        break;
+      case "svm" :
+        fieldSelection = EventConfigBuilder.resolveSvmInlineFieldSelection(fields$1, contractName, eventName);
+        break;
+    }
+  } else {
+    fieldSelection = undefined;
+  }
   let matched = {
     contents: false
   };
@@ -220,19 +237,22 @@ function addOnEventRegistration(registration, contractName, eventName, handler,
   if (matched.contents) {
     return;
   }
-  let match = describeConfigured(registration, contractName);
-  let eventNames = match[1];
+  let match$1 = describeConfigured(registration, contractName);
+  let eventNames = match$1[1];
   let listOr = (names, empty) => {
     if (Utils.$$Array.isEmpty(names)) {
       return empty;
     } else {
-      return names.join(", ");
+      return Utils.$$Array.quotedJoin(names);
     }
   };
+  let match$2 = registration.config.ecosystem;
+  let eventNoun = match$2.eventNoun;
+  let contractNoun = match$2.contractNoun;
   if (Utils.$$Array.isEmpty(eventNames)) {
-    return Stdlib_JsError.throwWithMessage(`Contract "` + contractName + `" is not configured on any chain, so its handler for "` + eventName + `" would never run. Add it to your config, or remove the registration. Configured contracts: ` + listOr(match[0], "none") + `.`);
+    return Stdlib_JsError.throwWithMessage(Utils.$$String.capitalize(contractNoun) + ` "` + contractName + `" is not configured on any chain, so its handler for "` + eventName + `" would never run. Add it to your config, or remove the registration. Configured ` + contractNoun + `s: ` + listOr(match$1[0], "none") + `.`);
   } else {
-    return Stdlib_JsError.throwWithMessage(`Event "` + eventName + `" is not configured on contract "` + contractName + `", so its handler would never run. Add it to your config, or remove the registration. Configured events on "` + contractName + `": ` + listOr(eventNames, "none") + `.`);
+    return Stdlib_JsError.throwWithMessage(Utils.$$String.capitalize(eventNoun) + ` "` + eventName + `" is not configured on ` + contractNoun + ` "` + contractName + `", so its handler would never run. Add it to your config, or remove the registration. Configured ` + eventNoun + `s on "` + contractName + `": ` + listOr(eventNames, "none") + `.`);
   }
 }
 

package/src/Hasura.res

@@ -286,9 +286,7 @@ let createSelectPermission = async (
 // 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"`,
-  ]
+  let pairs = [isDerivedFrom ? `"id": "${relationalKey}"` : `"${relationalKey}": "id"`]
   switch chainIdColumn {
   | Some(column) => pairs->Array.push(`"${column}": "${column}"`)->ignore
   | None => ()
@@ -309,7 +307,6 @@ let createEntityRelationship = async (
   ~chainIdColumn: option<string>,
   ~comment: option<string>=?,
 ) => {
-
   let tableJson = {
     "schema": pgSchema,
     "name": tableName,
@@ -347,15 +344,7 @@ let createEntityRelationship = async (
   )
 }
 
-let trackDatabase = async (
-  ~endpoint,
-  ~auth,
-  ~pgSchema,
-  ~userEntities: array<Internal.entityConfig>,
-  ~aggregateEntities,
-  ~responseLimit,
-  ~schema,
-) => {
+let makeTableConfigs = (~userEntities: array<Internal.entityConfig>) => {
   let exposedInternalTableConfigs = [
     {
       tableName: InternalTable.RawEvents.table.tableName,
@@ -373,12 +362,27 @@ let trackDatabase = async (
       columnConfigs: dict{},
     },
   ]
-  let userTableConfigs = userEntities->Array.map(entity => {
-    tableName: entity.table.tableName,
-    description: entity.table.description,
-    columnConfigs: entity.table->makeColumnConfigs,
-  })
-  let tableConfigs = [exposedInternalTableConfigs, userTableConfigs]->Array.flat
+  let userTableConfigs =
+    userEntities
+    ->Array.filter(entity => !entity.internal)
+    ->Array.map(entity => {
+      tableName: entity.table.tableName,
+      description: entity.table.description,
+      columnConfigs: entity.table->makeColumnConfigs,
+    })
+  [exposedInternalTableConfigs, userTableConfigs]->Array.flat
+}
+
+let trackDatabase = async (
+  ~endpoint,
+  ~auth,
+  ~pgSchema,
+  ~userEntities: array<Internal.entityConfig>,
+  ~aggregateEntities,
+  ~responseLimit,
+  ~schema,
+) => {
+  let tableConfigs = makeTableConfigs(~userEntities)
   let tableNames = tableConfigs->Array.map(c => c.tableName)
 
   Logging.info("Tracking tables in Hasura")
@@ -414,10 +418,14 @@ let trackDatabase = async (
     ->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)
+  // Relationships to an @internal entity can't exist here: codegen rejects a
+  // reference from an exposed entity to an @internal one.
+  let exposedEntities = userEntities->Array.filter(e => !e.internal)
+  for i in 0 to exposedEntities->Array.length - 1 {
+    let entityConfig = exposedEntities->Array.getUnsafe(i)
     let {tableName} = entityConfig.table
-    let ownChainIdColumn = entityConfig.table->Table.getChainIdField->Option.map(Table.getPgDbFieldName)
+    let ownChainIdColumn =
+      entityConfig.table->Table.getChainIdField->Option.map(Table.getPgDbFieldName)
     let sharedChainIdColumn = mappedEntity =>
       switch (ownChainIdColumn, chainIdColumnOf(mappedEntity)) {
       | (Some(column), Some(_)) => Some(column)

package/src/Hasura.res.mjs

@@ -285,7 +285,7 @@ async function createEntityRelationship(endpoint, auth, pgSchema, tableName, rel
   });
 }
 
-async function trackDatabase(endpoint, auth, pgSchema, userEntities, aggregateEntities, responseLimit, schema) {
+function makeTableConfigs(userEntities) {
   let exposedInternalTableConfigs = [
     {
       tableName: InternalTable.RawEvents.table.tableName,
@@ -303,15 +303,19 @@ async function trackDatabase(endpoint, auth, pgSchema, userEntities, aggregateEn
       columnConfigs: {}
     }
   ];
-  let userTableConfigs = userEntities.map(entity => ({
+  let userTableConfigs = userEntities.filter(entity => !entity.internal).map(entity => ({
     tableName: entity.table.tableName,
     description: entity.table.description,
     columnConfigs: makeColumnConfigs(entity.table)
   }));
-  let tableConfigs = [
+  return [
     exposedInternalTableConfigs,
     userTableConfigs
   ].flat();
+}
+
+async function trackDatabase(endpoint, auth, pgSchema, userEntities, aggregateEntities, responseLimit, schema) {
+  let tableConfigs = makeTableConfigs(userEntities);
   let tableNames = tableConfigs.map(c => c.tableName);
   Logging.info("Tracking tables in Hasura");
   await clearHasuraMetadata(endpoint, auth);
@@ -322,8 +326,9 @@ async function trackDatabase(endpoint, auth, pgSchema, userEntities, aggregateEn
     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 exposedEntities = userEntities.filter(e => !e.internal);
+  for (let i$1 = 0, i_finish$1 = exposedEntities.length; i$1 < i_finish$1; ++i$1) {
+    let entityConfig = exposedEntities[i$1];
     let match = entityConfig.table;
     let tableName$1 = match.tableName;
     let ownChainIdColumn = Stdlib_Option.map(Table.getChainIdField(entityConfig.table), Table.getPgDbFieldName);
@@ -365,6 +370,7 @@ export {
   createSelectPermission,
   makeColumnMapping,
   createEntityRelationship,
+  makeTableConfigs,
   trackDatabase,
 }
 /* Rest Not a pure module */

package/src/InMemoryStore.res

@@ -127,14 +127,18 @@ let dropCommittedEffects = (
 
 let prepareRollbackDiff = async (
   state: IndexerState.t,
+  ~rollbackScope: RollbackScope.t,
   ~rollbackTargetCheckpointId,
   ~rollbackDiffCheckpointId,
-  ~progressBlockNumberByChainId,
+  ~progressedChains,
+  ~rolledBackAddresses,
 ) => {
   state->IndexerState.beginRollbackDiff(
     ~targetCheckpointId=rollbackTargetCheckpointId,
     ~diffCheckpointId=rollbackDiffCheckpointId,
-    ~progressBlockNumberByChainId,
+    ~scope=rollbackScope,
+    ~progressedChains,
+    ~rolledBackAddresses,
   )
   let persistence = state->IndexerState.persistence
   let committedCheckpointId = state->IndexerState.committedCheckpointId
@@ -148,8 +152,9 @@ let prepareRollbackDiff = async (
   let _ = await persistence.allEntities
   ->Array.filter(entityConfig => entityConfig.storage.postgres)
   ->Array.map(async entityConfig => {
-    let (removals, restoredEntitiesResult) = await persistence.storage.getRollbackData(
+    let (removals, restoredEntities) = await persistence.storage.getRollbackData(
       ~entityConfig,
+      ~scope=rollbackScope,
       ~rollbackTargetCheckpointId,
     )
 
@@ -166,11 +171,6 @@ let prepareRollbackDiff = async (
       )
     })
 
-    let restoredEntities =
-      restoredEntitiesResult
-      ->S.parseOrThrow(entityConfig.table->Table.pgRowsSchema)
-      ->(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)
@@ -199,12 +199,6 @@ 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,
-    ~scope=CrossChain,
-  )
-  let committedCheckpointId = state->IndexerState.committedCheckpointId
-
   batch.progressedChainsById->Utils.Dict.forEach(progressedChain => {
     let chainId = progressedChain.fetchState.chainId
     let chainState = state->IndexerState.getChainState(~chainId)
@@ -224,31 +218,22 @@ let setBatchDcs = (state: IndexerState.t, ~batch: Batch.t) => {
         }
       }
 
-      chainState
-      ->ChainState.drainAddressesForWrite(
-        ~toBlockInclusive=progressedChain.progressBlockNumber,
-        ~checkpointBlockNumbers,
-      )
-      ->Array.forEach(dc => {
-        let entity: InternalTable.EnvioAddresses.t = {
-          id: InternalTable.EnvioAddresses.makeId(~chainId, ~address=dc.address),
-          chainId,
-          contractName: dc.contractName,
-          registrationBlock: dc.registrationBlock,
-          // Only ever written, never read back. Kept on the table so the column
-          // doesn't need a migration.
-          registrationLogIndex: -1,
-        }
-
-        inMemTable->InMemoryTable.Entity.set(
-          ~committedCheckpointId,
-          Set({
-            entityId: entity.id->EntityId.unsafeOfString,
-            checkpointId: checkpointIds->Array.getUnsafe(dc.checkpointIdx),
-            entity: entity->InternalTable.EnvioAddresses.castToInternal,
-          }),
+      batch.registeredAddresses->Array.pushMany(
+        chainState
+        ->ChainState.drainAddressesForWrite(
+          ~toBlockInclusive=progressedChain.progressBlockNumber,
+          ~checkpointBlockNumbers,
         )
-      })
+        ->Array.map((dc): AddressRows.staged => {
+          row: {
+            chainId,
+            address: dc.address,
+            contractId: dc.contractId,
+            registrationBlock: dc.registrationBlock,
+          },
+          checkpointId: checkpointIds->Array.getUnsafe(dc.checkpointIdx),
+        }),
+      )->ignore
     }
   })
 }

package/src/InMemoryStore.res.mjs

@@ -2,7 +2,6 @@
 
 import * as Table from "./db/Table.res.mjs";
 import * as Utils from "./Utils.res.mjs";
-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";
@@ -12,7 +11,6 @@ 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, scope) {
   let tmp;
@@ -108,14 +106,14 @@ function dropCommittedEffects(inMemTable, committedCheckpointId, keepLoadedFromD
   inMemTable.changesCount = inMemTable.changesCount - keysToDelete.length;
 }
 
-async function prepareRollbackDiff(state, rollbackTargetCheckpointId, rollbackDiffCheckpointId, progressBlockNumberByChainId) {
-  IndexerState.beginRollbackDiff(state, rollbackTargetCheckpointId, rollbackDiffCheckpointId, progressBlockNumberByChainId);
+async function prepareRollbackDiff(state, rollbackScope, rollbackTargetCheckpointId, rollbackDiffCheckpointId, progressedChains, rolledBackAddresses) {
+  IndexerState.beginRollbackDiff(state, rollbackTargetCheckpointId, rollbackDiffCheckpointId, rollbackScope, progressedChains, rolledBackAddresses);
   let persistence = IndexerState.persistence(state);
   let committedCheckpointId = IndexerState.committedCheckpointId(state);
   let deletedEntities = {};
   let setEntities = {};
   await Promise.all(persistence.allEntities.filter(entityConfig => entityConfig.storage.postgres).map(async entityConfig => {
-    let match = await persistence.storage.getRollbackData(entityConfig, rollbackTargetCheckpointId);
+    let match = await persistence.storage.getRollbackData(entityConfig, rollbackScope, rollbackTargetCheckpointId);
     match[0].forEach(param => {
       let entityId = param.entityId;
       Utils.Dict.push(deletedEntities, entityConfig.name, entityId);
@@ -125,8 +123,7 @@ async function prepareRollbackDiff(state, rollbackTargetCheckpointId, rollbackDi
         checkpointId: rollbackDiffCheckpointId
       });
     });
-    let restoredEntities = S$RescriptSchema.parseOrThrow(match[1], Table.pgRowsSchema(entityConfig.table));
-    restoredEntities.forEach(entity => {
+    match[1].forEach(entity => {
       let scope = takeRowScope(entity, entityConfig);
       Utils.Dict.push(setEntities, entityConfig.name, entity.id);
       InMemoryTable.Entity.set(getInMemTable(state, entityConfig, scope), committedCheckpointId, {
@@ -144,8 +141,6 @@ async function prepareRollbackDiff(state, rollbackTargetCheckpointId, rollbackDi
 }
 
 function setBatchDcs(state, batch) {
-  let inMemTable = getInMemTable(state, Config.EnvioAddresses.entityConfig, "crossChain");
-  let committedCheckpointId = IndexerState.committedCheckpointId(state);
   Utils.Dict.forEach(batch.progressedChainsById, progressedChain => {
     let chainId = progressedChain.fetchState.chainId;
     let chainState = IndexerState.getChainState(state, chainId);
@@ -160,24 +155,15 @@ function setBatchDcs(state, batch) {
         checkpointBlockNumbers.push(batch.checkpointBlockNumbers[idx]);
       }
     }
-    ChainState.drainAddressesForWrite(chainState, progressedChain.progressBlockNumber, checkpointBlockNumbers).forEach(dc => {
-      let entity_id = Config.EnvioAddresses.makeId(chainId, dc.address);
-      let entity_registration_block = dc.registrationBlock;
-      let entity_contract_name = dc.contractName;
-      let entity = {
-        id: entity_id,
-        chain_id: chainId,
-        registration_block: entity_registration_block,
-        registration_log_index: -1,
-        contract_name: entity_contract_name
-      };
-      InMemoryTable.Entity.set(inMemTable, committedCheckpointId, {
-        type: "SET",
-        entityId: entity_id,
-        entity: entity,
-        checkpointId: checkpointIds[dc.checkpointIdx]
-      });
-    });
+    batch.registeredAddresses.push(...ChainState.drainAddressesForWrite(chainState, progressedChain.progressBlockNumber, checkpointBlockNumbers).map(dc => ({
+      row: {
+        chainId: chainId,
+        address: dc.address,
+        contractId: dc.contractId,
+        registrationBlock: dc.registrationBlock
+      },
+      checkpointId: checkpointIds[dc.checkpointIdx]
+    })));
   });
 }