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))
@@ -460,7 +460,7 @@ public static MmuCall forIdentityReturnData(
     return new MmuCall(hub, MMU_INST_RAM_TO_RAM_SANS_PADDING)
         .sourceId(precompileSubsection.exoModuleOperationId())
         .sourceRamBytes(Optional.of(precompileSubsection.returnData()))
-        .targetId(precompileSubsection.callSection.hubStamp())
+        .targetId(hub.currentFrame().contextNumber())
         .targetRamBytes(Optional.of(precompileSubsection.callerMemorySnapshot))
         .sourceOffset(EWord.ZERO)
         .targetOffset(EWord.of(precompileSubsection.parentReturnDataTarget.offset()))
@@ -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/fragment/imc/mmu/opcode/Create.java

@@ -24,6 +24,7 @@
 import net.consensys.linea.zktracer.module.romlex.ContractMetadata;
 import net.consensys.linea.zktracer.module.romlex.RomLexDefer;
 import net.consensys.linea.zktracer.types.EWord;
+import org.apache.tuweni.bytes.Bytes;
 import org.hyperledger.besu.evm.internal.Words;
 
 /**
@@ -45,7 +46,6 @@ public Create(final Hub hub) {
                 hub.currentFrame()
                     .frame()
                     .shadowReadMemory(0, hub.currentFrame().frame().memoryByteSize())))
-        .exoBytes(Optional.of(hub.romLex().getCodeByMetadata(contract)))
         .sourceOffset(EWord.of(hub.messageFrame().getStackItem(1)))
         .size(Words.clampedToLong(hub.messageFrame().getStackItem(2)))
         .referenceSize(Words.clampedToLong(hub.messageFrame().getStackItem(2)))
@@ -54,11 +54,16 @@ public Create(final Hub hub) {
 
   @Override
   public int targetId() {
-    return hub.romLex().getCodeFragmentIndexByMetadata(this.contract);
+    return hub.romLex().getCodeFragmentIndexByMetadata(contract);
+  }
+
+  @Override
+  public Optional<Bytes> exoBytes() {
+    return Optional.of(hub.romLex().getCodeByMetadata(contract));
   }
 
   @Override
   public void updateContractMetadata(ContractMetadata metadata) {
-    this.contract = metadata;
+    contract = metadata;
   }
 }

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

@@ -24,6 +24,7 @@
 import net.consensys.linea.zktracer.module.romlex.ContractMetadata;
 import net.consensys.linea.zktracer.module.romlex.RomLexDefer;
 import net.consensys.linea.zktracer.types.EWord;
+import org.apache.tuweni.bytes.Bytes;
 import org.hyperledger.besu.evm.internal.Words;
 
 /**
@@ -45,7 +46,6 @@ public Create2(final Hub hub, final boolean failedCreate) {
                 hub.currentFrame()
                     .frame()
                     .shadowReadMemory(0, hub.currentFrame().frame().memoryByteSize())))
-        .exoBytes(Optional.of(hub.romLex().getCodeByMetadata(contract)))
         .auxId(hub.state().stamps().hub())
         .sourceOffset(EWord.of(hub.messageFrame().getStackItem(1)))
         .size(Words.clampedToLong(hub.messageFrame().getStackItem(2)))
@@ -62,8 +62,13 @@ public int targetId() {
     return exoIsRom ? hub.romLex().getCodeFragmentIndexByMetadata(contract) : 0;
   }
 
+  @Override
+  public Optional<Bytes> exoBytes() {
+    return Optional.of(hub.romLex().getCodeByMetadata(contract));
+  }
+
   @Override
   public void updateContractMetadata(ContractMetadata metadata) {
-    this.contract = metadata;
+    contract = metadata;
   }
 }

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

@@ -24,6 +24,7 @@
 import net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall;
 import net.consensys.linea.zktracer.module.romlex.ContractMetadata;
 import net.consensys.linea.zktracer.types.EWord;
+import org.apache.tuweni.bytes.Bytes;
 import org.hyperledger.besu.datatypes.Address;
 import org.hyperledger.besu.evm.internal.Words;
 import org.hyperledger.besu.evm.worldstate.WorldView;
@@ -44,8 +45,7 @@ public ExtCodeCopy(final Hub hub) {
     final Address foreignCodeAddress = Words.toAddress(hub.messageFrame().getStackItem(0));
     this.contract = ContractMetadata.canonical(hub, foreignCodeAddress);
 
-    this.exoBytes(Optional.of(hub.romLex().getCodeByMetadata(contract)))
-        .targetId(hub.currentFrame().contextNumber())
+    this.targetId(hub.currentFrame().contextNumber())
         .targetRamBytes(
             Optional.of(
                 hub.currentFrame()
@@ -57,12 +57,24 @@ public ExtCodeCopy(final Hub hub) {
         .setRom();
   }
 
+  @Override
+  public Optional<Bytes> exoBytes() {
+    try {
+      return Optional.of(hub.romLex().getCodeByMetadata(contract));
+    } catch (Exception ignored) {
+      // Can be empty Bytes in case the ext account is empty. In this case, no associated CFI
+      return Optional.of(Bytes.EMPTY);
+    }
+  }
+
   @Override
   public long referenceSize() {
-    return hub.romLex()
-        .getChunkByMetadata(contract)
-        .map(chunk -> chunk.byteCode().size())
-        .orElse(0);
+    try {
+      return (hub.romLex().getCodeByMetadata(contract).size());
+    } catch (Exception ignored) {
+      // Can be 0 in case the ext account is empty. In this case, no associated CFI
+      return 0;
+    }
   }
 
   @Override

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

@@ -26,6 +26,7 @@
 import net.consensys.linea.zktracer.module.romlex.ContractMetadata;
 import net.consensys.linea.zktracer.module.shakiradata.ShakiraDataOperation;
 import net.consensys.linea.zktracer.types.EWord;
+import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
 import org.hyperledger.besu.datatypes.Address;
 import org.hyperledger.besu.evm.internal.Words;
@@ -37,7 +38,7 @@
 @Accessors(fluent = true)
 public class ReturnFromDeploymentMmuCall extends MmuCall {
   private final Hub hub;
-  private final ContractMetadata contract;
+  private ContractMetadata contract;
   @Getter private final Bytes32 hashResult;
 
   public ReturnFromDeploymentMmuCall(final Hub hub) {
@@ -47,10 +48,10 @@ public ReturnFromDeploymentMmuCall(final Hub hub) {
 
     final Address contractAddress = hub.messageFrame().getContractAddress();
     final int depNumber = hub.deploymentNumberOf(contractAddress);
-    contract = ContractMetadata.underDeployment(contractAddress, depNumber);
+    contract = ContractMetadata.make(contractAddress, depNumber, false);
 
     final ShakiraDataOperation shakiraDataOperation =
-        new ShakiraDataOperation(hub.stamp(), hub.romLex().getCodeByMetadata(contract));
+        new ShakiraDataOperation(hub.stamp(), hub.romLex().byteCode());
     hub.shakiraData().call(shakiraDataOperation);
 
     hashResult = shakiraDataOperation.result();
@@ -61,7 +62,6 @@ public ReturnFromDeploymentMmuCall(final Hub hub) {
                 hub.currentFrame()
                     .frame()
                     .shadowReadMemory(0, hub.currentFrame().frame().memoryByteSize())))
-        .exoBytes(Optional.of(hub.romLex().getCodeByMetadata(contract)))
         .auxId(hub.state().stamps().hub())
         .sourceOffset(EWord.of(hub.messageFrame().getStackItem(0)))
         .size(Words.clampedToLong(hub.messageFrame().getStackItem(1)))
@@ -72,6 +72,11 @@ public ReturnFromDeploymentMmuCall(final Hub hub) {
 
   @Override
   public int targetId() {
-    return this.hub.romLex().getCodeFragmentIndexByMetadata(this.contract);
+    return hub.romLex().getCodeFragmentIndexByMetadata(contract);
+  }
+
+  @Override
+  public Optional<Bytes> exoBytes() {
+    return Optional.of(hub.romLex().getCodeByMetadata(contract));
   }
 }

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

@@ -19,6 +19,8 @@
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_RIGHT_PADDED_WORD_EXTRACTION;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.WORD_SIZE;
 import static net.consensys.linea.zktracer.module.hub.fragment.ContextFragment.readCurrentContextData;
+import static net.consensys.linea.zktracer.runtime.callstack.CallFrame.extractContiguousLimbsFromMemory;
+import static org.hyperledger.besu.evm.internal.Words.clampedToLong;
 
 import java.util.Arrays;
 import java.util.Optional;
@@ -30,40 +32,33 @@
 import net.consensys.linea.zktracer.module.hub.fragment.imc.oob.opcodes.CallDataLoadOobCall;
 import net.consensys.linea.zktracer.module.hub.signals.Exceptions;
 import net.consensys.linea.zktracer.opcode.OpCode;
+import net.consensys.linea.zktracer.runtime.callstack.CallFrame;
+import net.consensys.linea.zktracer.runtime.callstack.CallFrameType;
 import net.consensys.linea.zktracer.types.EWord;
+import net.consensys.linea.zktracer.types.MemorySpan;
 import org.apache.tuweni.bytes.Bytes;
 import org.hyperledger.besu.evm.internal.Words;
 
 public class CallDataLoadSection extends TraceSection {
-  final short exception;
-  final Bytes callDataRam;
-  final int currentContextNumber;
-  final long callDataCN;
-  final EWord sourceOffset;
-
-  long callDataOffset = -1;
-  long callDataSize = -1;
 
   public CallDataLoadSection(Hub hub) {
     super(hub, (short) (hub.opCode().equals(OpCode.CALLDATALOAD) ? 4 : 3));
     this.addStack(hub);
 
-    this.exception = hub.pch().exceptions();
-    this.currentContextNumber = hub.currentFrame().contextNumber();
-    this.callDataSize = hub.currentFrame().callDataInfo().memorySpan().length();
-    this.callDataOffset = hub.currentFrame().callDataInfo().memorySpan().offset();
-    this.sourceOffset = EWord.of(hub.currentFrame().frame().getStackItem(0));
-    this.callDataCN = hub.currentFrame().callDataInfo().callDataContextNumber();
-    this.callDataRam = hub.currentFrame().callDataInfo().data();
+    final short exception = hub.pch().exceptions();
 
     final ImcFragment imcFragment = ImcFragment.empty(hub);
+    this.addFragment(imcFragment);
 
     final CallDataLoadOobCall oobCall = new CallDataLoadOobCall();
     imcFragment.callOob(oobCall);
 
     if (Exceptions.none(exception)) {
-
       if (!oobCall.isCdlOutOfBounds()) {
+        final long callDataSize = hub.currentFrame().callDataInfo().memorySpan().length();
+        final long callDataOffset = hub.currentFrame().callDataInfo().memorySpan().offset();
+        final EWord sourceOffset = EWord.of(hub.currentFrame().frame().getStackItem(0));
+        final long callDataCN = hub.currentFrame().callDataInfo().callDataContextNumber();
 
         final EWord read =
             EWord.of(
@@ -73,24 +68,32 @@ public CallDataLoadSection(Hub hub) {
                         Words.clampedToInt(sourceOffset),
                         Words.clampedToInt(sourceOffset) + WORD_SIZE)));
 
+        final CallFrame callDataCallFrame = hub.callStack().getByContextNumber(callDataCN);
+
         final MmuCall call =
             new MmuCall(hub, MMU_INST_RIGHT_PADDED_WORD_EXTRACTION)
                 .sourceId((int) callDataCN)
+                .sourceRamBytes(
+                    Optional.of(
+                        callDataCallFrame.type() == CallFrameType.TRANSACTION_CALL_DATA_HOLDER
+                            ? callDataCallFrame.callDataInfo().data()
+                            : extractContiguousLimbsFromMemory(
+                                callDataCallFrame.frame(),
+                                MemorySpan.fromStartLength(
+                                    clampedToLong(sourceOffset) + callDataOffset, WORD_SIZE))))
                 .sourceOffset(sourceOffset)
                 .referenceOffset(callDataOffset)
                 .referenceSize(callDataSize)
                 .limb1(read.hi())
-                .limb2(read.lo())
-                .sourceRamBytes(Optional.of(callDataRam));
+                .limb2(read.lo());
 
         imcFragment.callMmu(call);
       }
     } else {
+      // Sanity check
       checkArgument(Exceptions.outOfGasException(exception));
     }
 
-    this.addFragment(imcFragment);
-
     final ContextFragment context = readCurrentContextData(hub);
     this.addFragment(context);
   }