debug memory replay tests

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/fragment/common/CommonFragment.java

@@ -51,12 +51,20 @@ public final class CommonFragment implements TraceFragment {
   private final CommonFragmentValues commonFragmentValues;
   private final int nonStackRowsCounter;
   private final boolean twoLineInstructionCounter;
+  private final int mmuStamp;
+  private final int mxpStamp;
 
   public CommonFragment(
-      CommonFragmentValues commonValues, int stackLineCounter, int nonStackLineCounter) {
+      CommonFragmentValues commonValues,
+      int stackLineCounter,
+      int nonStackLineCounter,
+      int mmuStamp,
+      int mxpStamp) {
     this.commonFragmentValues = commonValues;
     this.twoLineInstructionCounter = stackLineCounter == 1;
     this.nonStackRowsCounter = nonStackLineCounter;
+    this.mmuStamp = mmuStamp;
+    this.mxpStamp = mxpStamp;
   }
 
   private boolean isUnexceptional() {
@@ -87,8 +95,8 @@ public Trace trace(Trace trace) {
         .contextMayChange(commonFragmentValues.contextMayChange)
         .exceptionAhoy(Exceptions.any(commonFragmentValues.exceptions) && isExec)
         .logInfoStamp(commonFragmentValues.logStamp)
-        .mmuStamp(commonFragmentValues.stamps.mmu())
-        .mxpStamp(commonFragmentValues.stamps.mxp())
+        .mmuStamp(mmuStamp)
+        .mxpStamp(mxpStamp)
         // nontrivial dom / sub are traced in storage or account fragments only
         .contextNumber(isExec ? frame.contextNumber() : 0)
         .contextNumberNew(commonFragmentValues.contextNumberNew)

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/fragment/common/CommonFragmentValues.java

@@ -47,7 +47,7 @@ public class CommonFragmentValues {
   public final HubProcessingPhase hubProcessingPhase;
   public final int hubStamp;
   public final CallStack callStack;
-  public final State.TxState.Stamps stamps; // for MMU and MXP stamps
+  public final State.TxState.Stamps stamps;
   @Setter public int logStamp = -1;
   @Getter final CallFrame callFrame;
   public final short exceptions;

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

@@ -0,0 +1,58 @@
+/*
+ * Copyright ConsenSys Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.opcode;
+
+import static net.consensys.linea.zktracer.module.constants.GlobalConstants.MMU_INST_RIGHT_PADDED_WORD_EXTRACTION;
+
+import java.util.Optional;
+
+import net.consensys.linea.zktracer.module.hub.Hub;
+import net.consensys.linea.zktracer.module.hub.defer.PostOpcodeDefer;
+import net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall;
+import net.consensys.linea.zktracer.runtime.callstack.CallFrame;
+import net.consensys.linea.zktracer.types.EWord;
+import org.apache.tuweni.bytes.Bytes;
+import org.hyperledger.besu.evm.frame.MessageFrame;
+import org.hyperledger.besu.evm.operation.Operation;
+
+public class CallDataLoad extends MmuCall implements PostOpcodeDefer {
+
+  public CallDataLoad(final Hub hub) {
+    super(hub, MMU_INST_RIGHT_PADDED_WORD_EXTRACTION);
+    hub.defers().scheduleForPostExecution(this);
+
+    final CallFrame currentFrame = hub.currentFrame();
+    final long callDataSize = currentFrame.callDataInfo().memorySpan().length();
+    final long callDataOffset = currentFrame.callDataInfo().memorySpan().offset();
+    final EWord sourceOffset = EWord.of(currentFrame.frame().getStackItem(0));
+    final long callDataCN = currentFrame.callDataInfo().callDataContextNumber();
+    final Bytes sourceBytes = hub.callStack().getFullMemoryOfCaller(hub);
+
+    this.sourceId((int) callDataCN)
+        .sourceRamBytes(Optional.of(sourceBytes))
+        .sourceOffset(sourceOffset)
+        .referenceOffset(callDataOffset)
+        .referenceSize(callDataSize);
+  }
+
+  @Override
+  public void resolvePostExecution(
+      Hub hub, MessageFrame frame, Operation.OperationResult operationResult) {
+    final EWord stack = EWord.of(frame.getStackItem(0));
+    this.limb1(stack.hi());
+    this.limb2(stack.lo());
+  }
+}

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

@@ -39,7 +39,7 @@ public class Create2 extends MmuCall implements RomLexDefer {
   private final Hub hub;
   private ContractMetadata contract;
 
-  public Create2(final Hub hub, final boolean failedCreate) {
+  public Create2(final Hub hub, final Bytes create2initCode, final boolean failedCreate) {
     super(hub, MMU_INST_RAM_TO_EXO_WITH_PADDING);
     this.hub = hub;
     this.hub.romLex().createDefers().register(this);
@@ -55,6 +55,7 @@ public Create2(final Hub hub, final boolean failedCreate) {
                     currentFrame.frame(),
                     MemorySpan.fromStartLength(clampedToLong(sourceOffset), size))))
         .auxId(newIdentifierFromStamp(hub.stamp()))
+        .exoBytes(Optional.of(create2initCode))
         .sourceOffset(sourceOffset)
         .size(size)
         .referenceSize(size)
@@ -70,11 +71,6 @@ public int targetId() {
     return exoIsRom ? hub.romLex().getCodeFragmentIndexByMetadata(contract) : 0;
   }
 
-  @Override
-  public Optional<Bytes> exoBytes() {
-    return exoIsRom ? Optional.of(hub.romLex().getCodeByMetadata(contract)) : Optional.empty();
-  }
-
   @Override
   public void updateContractMetadata(ContractMetadata metadata) {
     contract = metadata;

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

@@ -16,28 +16,16 @@
 package net.consensys.linea.zktracer.module.hub.section;
 
 import static com.google.common.base.Preconditions.checkArgument;
-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;
 
 import net.consensys.linea.zktracer.module.hub.Hub;
 import net.consensys.linea.zktracer.module.hub.fragment.ContextFragment;
 import net.consensys.linea.zktracer.module.hub.fragment.imc.ImcFragment;
 import net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall;
+import net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.opcode.CallDataLoad;
 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 {
 
@@ -55,39 +43,8 @@ public CallDataLoadSection(Hub hub) {
 
     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(
-                Bytes.wrap(
-                    Arrays.copyOfRange(
-                        hub.currentFrame().callDataInfo().data().toArray(),
-                        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());
-
-        imcFragment.callMmu(call);
+        final CallDataLoad mmuCall = (CallDataLoad) MmuCall.callDataLoad(hub);
+        imcFragment.callMmu(mmuCall);
       }
     } else {
       // Sanity check

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

@@ -181,23 +181,24 @@ public CreateSection(Hub hub) {
     final boolean failedCreate = createdAddressHasNonZeroNonce || createdAddressHasNonEmptyCode;
     final boolean emptyInitCode = hub.transients().op().initCodeSegment().isEmpty();
 
-    final long offset = Words.clampedToLong(hub.messageFrame().getStackItem(1));
-    final long size = Words.clampedToLong(hub.messageFrame().getStackItem(2));
+    final long offset = Words.clampedToLong(messageFrame.getStackItem(1));
+    final long size = Words.clampedToLong(messageFrame.getStackItem(2));
 
     // Trigger MMU & SHAKIRA to hash the (non-empty) InitCode of CREATE2 - even for failed CREATE2
     if (hub.opCode() == CREATE2 && !emptyInitCode) {
-      final MmuCall mmuCall = MmuCall.create2(hub, failedCreate);
+      final Bytes create2InitCode = messageFrame.shadowReadMemory(offset, size);
+
+      final MmuCall mmuCall = MmuCall.create2(hub, create2InitCode, failedCreate);
       imcFragment.callMmu(mmuCall);
 
-      final Bytes create2InitCode = messageFrame.shadowReadMemory(offset, size);
       final ShakiraDataOperation shakiraDataOperation =
           new ShakiraDataOperation(hub.stamp(), create2InitCode);
       hub.shakiraData().call(shakiraDataOperation);
 
       triggerHashInfo(shakiraDataOperation.result());
     }
 
-    value = failedCreate ? Wei.ZERO : Wei.of(UInt256.fromBytes(hub.messageFrame().getStackItem(0)));
+    value = failedCreate ? Wei.ZERO : Wei.of(UInt256.fromBytes(messageFrame.getStackItem(0)));
 
     if (failedCreate || emptyInitCode) {
       finalContextFragment = ContextFragment.nonExecutionProvidesEmptyReturnData(hub);
@@ -219,16 +220,14 @@ public CreateSection(Hub hub) {
 
     // Finally, non-exceptional, non-aborting, non-failing, non-emptyInitCode create
     hub.defers()
-        .scheduleForContextReEntry(
-            this, hub.currentFrame()); // To get the success bit of the CREATE(2)
+        .scheduleForContextReEntry(this, callFrame); // To get the success bit of the CREATE(2)
 
     requiresRomLex = true;
     hub.romLex().callRomLex(messageFrame);
     hub.transients()
         .conflation()
         .deploymentInfo()
-        .newDeploymentWithExecutionAt(
-            createeAddress, hub.messageFrame().shadowReadMemory(offset, size));
+        .newDeploymentWithExecutionAt(createeAddress, messageFrame.shadowReadMemory(offset, size));
 
     // Note: the case CREATE2 has been set before, we need to do it even in the failure case
     if (hub.opCode() == CREATE) {

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

@@ -30,6 +30,7 @@
 import net.consensys.linea.zktracer.module.hub.fragment.imc.mmu.MmuCall;
 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.types.EWord;
 import net.consensys.linea.zktracer.types.MemorySpan;
 import org.apache.tuweni.bytes.Bytes;
@@ -62,47 +63,45 @@ public StackRamSection(Hub hub) {
     // the unexceptional case
     checkArgument(Exceptions.none(exceptions));
 
-    final EWord offset = EWord.of(hub.currentFrame().frame().getStackItem(0));
+    final CallFrame currentFrame = hub.currentFrame();
+    final EWord offset = EWord.of(currentFrame.frame().getStackItem(0));
     final long longOffset = Words.clampedToLong(offset);
     final Bytes currentRam =
         extractContiguousLimbsFromMemory(
-            hub.currentFrame().frame(), new MemorySpan(longOffset, WORD_SIZE));
-    final int currentContextNumber = hub.currentFrame().contextNumber();
+            currentFrame.frame(), new MemorySpan(longOffset, WORD_SIZE));
+    final int currentContextNumber = currentFrame.contextNumber();
     final EWord value =
         instruction.equals(OpCode.MLOAD)
-            ? EWord.of(hub.messageFrame().shadowReadMemory(Words.clampedToLong(offset), WORD_SIZE))
-            : EWord.of(hub.currentFrame().frame().getStackItem(1));
+            ? EWord.of(currentFrame.frame().shadowReadMemory(longOffset, WORD_SIZE))
+            : EWord.of(currentFrame.frame().getStackItem(1));
 
     MmuCall mmuCall;
 
     switch (instruction) {
-      case MSTORE -> {
-        mmuCall =
-            new MmuCall(hub, MMU_INST_MSTORE)
-                .targetId(currentContextNumber)
-                .targetOffset(offset)
-                .limb1(value.hi())
-                .limb2(value.lo())
-                .targetRamBytes(Optional.of(currentRam));
-      }
-      case MSTORE8 -> {
-        mmuCall =
-            new MmuCall(hub, MMU_INST_MSTORE8)
-                .targetId(currentContextNumber)
-                .targetOffset(offset)
-                .limb1(value.hi())
-                .limb2(value.lo())
-                .targetRamBytes(Optional.of(currentRam));
-      }
-      case MLOAD -> {
-        mmuCall =
-            new MmuCall(hub, MMU_INST_MLOAD)
-                .sourceId(currentContextNumber)
-                .sourceOffset(offset)
-                .limb1(value.hi())
-                .limb2(value.lo())
-                .sourceRamBytes(Optional.of(currentRam));
-      }
+      case MSTORE -> mmuCall =
+          new MmuCall(hub, MMU_INST_MSTORE)
+              .targetId(currentContextNumber)
+              .targetOffset(offset)
+              .limb1(value.hi())
+              .limb2(value.lo())
+              .targetRamBytes(Optional.of(currentRam));
+
+      case MSTORE8 -> mmuCall =
+          new MmuCall(hub, MMU_INST_MSTORE8)
+              .targetId(currentContextNumber)
+              .targetOffset(offset)
+              .limb1(value.hi())
+              .limb2(value.lo())
+              .targetRamBytes(Optional.of(currentRam));
+
+      case MLOAD -> mmuCall =
+          new MmuCall(hub, MMU_INST_MLOAD)
+              .sourceId(currentContextNumber)
+              .sourceOffset(offset)
+              .limb1(value.hi())
+              .limb2(value.lo())
+              .sourceRamBytes(Optional.of(currentRam));
+
       default -> throw new IllegalStateException("Not a STACK_RAM instruction");
     }
 

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

@@ -224,7 +224,12 @@ public void trace(Trace hubTrace) {
 
       specificFragment.trace(hubTrace);
       final CommonFragment commonFragment =
-          new CommonFragment(commonValues, stackLineCounter, nonStackLineCounter);
+          new CommonFragment(
+              commonValues,
+              stackLineCounter,
+              nonStackLineCounter,
+              hub.state.stamps().mmu(),
+              hub.state.stamps().mxp());
       commonFragment.trace(hubTrace);
       hubTrace.fillAndValidateRow();
     }

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

@@ -96,6 +96,8 @@ public PrecompileSubsection(final Hub hub, final CallSection callSection) {
     this.callSection = callSection;
     fragments = new ArrayList<>(maxNumberOfLines());
 
+    final MessageFrame messageFrame = hub.messageFrame();
+
     hub.defers().scheduleForImmediateContextEntry(this); // gas & input data, ...
     hub.defers().scheduleForContextExit(this, hub.callStack().futureId());
     hub.defers().scheduleForContextReEntry(this, hub.currentFrame()); // success bit & return data
@@ -114,16 +116,15 @@ public PrecompileSubsection(final Hub hub, final CallSection callSection) {
     final long offset =
         Words.clampedToLong(
             opCode.callCanTransferValue()
-                ? hub.messageFrame().getStackItem(3)
-                : hub.messageFrame().getStackItem(2));
+                ? messageFrame.getStackItem(3)
+                : messageFrame.getStackItem(2));
     final long length =
         Words.clampedToLong(
             opCode.callCanTransferValue()
-                ? hub.messageFrame().getStackItem(4)
-                : hub.messageFrame().getStackItem(3));
+                ? messageFrame.getStackItem(4)
+                : messageFrame.getStackItem(3));
     callDataMemorySpan = new MemorySpan(offset, length);
-    callerMemorySnapshot =
-        extractContiguousLimbsFromMemory(hub.currentFrame().frame(), callDataMemorySpan);
+    callerMemorySnapshot = extractContiguousLimbsFromMemory(messageFrame, callDataMemorySpan);
     final int lengthToExtract =
         (int) Math.min(length, Math.max(callerMemorySnapshot.size() - offset, 0));
     callData = rightPadTo(callerMemorySnapshot.slice((int) offset, lengthToExtract), (int) length);
@@ -150,10 +151,13 @@ public void resolveUponContextExit(Hub hub, CallFrame callFrame) {
 
   @Override
   public void resolveAtContextReEntry(Hub hub, CallFrame frame) {
-    callSuccess = bytesToBoolean(hub.messageFrame().getStackItem(0));
+    callSuccess = bytesToBoolean(frame.frame().getStackItem(0));
     returnData = frame.frame().getReturnData();
 
-    frame.returnDataContextNumber(exoModuleOperationId());
+    final int returnerCn = exoModuleOperationId();
+    final CallFrame returnerFrame = hub.callStack().getByContextNumber(returnerCn);
+    returnerFrame.returnData(returnData);
+    frame.returnDataContextNumber(returnerCn);
     frame.returnDataSpan(new MemorySpan(0, returnData.size()));
 
     if (callSuccess) {