Fix: adding the chainId to the HUB and BLOCK_DATA modules (#1359)

arithmetization/src/main/java/net/consensys/linea/zktracer/ZkTracer.java

@@ -60,8 +60,6 @@ public class ZkTracer implements ConflationAwareOperationTracer {
 
   private static final Map<String, Integer> spillings;
 
-  public final BigInteger chainId;
-
   static {
     try {
       // Load spillings configured in src/main/resources/spillings.toml.
@@ -82,17 +80,20 @@ public class ZkTracer implements ConflationAwareOperationTracer {
   @Getter private final List<Exception> tracingExceptions = new FiniteList<>(50);
 
   public ZkTracer() {
-    this(LineaL1L2BridgeSharedConfiguration.EMPTY, Bytes.fromHexString("c0ffee").toBigInteger());
+    this(
+        LineaL1L2BridgeSharedConfiguration.EMPTY,
+        Bytes.fromHexString("c0ffee").toBigInteger().abs());
   }
 
-  public ZkTracer(BigInteger chainId) {
-    this(LineaL1L2BridgeSharedConfiguration.EMPTY, chainId);
+  public ZkTracer(BigInteger nonnegativeChainId) {
+    this(LineaL1L2BridgeSharedConfiguration.EMPTY, nonnegativeChainId);
   }
 
   public ZkTracer(
       final LineaL1L2BridgeSharedConfiguration bridgeConfiguration, BigInteger chainId) {
-    this.hub = new Hub(bridgeConfiguration.contract(), bridgeConfiguration.topic());
-    this.chainId = chainId;
+    BigInteger nonnegativeChainId = chainId.abs();
+    this.hub =
+        new Hub(bridgeConfiguration.contract(), bridgeConfiguration.topic(), nonnegativeChainId);
     for (Module m : this.hub.getModulesToCount()) {
       if (!spillings.containsKey(m.moduleKey())) {
         throw new IllegalStateException(

arithmetization/src/main/java/net/consensys/linea/zktracer/module/blockdata/Blockdata.java

@@ -16,8 +16,8 @@
 package net.consensys.linea.zktracer.module.blockdata;
 
 import static net.consensys.linea.zktracer.module.blockdata.Trace.MAX_CT;
-import static net.consensys.linea.zktracer.types.TransactionUtils.getChainIdFromTransaction;
 
+import java.math.BigInteger;
 import java.nio.MappedByteBuffer;
 import java.util.ArrayDeque;
 import java.util.Deque;
@@ -27,7 +27,6 @@
 import net.consensys.linea.zktracer.ColumnHeader;
 import net.consensys.linea.zktracer.container.module.Module;
 import net.consensys.linea.zktracer.module.rlptxn.RlpTxn;
-import net.consensys.linea.zktracer.module.rlptxn.RlpTxnOperation;
 import net.consensys.linea.zktracer.module.txndata.TxnData;
 import net.consensys.linea.zktracer.module.wcp.Wcp;
 import org.hyperledger.besu.evm.worldstate.WorldView;
@@ -39,6 +38,7 @@ public class Blockdata implements Module {
   private final Wcp wcp;
   private final TxnData txnData;
   private final RlpTxn rlpTxn;
+  private final BigInteger chainId;
   private final Deque<BlockdataOperation> operations = new ArrayDeque<>();
   private boolean conflationFinished = false;
   private static final int TIMESTAMP_BYTESIZE = 4;
@@ -96,27 +96,9 @@ public void commit(List<MappedByteBuffer> buffers) {
     final Trace trace = new Trace(buffers);
 
     final long firstBlockNumber = operations.getFirst().absoluteBlockNumber();
-    final long chainId = getChainIdFromConflation();
     int relblock = 0;
     for (BlockdataOperation blockData : operations) {
       blockData.trace(trace, ++relblock, firstBlockNumber, chainId);
     }
   }
-
-  private long getChainIdFromConflation() {
-    // TODO: this doesn't work if all transaction of the batch are WO ChainId
-    long chainId = -1;
-    for (RlpTxnOperation tx : rlpTxn.operations().getAll()) {
-      try {
-        chainId = getChainIdFromTransaction(tx.tx());
-        break;
-      } catch (Exception e) {
-        continue;
-      }
-    }
-    if (chainId == -1) {
-      throw new RuntimeException("No chainId found in the batch");
-    }
-    return chainId;
-  }
 }

arithmetization/src/main/java/net/consensys/linea/zktracer/module/blockdata/BlockdataOperation.java

@@ -55,7 +55,7 @@ protected int computeLineCount() {
   }
 
   public void trace(
-      Trace trace, final int relBlock, final long firstBlockNumber, final long chainId) {
+      Trace trace, final int relBlock, final long firstBlockNumber, final BigInteger chainId) {
     for (short ct = 0; ct <= MAX_CT; ct++) {
       traceBlockConstant(trace, relBlock, firstBlockNumber);
       traceRowDependant(trace, ct, relBlock, chainId);
@@ -64,7 +64,7 @@ public void trace(
   }
 
   private void traceRowDependant(
-      Trace trace, final short ct, final int relBlock, final long chainId) {
+      Trace trace, final short ct, final int relBlock, final BigInteger chainId) {
     trace.ct(ct);
 
     Bytes32 data;
@@ -90,7 +90,7 @@ private void traceRowDependant(
         trace.inst(UnsignedByte.of(EVM_INST_GASLIMIT)).wcpFlag(false);
       }
       case 5 -> {
-        data = Bytes32.leftPad(Bytes.ofUnsignedLong(chainId));
+        data = Bytes32.leftPad(bigIntegerToBytes(chainId));
         trace.inst(UnsignedByte.of(EVM_INST_CHAINID)).wcpFlag(false);
       }
       case 6 -> {

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/Hub.java

@@ -26,13 +26,15 @@
 import static net.consensys.linea.zktracer.types.AddressUtils.effectiveToAddress;
 import static org.hyperledger.besu.evm.frame.MessageFrame.Type.*;
 
+import java.math.BigInteger;
 import java.nio.MappedByteBuffer;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Stream;
 
+import com.google.common.base.Preconditions;
 import lombok.Getter;
 import lombok.experimental.Accessors;
 import lombok.extern.slf4j.Slf4j;
@@ -196,6 +198,8 @@ public int lineCount() {
     return state.lineCounter().lineCount();
   }
 
+  @Getter private final BigInteger chainId;
+
   /** List of all modules of the ZK-evm */
   // stateless modules
   @Getter private final Wcp wcp = new Wcp();
@@ -394,12 +398,17 @@ public List<Module> getModulesToCount() {
         .toList();
   }
 
-  public Hub(final Address l2l1ContractAddress, final Bytes l2l1Topic) {
+  public Hub(
+      final Address l2l1ContractAddress,
+      final Bytes l2l1Topic,
+      final BigInteger nonnegativeChainId) {
+    Preconditions.checkState(nonnegativeChainId.signum() >= 0);
+    chainId = nonnegativeChainId;
     l2Block = new L2Block(l2l1ContractAddress, LogTopic.of(l2l1Topic));
     l2L1Logs = new L2L1Logs(l2Block);
     keccak = new Keccak(ecRecoverEffectiveCall, l2Block);
     shakiraData = new ShakiraData(wcp, sha256Blocks, keccak, ripemdBlocks);
-    blockdata = new Blockdata(wcp, txnData, rlpTxn);
+    blockdata = new Blockdata(wcp, txnData, rlpTxn, chainId);
     mmu = new Mmu(euc, wcp);
     mmio = new Mmio(mmu);
 

reference-tests/src/test/java/net/consensys/linea/BlockchainReferenceTestTools.java

@@ -20,6 +20,7 @@
 import static net.consensys.linea.ReferenceTestWatcher.JSON_INPUT_FILENAME;
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.math.BigInteger;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -178,7 +179,9 @@ public static void executeTest(final BlockchainReferenceTestCaseSpec spec) {
     final MutableBlockchain blockchain = spec.getBlockchain();
     final ProtocolContext context = spec.getProtocolContext();
 
-    final ZkTracer zkTracer = new ZkTracer(schedule.getChainId().get());
+    final BigInteger nonnegativeChainId = schedule.getChainId().get().abs();
+
+    final ZkTracer zkTracer = new ZkTracer(nonnegativeChainId);
     zkTracer.traceStartConflation(spec.getCandidateBlocks().length);
 
     for (var candidateBlock : spec.getCandidateBlocks()) {

testing/src/main/java/net/consensys/linea/testing/ToyExecutionEnvironmentV2.java

@@ -58,7 +58,7 @@ public class ToyExecutionEnvironmentV2 {
 
   @Builder.Default private final Consumer<ZkTracer> zkTracerValidator = x -> {};
 
-  private final ZkTracer tracer = new ZkTracer();
+  private final ZkTracer tracer = new ZkTracer(CHAIN_ID);
 
   public void run() {
     ProtocolSpec protocolSpec = ExecutionEnvironment.getProtocolSpec(CHAIN_ID);