HUB debugging continued (#1481)

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

@@ -57,11 +57,14 @@ public class AccountSnapshot {
    * @return
    */
   public static AccountSnapshot canonical(Hub hub, Address address) {
-    return fromArguments(
-        hub.messageFrame().getWorldUpdater(),
-        address,
-        hub.transients.conflation().deploymentInfo(),
-        isAddressWarm(hub.messageFrame(), address));
+    AccountSnapshot canonicalSnapshot =
+        fromArguments(
+            hub.messageFrame().getWorldUpdater(),
+            address,
+            hub.transients.conflation().deploymentInfo(),
+            isAddressWarm(hub.messageFrame(), address));
+
+    return canonicalSnapshot;
   }
 
   public static AccountSnapshot canonical(
@@ -76,13 +79,17 @@ private static AccountSnapshot fromArguments(
       final boolean warmth) {
 
     final Account account = worldView.get(address);
+    Bytecode bytecode =
+        deploymentInfo.getDeploymentStatus(address)
+            ? new Bytecode(deploymentInfo.getInitializationCode(address))
+            : (account == null) ? new Bytecode(Bytes.EMPTY) : new Bytecode(account.getCode());
     if (account != null) {
       return new AccountSnapshot(
           account.getAddress(),
           account.getNonce(),
           account.getBalance(),
           warmth,
-          new Bytecode(account.getCode()),
+          bytecode,
           deploymentInfo.deploymentNumber(address),
           deploymentInfo.getDeploymentStatus(address));
     } else {
@@ -91,7 +98,7 @@ private static AccountSnapshot fromArguments(
           0,
           Wei.ZERO,
           warmth,
-          new Bytecode(Bytes.EMPTY),
+          bytecode,
           deploymentInfo.deploymentNumber(address),
           deploymentInfo.getDeploymentStatus(address));
     }

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

@@ -72,7 +72,7 @@ public class CommonFragmentValues {
 
   public CommonFragmentValues(Hub hub) {
     final short exceptions = hub.pch().exceptions();
-    final boolean noStackException = !stackException(exceptions);
+    final boolean stackException = stackException(exceptions);
 
     final boolean isExec = hub.state.getProcessingPhase() == TX_EXEC;
 
@@ -86,7 +86,7 @@ public CommonFragmentValues(Hub hub) {
     this.exceptions = exceptions;
     // this.contextNumberNew = hub.contextNumberNew(callFrame);
     this.pc = isExec ? hub.currentFrame().pc() : 0;
-    this.pcNew = computePcNew(hub, pc, noStackException, isExec);
+    this.pcNew = computePcNew(hub, pc, stackException, isExec);
     this.height = callFrame.stack().getHeight();
     this.heightNew = callFrame.stack().getHeightNew();
 
@@ -166,15 +166,16 @@ public void setTracedException(TracedException tracedException) {
     this.tracedException = tracedException;
   }
 
-  static int computePcNew(
-      final Hub hub, final int pc, boolean noStackException, boolean hubInExecPhase) {
+  static int computePcNew(final Hub hub, final int pc, boolean stackException, boolean isExec) {
     final OpCode opCode = hub.opCode();
-    if (!(noStackException && hubInExecPhase)) {
+    if (!isExec || stackException) {
       return 0;
     }
 
+    if (!opCode.isPush() && !opCode.isJump()) return pc + 1;
+
     if (opCode.getData().isPush()) {
-      return pc + opCode.byteValue() - OpCode.PUSH1.byteValue() + 2;
+      return pc + 1 + (opCode.byteValue() - OpCode.PUSH1.byteValue() + 1);
     }
 
     if (opCode.isJump()) {
@@ -193,12 +194,14 @@ static int computePcNew(
         BigInteger condition = hub.currentFrame().frame().getStackItem(1).toUnsignedBigInteger();
         if (!condition.equals(BigInteger.ZERO)) {
           return attemptedPcNew;
+        } else {
+          return pc + 1;
         }
       }
     }
-    ;
 
-    return pc + 1;
+    throw new RuntimeException(
+        "Instruction not covered " + opCode.getData().mnemonic() + " unable to compute pcNew.");
   }
 
   private long computeGasRemaining() {

arithmetization/src/main/java/net/consensys/linea/zktracer/module/hub/fragment/scenario/CreateScenarioFragment.java

@@ -54,7 +54,7 @@ public Trace trace(Trace trace) {
         .pScenarioCreateException(scenario == CreateScenario.CREATE_EXCEPTION)
         .pScenarioCreateAbort(scenario == CreateScenario.CREATE_ABORT)
         .pScenarioCreateFailureConditionWillRevert(
-            scenario == CreateScenario.CREATE_NON_EMPTY_INIT_CODE_FAILURE_WILL_REVERT)
+            scenario == CreateScenario.CREATE_FAILURE_CONDITION_WILL_REVERT)
         .pScenarioCreateFailureConditionWontRevert(
             scenario == CreateScenario.CREATE_FAILURE_CONDITION_WONT_REVERT)
         .pScenarioCreateEmptyInitCodeWillRevert(

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

@@ -26,8 +26,7 @@
 import static net.consensys.linea.zktracer.module.hub.fragment.scenario.CreateScenarioFragment.CreateScenario.CREATE_NON_EMPTY_INIT_CODE_FAILURE_WONT_REVERT;
 import static net.consensys.linea.zktracer.module.hub.fragment.scenario.CreateScenarioFragment.CreateScenario.CREATE_NON_EMPTY_INIT_CODE_SUCCESS_WILL_REVERT;
 import static net.consensys.linea.zktracer.module.hub.fragment.scenario.CreateScenarioFragment.CreateScenario.CREATE_NON_EMPTY_INIT_CODE_SUCCESS_WONT_REVERT;
-import static net.consensys.linea.zktracer.opcode.OpCode.CREATE;
-import static net.consensys.linea.zktracer.opcode.OpCode.CREATE2;
+import static net.consensys.linea.zktracer.opcode.OpCode.*;
 import static net.consensys.linea.zktracer.types.AddressUtils.getDeploymentAddress;
 
 import java.util.Optional;
@@ -53,8 +52,10 @@
 import net.consensys.linea.zktracer.module.shakiradata.ShakiraDataOperation;
 import net.consensys.linea.zktracer.runtime.callstack.CallFrame;
 import org.apache.tuweni.bytes.Bytes;
+import org.apache.tuweni.units.bigints.UInt256;
 import org.hyperledger.besu.datatypes.Address;
 import org.hyperledger.besu.datatypes.Transaction;
+import org.hyperledger.besu.datatypes.Wei;
 import org.hyperledger.besu.evm.account.Account;
 import org.hyperledger.besu.evm.account.AccountState;
 import org.hyperledger.besu.evm.frame.MessageFrame;
@@ -86,6 +87,9 @@ public class CreateSection extends TraceSection
   // row i+?
   private ContextFragment finalContextFragment;
 
+  private boolean requiresRomLex;
+  private Wei value;
+
   // TODO: according to our preliminary conclusion in issue #866
   //  CREATE's that raise a failure condition _do spawn a child context_.
   public CreateSection(Hub hub) {
@@ -177,13 +181,14 @@ 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));
+
     // 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);
       imcFragment.callMmu(mmuCall);
 
-      final long offset = Words.clampedToLong(hub.messageFrame().getStackItem(1));
-      final long size = Words.clampedToLong(hub.messageFrame().getStackItem(2));
       final Bytes create2InitCode = messageFrame.shadowReadMemory(offset, size);
       final ShakiraDataOperation shakiraDataOperation =
           new ShakiraDataOperation(hub.stamp(), create2InitCode);
@@ -192,6 +197,8 @@ public CreateSection(Hub hub) {
       triggerHashInfo(shakiraDataOperation.result());
     }
 
+    value = failedCreate ? Wei.ZERO : Wei.of(UInt256.fromBytes(hub.messageFrame().getStackItem(0)));
+
     if (failedCreate || emptyInitCode) {
       finalContextFragment = ContextFragment.nonExecutionProvidesEmptyReturnData(hub);
 
@@ -215,8 +222,13 @@ public CreateSection(Hub hub) {
         .scheduleForContextReEntry(
             this, hub.currentFrame()); // To get the success bit of the CREATE(2)
 
+    requiresRomLex = true;
     hub.romLex().callRomLex(messageFrame);
-    hub.transients().conflation().deploymentInfo().newDeploymentWithExecutionAt(createeAddress);
+    hub.transients()
+        .conflation()
+        .deploymentInfo()
+        .newDeploymentWithExecutionAt(
+            createeAddress, hub.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) {
@@ -229,8 +241,14 @@ public CreateSection(Hub hub) {
 
   @Override
   public void resolveUponContextEntry(Hub hub) {
-    childEntryCreatorSnapshot = AccountSnapshot.canonical(hub, preOpcodeCreatorSnapshot.address());
-    childEntryCreateeSnapshot = AccountSnapshot.canonical(hub, preOpcodeCreateeSnapshot.address());
+    childEntryCreatorSnapshot =
+        AccountSnapshot.canonical(hub, preOpcodeCreatorSnapshot.address())
+            // .raiseNonceByOne() // the nonce was already raised
+            .decrementBalanceBy(value);
+    childEntryCreateeSnapshot =
+        AccountSnapshot.canonical(hub, preOpcodeCreateeSnapshot.address())
+            .raiseNonceByOne()
+            .incrementBalanceBy(value);
 
     final AccountFragment.AccountFragmentFactory accountFragmentFactory =
         hub.factories().accountFragment();
@@ -243,11 +261,14 @@ public void resolveUponContextEntry(Hub hub) {
     creatorAccountFragment.rlpAddrSubFragment(rlpAddrSubFragment);
 
     final AccountFragment createeAccountFragment =
-        accountFragmentFactory.make(
+        accountFragmentFactory.makeWithTrm(
             preOpcodeCreateeSnapshot,
             childEntryCreateeSnapshot,
+            preOpcodeCreateeSnapshot.address().trimLeadingZeros(),
             DomSubStampsSubFragment.standardDomSubStamps(this.hubStamp(), 1));
 
+    createeAccountFragment.requiresRomlex(requiresRomLex);
+
     this.addFragments(creatorAccountFragment, createeAccountFragment);
   }
 

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

@@ -88,7 +88,8 @@ public TxInitializationSection(Hub hub, WorldView world) {
     }
 
     if (isDeployment) {
-      deploymentInfo.newDeploymentWithExecutionAt(recipientAddress);
+      deploymentInfo.newDeploymentWithExecutionAt(
+          recipientAddress, tx.getBesuTransaction().getInit().orElse(Bytes.EMPTY));
     }
 
     final Bytecode initCode = new Bytecode(tx.getBesuTransaction().getInit().orElse(Bytes.EMPTY));

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

@@ -41,10 +41,12 @@
 import net.consensys.linea.zktracer.module.hub.signals.Exceptions;
 import net.consensys.linea.zktracer.module.hub.signals.TracedException;
 import net.consensys.linea.zktracer.runtime.callstack.CallFrame;
+import net.consensys.linea.zktracer.types.Bytecode;
 import org.apache.tuweni.bytes.Bytes;
 import org.hyperledger.besu.datatypes.Address;
 import org.hyperledger.besu.datatypes.Transaction;
 import org.hyperledger.besu.evm.frame.MessageFrame;
+import org.hyperledger.besu.evm.internal.Words;
 import org.hyperledger.besu.evm.worldstate.WorldView;
 
 @Getter
@@ -209,6 +211,9 @@ public ReturnSection(Hub hub) {
 
       // Empty deployments
       if (!nonemptyByteCode) {
+        if (hub.messageFrame().getDepth() == 0) {
+          this.addDeploymentAccountFragmentIfRoot(hub, mxpCall);
+        }
         return;
       }
 
@@ -232,6 +237,10 @@ public ReturnSection(Hub hub) {
       secondImcFragment.callMmu(nonemptyDeploymentMmuCall);
 
       triggerHashInfo(nonemptyDeploymentMmuCall.hashResult());
+
+      if (hub.messageFrame().getDepth() == 0) {
+        this.addDeploymentAccountFragmentIfRoot(hub, mxpCall);
+      }
     }
   }
 
@@ -302,6 +311,28 @@ public void resolvePostTransaction(
     this.addFragment(squashParentContextReturnData);
   }
 
+  private void addDeploymentAccountFragmentIfRoot(Hub hub, MxpCall mxpCall) {
+    // in case of zero depth we don't have a ContextReEntry step so we have to add the
+    // deployment account fragment manually
+    postDeploymentAccountSnapshot = AccountSnapshot.canonical(hub, deploymentAddress);
+    postDeploymentAccountSnapshot.code(
+        new Bytecode(
+            hub.messageFrame()
+                .shadowReadMemory(
+                    Words.clampedToLong(mxpCall.offset1), Words.clampedToLong(mxpCall.size1))));
+    postDeploymentAccountSnapshot.deploymentStatus(false);
+
+    final AccountFragment deploymentAccountFragment =
+        hub.factories()
+            .accountFragment()
+            .make(
+                preDeploymentAccountSnapshot,
+                postDeploymentAccountSnapshot,
+                DomSubStampsSubFragment.standardDomSubStamps(this.hubStamp(), 0));
+
+    this.addFragment(deploymentAccountFragment);
+  }
+
   private static short maxNumberOfRows(Hub hub) {
     return (short)
         (hub.opCode().numberOfStackRows() + (Exceptions.any(hub.pch().exceptions()) ? 4 : 7));

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

@@ -18,12 +18,14 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import org.apache.tuweni.bytes.Bytes;
 import org.hyperledger.besu.datatypes.Address;
 
 /** Stores information relative to contract deployment. */
 public class DeploymentInfo {
   private final Map<Address, Integer> deploymentNumber = new HashMap<>();
   private final Map<Address, Boolean> deploymentStatus = new HashMap<>();
+  private final Map<Address, Bytes> initializationCodes = new HashMap<>();
 
   /**
    * Returns the deployment number of the given address; sets it to zero if it is the first
@@ -36,19 +38,22 @@ public final int deploymentNumber(Address address) {
     return this.getDeploymentNumber(address);
   }
 
-  public void newDeploymentWithExecutionAt(Address address) {
+  public void newDeploymentWithExecutionAt(Address address, Bytes bytecode) {
     this.incrementDeploymentNumber(address);
     this.markAsUnderDeployment(address);
+    this.setInitializationCode(address, bytecode);
   }
 
   public void newDeploymentSansExecutionAt(Address address) {
     this.incrementDeploymentNumber(address);
     this.markAsNotUnderDeployment(address);
+    this.setInitializationCode(address, Bytes.EMPTY);
   }
 
   public void freshDeploymentNumberFinishingSelfdestruct(Address address) {
     this.incrementDeploymentNumber(address);
     this.markAsNotUnderDeployment(address);
+    this.setInitializationCode(address, Bytes.EMPTY);
   }
 
   private int getDeploymentNumber(Address address) {
@@ -59,6 +64,10 @@ public final boolean getDeploymentStatus(Address address) {
     return this.deploymentStatus.getOrDefault(address, false);
   }
 
+  public final Bytes getInitializationCode(Address address) {
+    return this.initializationCodes.get(address);
+  }
+
   private void incrementDeploymentNumber(Address address) {
     int currentDeploymentNumber = getDeploymentNumber(address);
     deploymentNumber.put(address, currentDeploymentNumber + 1);
@@ -71,4 +80,8 @@ private void markAsUnderDeployment(Address address) {
   public final void markAsNotUnderDeployment(Address address) {
     this.deploymentStatus.put(address, false);
   }
+
+  public void setInitializationCode(Address address, Bytes bytecode) {
+    this.initializationCodes.put(address, bytecode);
+  }
 }

arithmetization/src/main/java/net/consensys/linea/zktracer/opcode/OpCode.java

@@ -213,6 +213,11 @@ public UnsignedByte unsignedByteValue() {
     return UnsignedByte.of(byteValue());
   }
 
+  /** Returns true for PUSH-type instructions */
+  public boolean isPush() {
+    return getData().isPush();
+  }
+
   /** Returns true for JUMP-type instructions */
   public boolean isJump() {
     return getData().isJump();