Gas module implementation (#1122)

arithmetization/src/main/java/net/consensys/linea/zktracer/module/gas/Gas.java

@@ -20,48 +20,63 @@
 import java.util.List;
 
 import lombok.Getter;
+import lombok.RequiredArgsConstructor;
 import lombok.experimental.Accessors;
 import net.consensys.linea.zktracer.ColumnHeader;
 import net.consensys.linea.zktracer.container.module.OperationSetModule;
 import net.consensys.linea.zktracer.container.stacked.ModuleOperationStackedSet;
+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.common.CommonFragmentValues;
+import net.consensys.linea.zktracer.module.hub.signals.Exceptions;
+import net.consensys.linea.zktracer.module.hub.signals.TracedException;
+import net.consensys.linea.zktracer.module.wcp.Wcp;
 import org.hyperledger.besu.evm.frame.MessageFrame;
+import org.hyperledger.besu.evm.operation.Operation;
 
+@RequiredArgsConstructor
 @Accessors(fluent = true)
-public class Gas implements OperationSetModule<GasOperation> {
+public class Gas implements OperationSetModule<GasOperation>, PostOpcodeDefer {
   /** A list of the operations to trace */
   @Getter
   private final ModuleOperationStackedSet<GasOperation> operations =
       new ModuleOperationStackedSet<>();
 
+  private CommonFragmentValues commonValues;
+  private GasParameters gasParameters;
+  private final Wcp wcp;
+
   @Override
   public String moduleKey() {
     return "GAS";
   }
 
   @Override
   public List<ColumnHeader> columnsHeaders() {
-    return null;
-  }
-
-  @Override
-  public void tracePreOpcode(MessageFrame frame) {
-    GasParameters gasParameters = extractGasParameters(frame);
-    this.operations.add(new GasOperation(gasParameters));
+    return Trace.headers(this.lineCount());
   }
 
-  private GasParameters extractGasParameters(MessageFrame frame) {
-    // TODO: fill it with the actual values
-    return new GasParameters(0, BigInteger.ZERO, BigInteger.ZERO, false, false);
+  public void call(GasParameters gasParameters, Hub hub, CommonFragmentValues commonValues) {
+    this.commonValues = commonValues;
+    this.gasParameters = gasParameters;
+    hub.defers().scheduleForPostExecution(this);
   }
 
   @Override
   public void commit(List<MappedByteBuffer> buffers) {
     final Trace trace = new Trace(buffers);
-    int stamp = 0;
     for (GasOperation gasOperation : operations.sortOperations(new GasOperationComparator())) {
-      // TODO: I thought we don't have stamp for gas anymore ?
-      stamp++;
-      gasOperation.trace(stamp, trace);
+      gasOperation.trace(trace);
     }
   }
+
+  @Override
+  public void resolvePostExecution(
+      Hub hub, MessageFrame frame, Operation.OperationResult operationResult) {
+    gasParameters.gasActual(BigInteger.valueOf(commonValues.gasActual));
+    gasParameters.gasCost(BigInteger.valueOf(commonValues.gasCost()));
+    gasParameters.xahoy(Exceptions.any(commonValues.exceptions));
+    gasParameters.oogx(commonValues.tracedException() == TracedException.OUT_OF_GAS_EXCEPTION);
+    this.operations.add(new GasOperation(gasParameters, wcp));
+  }
 }

arithmetization/src/main/java/net/consensys/linea/zktracer/module/gas/GasOperation.java

@@ -18,13 +18,15 @@
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.EVM_INST_LT;
 import static net.consensys.linea.zktracer.module.constants.GlobalConstants.WCP_INST_LEQ;
 import static net.consensys.linea.zktracer.types.Conversions.bigIntegerToBytes;
+import static net.consensys.linea.zktracer.types.Utils.initArray;
 
 import java.math.BigInteger;
 
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.experimental.Accessors;
 import net.consensys.linea.zktracer.container.ModuleOperation;
+import net.consensys.linea.zktracer.module.wcp.Wcp;
 import net.consensys.linea.zktracer.types.UnsignedByte;
 
 @Accessors(fluent = true)
@@ -35,58 +37,59 @@ public class GasOperation extends ModuleOperation {
   BigInteger[] wcpArg2Lo;
   UnsignedByte[] wcpInst;
   boolean[] wcpRes;
-  int CT_MAX;
+  int ctMax;
 
-  public GasOperation(GasParameters gasParameters) {
+  public GasOperation(GasParameters gasParameters, Wcp wcp) {
     this.gasParameters = gasParameters;
-    CT_MAX = gasParameters.ctMax();
+    ctMax = compareGasActualAndGasCost() ? 2 : 1;
 
     // init arrays
-    wcpArg1Lo = new BigInteger[CT_MAX + 1];
-    wcpArg2Lo = new BigInteger[CT_MAX + 1];
-    wcpInst = new UnsignedByte[CT_MAX + 1];
-    wcpRes = new boolean[CT_MAX + 1];
+    wcpArg1Lo = initArray(BigInteger.ZERO, ctMax + 1);
+    wcpArg2Lo = initArray(BigInteger.ZERO, ctMax + 1);
+    wcpInst = initArray(UnsignedByte.of(0), ctMax + 1);
+    wcpRes = new boolean[ctMax + 1];
 
     // row 0
     wcpArg1Lo[0] = BigInteger.ZERO;
     wcpArg2Lo[0] = gasParameters.gasActual();
     wcpInst[0] = UnsignedByte.of(WCP_INST_LEQ);
-    wcpRes[0] = true;
+    final boolean gasActualIsNonNegative = wcp.callLEQ(0, gasParameters.gasActual().longValue());
+    wcpRes[0] = gasActualIsNonNegative; // supposed to be true
 
     // row 1
     wcpArg1Lo[1] = BigInteger.ZERO;
     wcpArg2Lo[1] = gasParameters.gasCost();
     wcpInst[1] = UnsignedByte.of(WCP_INST_LEQ);
-    wcpRes[1] = true;
+    final boolean gasCostIsNonNegative = wcp.callLEQ(0, gasParameters.gasCost().longValue());
+    wcpRes[1] = gasCostIsNonNegative; // supposed to be true
 
     // row 2
-    if (gasParameters.oogx()) {
+    if (compareGasActualAndGasCost()) {
       wcpArg1Lo[2] = gasParameters.gasActual();
       wcpArg2Lo[2] = gasParameters.gasCost();
       wcpInst[2] = UnsignedByte.of(EVM_INST_LT);
-      wcpRes[2] = gasParameters.oogx();
-    } else {
-      // TODO: init the lists with zeros (or something equivalent) instead
-      wcpArg1Lo[2] = BigInteger.ZERO;
-      wcpArg2Lo[2] = BigInteger.ZERO;
-      wcpInst[2] = UnsignedByte.of(0);
-      wcpRes[2] = false;
+      final boolean gasActualLTGasCost =
+          wcp.callLT(gasParameters.gasActual().longValue(), gasParameters.gasCost().longValue());
+      wcpRes[2] = gasActualLTGasCost; // supposed to be equal to gasParameters.isOogx()
     }
   }
 
+  private boolean compareGasActualAndGasCost() {
+    return !gasParameters.xahoy() || gasParameters.oogx();
+  }
+
   @Override
   protected int computeLineCount() {
-    return gasParameters.ctMax() + 1;
+    return ctMax + 1;
   }
 
-  public void trace(int stamp, Trace trace) {
-    for (short i = 0; i < CT_MAX + 1; i++) {
-      // TODO: review traced values
+  public void trace(Trace trace) {
+    for (short i = 0; i < ctMax + 1; i++) {
       trace
-          .inputsAndOutputsAreMeaningful(stamp != 0)
+          .inputsAndOutputsAreMeaningful(true)
           .first(i == 0)
           .ct(i)
-          .ctMax(CT_MAX)
+          .ctMax(ctMax)
           .gasActual(bigIntegerToBytes(gasParameters.gasActual()))
           .gasCost(bigIntegerToBytes(gasParameters.gasCost()))
           .exceptionsAhoy(gasParameters.xahoy())

arithmetization/src/main/java/net/consensys/linea/zktracer/module/gas/GasParameters.java

@@ -17,13 +17,18 @@
 
 import java.math.BigInteger;
 
-public record GasParameters(
-    int ctMax, // TODO @Lorenzo this shouldn't be in gasParameters, because it shouldn't be in the
-    // EqualAndHAsh, it's a consequence of oogx and xahoy
-    BigInteger gasActual,
-    BigInteger gasCost,
-    boolean xahoy,
-    boolean oogx) {
+import lombok.Getter;
+import lombok.Setter;
+import lombok.experimental.Accessors;
+
+@Getter
+@Setter
+@Accessors(fluent = true)
+public class GasParameters {
+  BigInteger gasActual;
+  BigInteger gasCost;
+  boolean xahoy;
+  boolean oogx;
 
   public int compareTo(GasParameters other) {
     if (oogx() != other.oogx()) {

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

@@ -211,7 +211,7 @@ public int lineCount() {
   private final Blockhash blockhash = new Blockhash(this, wcp);
   private final Euc euc = new Euc(wcp);
   @Getter private final Ext ext = new Ext(this);
-  private final Gas gas = new Gas();
+  @Getter private final Gas gas = new Gas(wcp);
   private final Mul mul = new Mul(this);
   private final Mod mod = new Mod();
   private final Shf shf = new Shf();
@@ -332,6 +332,7 @@ public List<Module> getModulesToTrace() {
                 exp,
                 ext,
                 euc,
+                gas,
                 logData,
                 logInfo,
                 mmu, // WARN: must be traced before the MMIO

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

@@ -27,6 +27,7 @@
 import lombok.RequiredArgsConstructor;
 import lombok.Setter;
 import lombok.experimental.Accessors;
+import net.consensys.linea.zktracer.module.gas.GasParameters;
 import net.consensys.linea.zktracer.module.hub.Hub;
 import net.consensys.linea.zktracer.module.hub.HubProcessingPhase;
 import net.consensys.linea.zktracer.module.hub.State;
@@ -105,6 +106,11 @@ public CommonFragmentValues(Hub hub) {
                     || instructionFamily == INVALID)
                 || any(this.exceptions));
 
+    if (contextMayChange) {
+      // Trigger the gas module in case contextMayChange is true
+      hub.gas().call(new GasParameters(), hub, this);
+    }
+
     if (none(exceptions)) {
       tracedException = TracedException.NONE;
       return;

arithmetization/src/main/java/net/consensys/linea/zktracer/types/Utils.java

@@ -21,7 +21,10 @@
 import java.math.BigInteger;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.toml.Toml;
@@ -112,4 +115,33 @@ public static Map<String, Integer> computeSpillings() throws IOException {
 
     return spillings;
   }
+
+  /**
+   * Initializes an array with a specified value and size.
+   *
+   * @param <T> The type of the elements in the array.
+   * @param initValue The value to initialize each element of the array with.
+   * @param size The size of the array to be created.
+   * @return An array of the specified size, with each element initialized to the specified value.
+   */
+  @SuppressWarnings("unchecked")
+  public static <T> T[] initArray(T initValue, int size) {
+    return Stream.generate(() -> initValue)
+        .limit(size)
+        .toArray(i -> (T[]) java.lang.reflect.Array.newInstance(initValue.getClass(), i));
+  }
+
+  /**
+   * Initializes a list with a specified value and size.
+   *
+   * @param <T> The type of the elements in the list.
+   * @param initValue The value to initialize each element of the list with.
+   * @param size The size of the list to be created.
+   * @return A list of the specified size, with each element initialized to the specified value.
+   */
+  public static <T> List<T> initList(T initValue, int size) {
+    return Stream.generate(() -> initValue)
+        .limit(size)
+        .collect(Collectors.toCollection(ArrayList::new));
+  }
 }