reconnect MMIO constraints

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/fragment/imc/mmu/MmuCall.java

@@ -63,6 +63,7 @@
 import static net.consensys.linea.zktracer.module.hub.precompiles.ModexpMetadata.EBS_MIN_OFFSET;
 import static net.consensys.linea.zktracer.module.hub.precompiles.ModexpMetadata.MBS_MIN_OFFSET;
 import static net.consensys.linea.zktracer.types.Conversions.bigIntegerToBytes;
+import static net.consensys.linea.zktracer.types.Utils.leftPadTo;
 import static org.apache.tuweni.bytes.Bytes.minimalBytes;
 
 import java.util.Optional;
@@ -141,7 +142,7 @@ public void dontTraceMe() {
   }
 
   private MmuCall updateExoSum(final int exoValue) {
-    this.exoSum += exoValue;
+    exoSum += exoValue;
     return this;
   }
 
@@ -173,7 +174,6 @@ final MmuCall setEcData() {
     return this.exoIsEcData(true).updateExoSum(EXO_SUM_WEIGHT_ECDATA);
   }
 
-  // TODO: make the instruction an enum
   public MmuCall(final Hub hub, final int instruction) {
     hub.defers().scheduleForPostTransaction(this);
     this.instruction = instruction;
@@ -445,7 +445,7 @@ public static MmuCall forIdentityExtractCallData(
       final Hub hub, PrecompileSubsection precompileSubsection) {
 
     return new MmuCall(hub, MMU_INST_RAM_TO_RAM_SANS_PADDING)
-        .sourceId(precompileSubsection.callSection.hubStamp())
+        .sourceId(hub.currentFrame().contextNumber()) // called at ContextReEntry
         .sourceRamBytes(Optional.of(precompileSubsection.callerMemorySnapshot))
         .targetId(precompileSubsection.exoModuleOperationId())
         .targetRamBytes(Optional.of(Bytes.EMPTY))
@@ -706,10 +706,10 @@ public static MmuCall forModexpExtractBase(
       final Hub hub, final ModexpSubsection modexpSubsection, final ModexpMetadata modExpMetadata) {
     if (modExpMetadata.extractBase()) {
       return new MmuCall(hub, MMU_INST_MODEXP_DATA)
-          .sourceId(modexpSubsection.callSection.hubStamp())
+          .sourceId(hub.currentFrame().contextNumber()) // called at ContextReEntry
           .sourceRamBytes(Optional.of(modexpSubsection.callerMemorySnapshot))
           .targetId(modexpSubsection.exoModuleOperationId())
-          .exoBytes(Optional.of(modExpMetadata.base()))
+          .exoBytes(Optional.of(leftPadTo(modExpMetadata.base(), MODEXP_COMPONENT_BYTE_SIZE)))
           .sourceOffset(EWord.of(BASE_MIN_OFFSET))
           .size(modExpMetadata.bbs().toInt())
           .referenceOffset(modexpSubsection.callDataMemorySpan.offset())
@@ -728,10 +728,10 @@ public static MmuCall forModexpExtractExponent(
       final Hub hub, final ModexpSubsection modexpSubsection, final ModexpMetadata modExpMetadata) {
     if (modExpMetadata.extractExponent()) {
       return new MmuCall(hub, MMU_INST_MODEXP_DATA)
-          .sourceId(modexpSubsection.callSection.hubStamp())
+          .sourceId(hub.currentFrame().contextNumber()) // called at ContextReEntry
           .sourceRamBytes(Optional.of(modexpSubsection.callerMemorySnapshot))
           .targetId(modexpSubsection.exoModuleOperationId())
-          .exoBytes(Optional.of(modExpMetadata.exp()))
+          .exoBytes(Optional.of(leftPadTo(modExpMetadata.exp(), MODEXP_COMPONENT_BYTE_SIZE)))
           .sourceOffset(EWord.of(BASE_MIN_OFFSET + modExpMetadata.bbs().toInt()))
           .size(modExpMetadata.ebs().toInt())
           .referenceOffset(modexpSubsection.callDataMemorySpan.offset())
@@ -749,10 +749,10 @@ public static MmuCall forModexpExtractExponent(
   public static MmuCall forModexpExtractModulus(
       final Hub hub, final ModexpSubsection modexpSubsection, final ModexpMetadata modExpMetadata) {
     return new MmuCall(hub, MMU_INST_MODEXP_DATA)
-        .sourceId(modexpSubsection.callSection.hubStamp())
+        .sourceId(hub.currentFrame().contextNumber()) // called at ContextReEntry
         .sourceRamBytes(Optional.of(modexpSubsection.callerMemorySnapshot))
         .targetId(modexpSubsection.exoModuleOperationId())
-        .exoBytes(Optional.of(modExpMetadata.mod()))
+        .exoBytes(Optional.of(leftPadTo(modExpMetadata.mod(), MODEXP_COMPONENT_BYTE_SIZE)))
         .sourceOffset(
             EWord.of(BASE_MIN_OFFSET + modExpMetadata.bbs().toInt() + modExpMetadata.ebs().toInt()))
         .size(modExpMetadata.mbs().toInt())
@@ -766,7 +766,7 @@ public static MmuCall forModexpFullResultCopy(
       final Hub hub, final ModexpSubsection modexpSubsection, final ModexpMetadata modExpMetadata) {
     return new MmuCall(hub, MMU_INST_EXO_TO_RAM_TRANSPLANTS)
         .sourceId(modexpSubsection.exoModuleOperationId())
-        .exoBytes(Optional.of(modexpSubsection.returnData()))
+        .exoBytes(Optional.of(leftPadTo(modexpSubsection.returnData(), MODEXP_COMPONENT_BYTE_SIZE)))
         .targetId(modexpSubsection.returnDataContextNumber())
         .targetRamBytes(Optional.of(Bytes.EMPTY))
         .size(MODEXP_COMPONENT_BYTE_SIZE)

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/fragment/imc/mmu/opcode/CodeCopy.java

@@ -38,7 +38,7 @@ public CodeCopy(final Hub hub) {
     this.hub = hub;
     this.contract = hub.currentFrame().metadata();
 
-    this.exoBytes(Optional.of(hub.romLex().getCodeByMetadata(contract)))
+    this.exoBytes(Optional.of(hub.currentFrame().code().bytecode()))
         .targetId(hub.currentFrame().contextNumber())
         .targetRamBytes(
             Optional.of(
@@ -54,6 +54,6 @@ public CodeCopy(final Hub hub) {
 
   @Override
   public int sourceId() {
-    return this.hub.romLex().getCodeFragmentIndexByMetadata(this.contract);
+    return hub.romLex().getCodeFragmentIndexByMetadata(contract);
   }
 }

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/section/call/precompileSubsection/ModexpSubsection.java

@@ -16,6 +16,7 @@
 package net.consensys.linea.zktracer.module.hub.section.call.precompileSubsection;
 
 import static com.google.common.base.Preconditions.*;
+import static net.consensys.linea.zktracer.module.blake2fmodexpdata.BlakeModexpDataOperation.MODEXP_COMPONENT_BYTE_SIZE;
 import static net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall.forModexpExtractBase;
 import static net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall.forModexpExtractBbs;
 import static net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall.forModexpExtractEbs;
@@ -29,7 +30,6 @@
 import static net.consensys.linea.zktracer.module.hub.fragment.imc.oob.precompiles.ModexpXbsCase.OOB_INST_MODEXP_EBS;
 import static net.consensys.linea.zktracer.module.hub.fragment.imc.oob.precompiles.ModexpXbsCase.OOB_INST_MODEXP_MBS;
 import static net.consensys.linea.zktracer.module.hub.fragment.scenario.PrecompileScenarioFragment.PrecompileScenario.PRC_FAILURE_KNOWN_TO_RAM;
-import static net.consensys.linea.zktracer.module.limits.precompiles.ModexpEffectiveCall.PROVER_MAX_INPUT_BYTE_SIZE;
 
 import java.math.BigInteger;
 
@@ -60,17 +60,17 @@ public ModexpSubsection(final Hub hub, final CallSection callSection) {
     if (modexpMetaData
                 .bbs()
                 .toUnsignedBigInteger()
-                .compareTo(BigInteger.valueOf(PROVER_MAX_INPUT_BYTE_SIZE))
+                .compareTo(BigInteger.valueOf(MODEXP_COMPONENT_BYTE_SIZE))
             >= 0
         || modexpMetaData
                 .mbs()
                 .toUnsignedBigInteger()
-                .compareTo(BigInteger.valueOf(PROVER_MAX_INPUT_BYTE_SIZE))
+                .compareTo(BigInteger.valueOf(MODEXP_COMPONENT_BYTE_SIZE))
             >= 0
         || modexpMetaData
                 .ebs()
                 .toUnsignedBigInteger()
-                .compareTo(BigInteger.valueOf(PROVER_MAX_INPUT_BYTE_SIZE))
+                .compareTo(BigInteger.valueOf(MODEXP_COMPONENT_BYTE_SIZE))
             >= 0) {
       hub.modexpEffectiveCall().addPrecompileLimit(Integer.MAX_VALUE);
       return;
@@ -144,7 +144,7 @@ public void resolveAtContextReEntry(Hub hub, CallFrame callFrame) {
     }
 
     modexpMetaData.rawResult(returnData);
-    hub.blakeModexpData().callModexp(modexpMetaData, this.exoModuleOperationId());
+    hub.blakeModexpData().callModexp(modexpMetaData, exoModuleOperationId());
 
     fragments().add(seventhImcFragment);
     if (modexpMetaData.extractModulus()) {

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/section/copy/CodeCopySection.java

@@ -28,9 +28,6 @@
 import net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall;
 import net.consensys.linea.zktracer.module.hub.section.TraceSection;
 import net.consensys.linea.zktracer.module.hub.signals.Exceptions;
-import org.hyperledger.besu.datatypes.Address;
-import org.hyperledger.besu.evm.account.Account;
-import org.hyperledger.besu.evm.frame.MessageFrame;
 
 public class CodeCopySection extends TraceSection {
 
@@ -66,19 +63,9 @@ public CodeCopySection(Hub hub) {
     this.addFragment(contextFragment);
 
     // Account row
-    final MessageFrame frame = hub.messageFrame();
-    final Address codeAddress = frame.getContractAddress();
-    final Account codeAccount = frame.getWorldUpdater().get(codeAddress);
-
-    final boolean warmth = frame.isAddressWarm(codeAddress);
-    checkArgument(warmth);
-
     final AccountSnapshot codeAccountSnapshot =
-        AccountSnapshot.fromAccount(
-            codeAccount,
-            warmth,
-            hub.deploymentNumberOf(codeAddress),
-            hub.deploymentStatusOf(codeAddress));
+        AccountSnapshot.canonical(hub, hub.messageFrame().getContractAddress());
+    checkArgument(codeAccountSnapshot.isWarm());
 
     final DomSubStampsSubFragment doingDomSubStamps =
         DomSubStampsSubFragment.standardDomSubStamps(this.hubStamp(), 0); // Specifics for CODECOPY
@@ -94,7 +81,7 @@ public CodeCopySection(Hub hub) {
 
     final boolean triggerMmu = mxpCall.mayTriggerNontrivialMmuOperation;
     if (triggerMmu) {
-      MmuCall mmuCall = MmuCall.codeCopy(hub);
+      final MmuCall mmuCall = MmuCall.codeCopy(hub);
       imcFragment.callMmu(mmuCall);
     }
   }

arithmetization/src/main/java/net/consensys/linea/zktracer/module/limits/precompiles/ModexpEffectiveCall.java

@@ -29,7 +29,6 @@
 @Accessors(fluent = true)
 public class ModexpEffectiveCall implements CountingOnlyModule {
   private final CountOnlyOperation counts = new CountOnlyOperation();
-  public static final int PROVER_MAX_INPUT_BYTE_SIZE = 4096 / 8;
 
   @Override
   public String moduleKey() {

arithmetization/src/main/java/net/consensys/linea/zktracer/module/mmio/Mmio.java

@@ -87,26 +87,23 @@ public void commit(List<MappedByteBuffer> buffers) {
     Trace trace = new Trace(buffers);
     int stamp = 0;
     for (MmuOperation mmuOperation : mmu.operations().getAll()) {
-      if (mmuOperation.traceMe()) {
-
-        final MmuData currentMmuData = mmuOperation.mmuData();
-
-        for (int currentMmioInstNumber = 0;
-            currentMmioInstNumber < currentMmuData.numberMmioInstructions();
-            currentMmioInstNumber++) {
-          stamp++;
-
-          final MmioInstructions mmioInstructions =
-              new MmioInstructions(currentMmuData, currentMmioInstNumber);
-          final MmioData mmioData =
-              mmioInstructions.compute(
-                  currentMmuData
-                      .mmuToMmioInstructions()
-                      .get(currentMmioInstNumber)
-                      .mmioInstruction());
-
-          trace(trace, mmioData, stamp);
-        }
+
+      final MmuData currentMmuData = mmuOperation.mmuData();
+
+      for (int currentMmioInstNumber = 0;
+          currentMmioInstNumber < currentMmuData.mmuToMmioInstructions().size();
+          currentMmioInstNumber++) {
+
+        final MmioInstructions mmioInstructions =
+            new MmioInstructions(currentMmuData, currentMmioInstNumber);
+        final MmioData mmioData =
+            mmioInstructions.compute(
+                currentMmuData
+                    .mmuToMmioInstructions()
+                    .get(currentMmioInstNumber)
+                    .mmioInstruction());
+
+        trace(trace, mmioData, ++stamp);
       }
     }
   }

arithmetization/src/main/java/net/consensys/linea/zktracer/module/mmu/Mmu.java

@@ -59,14 +59,10 @@ public void commit(List<MappedByteBuffer> buffers) {
 
     for (MmuOperation mmuOperation : operations.getAll()) {
       Preconditions.checkState(mmuOperation.traceMe(), "Cannot compute if traceMe is false");
-      if (mmuOperation.traceMe()) {
-        mmuOperation.getCFI();
-        mmuOperation.fillLimb();
+      mmuOperation.getCFI();
+      mmuOperation.fillLimb();
 
-        mmuStamp += 1;
-        mmuOperation.trace(mmuStamp, mmioStamp, trace);
-        mmioStamp += mmuOperation.mmuData().numberMmioInstructions();
-      }
+      mmioStamp = mmuOperation.trace(++mmuStamp, mmioStamp, trace);
     }
   }
 

arithmetization/src/main/java/net/consensys/linea/zktracer/module/mmu/MmuData.java

@@ -18,6 +18,9 @@
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_ANY_TO_RAM_WITH_PADDING;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_BLAKE;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_EXO_TO_RAM_TRANSPLANTS;
+import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_MODEXP_DATA;
+import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_MODEXP_ZERO;
+import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_RAM_TO_EXO_WITH_PADDING;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -26,7 +29,6 @@
 import lombok.Getter;
 import lombok.Setter;
 import lombok.experimental.Accessors;
-import net.consensys.linea.zktracer.module.constants.GlobalConstants;
 import net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall;
 import net.consensys.linea.zktracer.module.mmu.values.HubToMmuValues;
 import net.consensys.linea.zktracer.module.mmu.values.MmuEucCallRecord;
@@ -62,9 +64,9 @@ public class MmuData {
   private static final List<Integer> MMU_INST_EXO_IS_TARGET =
       List.of(
           MMU_INST_BLAKE,
-          GlobalConstants.MMU_INST_MODEXP_DATA,
-          GlobalConstants.MMU_INST_MODEXP_ZERO,
-          GlobalConstants.MMU_INST_RAM_TO_EXO_WITH_PADDING);
+          MMU_INST_MODEXP_DATA,
+          MMU_INST_MODEXP_ZERO,
+          MMU_INST_RAM_TO_EXO_WITH_PADDING);
 
   public MmuData(final MmuCall mmuCall) {
     this(

arithmetization/src/main/java/net/consensys/linea/zktracer/module/mmu/MmuOperation.java

@@ -20,6 +20,7 @@
 
 import static com.google.common.base.Preconditions.*;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.LLARGE;
+import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMIO_INST_LIMB_VANISHES;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_ANY_TO_RAM_WITH_PADDING;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_BLAKE;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_EXO_TO_RAM_TRANSPLANTS;
@@ -34,6 +35,7 @@
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_RIGHT_PADDED_WORD_EXTRACTION;
 import static net.consensys.linea.zktracer.module.mmio.MmioData.lineCountOfMmioInstruction;
 import static net.consensys.linea.zktracer.types.Bytecodes.readBytes;
+import static net.consensys.linea.zktracer.types.Bytecodes.readLimb;
 import static net.consensys.linea.zktracer.types.Conversions.*;
 
 import java.util.ArrayList;
@@ -93,7 +95,7 @@ public int mmioLineCount() {
     return mmioLineCount;
   }
 
-  void trace(final int mmuStamp, final int mmioStamp, Trace trace) {
+  int trace(final int mmuStamp, final int mmioStamp, Trace trace) {
 
     setInstructionFlag();
 
@@ -104,7 +106,7 @@ void trace(final int mmuStamp, final int mmioStamp, Trace trace) {
     tracePreprocessingRows(mmuData, mmuStamp, mmioStamp, trace);
 
     // Trace Micro Instructions Rows
-    traceMicroRows(mmuStamp, mmioStamp, trace);
+    return traceMicroRows(mmuStamp, mmioStamp, trace);
   }
 
   private void setInstructionFlag() {
@@ -159,24 +161,28 @@ public void fillLimb() {
       final boolean exoIsTarget = mmuData.exoLimbIsTarget();
       checkArgument(exoIsSource == !exoIsTarget, "ExoLimb is either the source or the target");
 
-      for (MmuToMmioInstruction mmioInst : this.mmuData.mmuToMmioInstructions()) {
-        final int offset =
-            (int)
-                (exoIsSource
-                    ? mmioInst.sourceLimbOffset() * LLARGE + mmioInst.sourceByteOffset()
-                    : mmioInst.targetLimbOffset() * LLARGE + mmioInst.targetByteOffset());
-        final int sizeToExtract = mmioInst.size() == 0 ? LLARGE : mmioInst.size();
-        final Bytes16 exoLimb = readBytes(mmuData.exoBytes(), offset, sizeToExtract);
-        mmioInst.limb(exoLimb);
+      for (MmuToMmioInstruction mmioInst : mmuData.mmuToMmioInstructions()) {
+
+        // Limb remains zero for LIMB_VANISHES instructions
+        if (mmioInst.mmioInstruction() != MMIO_INST_LIMB_VANISHES) {
+
+          if (exoIsSource) {
+            mmioInst.limb(readLimb(mmuData.exoBytes(), mmioInst.sourceLimbOffset()));
+          } else {
+            final int offset =
+                (int) mmioInst.targetLimbOffset() * LLARGE + mmioInst.targetByteOffset();
+            final int sizeToExtract = mmioInst.size() == 0 ? LLARGE : mmioInst.size();
+            final int exoByteOffset = mmioInst.targetByteOffset();
+            final Bytes16 exoLimb =
+                readBytes(mmuData.exoBytes(), offset, sizeToExtract, exoByteOffset);
+            ;
+            mmioInst.limb(exoLimb);
+          }
+        }
       }
     }
   }
 
-  private boolean exoLimbIsSource() {
-    return List.of(MMU_INST_ANY_TO_RAM_WITH_PADDING, MMU_INST_EXO_TO_RAM_TRANSPLANTS)
-        .contains(this.mmuData.hubToMmuValues().mmuInstruction());
-  }
-
   private void traceFillMmuInstructionFlag(Trace trace) {
     trace
         .isMload(isMload)
@@ -195,7 +201,7 @@ private void traceFillMmuInstructionFlag(Trace trace) {
   }
 
   private void traceOutAndBin(Trace trace) {
-    MmuOutAndBinValues mmuOutAndBinRecord = mmuData.outAndBinValues();
+    final MmuOutAndBinValues mmuOutAndBinRecord = mmuData.outAndBinValues();
 
     trace
         .out1(Bytes.minimalBytes(mmuOutAndBinRecord.out1()))
@@ -337,15 +343,15 @@ private List<RowTypeRecord> generateRowTypeList() {
     return output;
   }
 
-  private void traceMicroRows(final long mmuStamp, int mmioStamp, Trace trace) {
+  private int traceMicroRows(final long mmuStamp, int mmioStamp, Trace trace) {
     final List<RowTypeRecord> rowType = generateRowTypeList();
     final HubToMmuValues mmuHubInput = mmuData.hubToMmuValues();
 
     final boolean successBit = mmuHubInput.successBit();
 
     final MmuToMmioConstantValues mmioConstantValues = mmuData.mmuToMmioConstantValues();
 
-    for (int i = 0; i < mmuData().numberMmioInstructions(); i++) {
+    for (int i = 0; i < mmuData().mmuToMmioInstructions().size(); i++) {
       mmioStamp += 1;
       traceFillMmuInstructionFlag(trace);
       traceOutAndBin(trace);
@@ -390,5 +396,6 @@ private void traceMicroRows(final long mmuStamp, int mmioStamp, Trace trace) {
           .pMicroTotalSize(Bytes.minimalBytes(mmioConstantValues.totalSize()))
           .fillAndValidateRow();
     }
+    return mmioStamp;
   }
 }