Optimize stack related WCP operations

arithmetization/src/main/java/net/consensys/linea/zktracer/container/stacked/NoLineCountStackedSet.java

@@ -0,0 +1,64 @@
+/*
+ * 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.container.stacked;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import lombok.Getter;
+import lombok.experimental.Accessors;
+
+@Accessors(fluent = true)
+@Getter
+public class NoLineCountStackedSet<E> {
+  private final Set<E> operationsCommitedToTheConflation;
+  private final Set<E> operationsInTransaction;
+
+  public NoLineCountStackedSet() {
+    operationsCommitedToTheConflation = new HashSet<>();
+    operationsInTransaction = new HashSet<>();
+  }
+
+  /** Prefer this constructor as we preallocate more needed memory */
+  public NoLineCountStackedSet(
+      final int expectedConflationNumberOperations, final int expectedTransactionNumberOperations) {
+    operationsCommitedToTheConflation = new HashSet<>(expectedConflationNumberOperations);
+    operationsInTransaction = new HashSet<>(expectedTransactionNumberOperations);
+  }
+
+  /**
+   * Upon entering a new transaction, the set of operations generated by the previous transaction
+   * {@link StackedSet#operationsInTransaction()} (which is empty if no such transaction exists) is
+   * added to the set of committed operations {@link
+   * StackedSet#operationsCommitedToTheConflation()}. {@link StackedSet#operationsInTransaction()}
+   * is further reset to be empty.
+   */
+  public void enter() {
+    operationsCommitedToTheConflation().addAll(operationsInTransaction());
+    operationsInTransaction().clear();
+  }
+
+  public void pop() {
+    operationsInTransaction().clear();
+  }
+
+  public boolean add(E e) {
+    if (!operationsCommitedToTheConflation().contains(e)) {
+      return operationsInTransaction().add(e);
+    }
+    return false;
+  }
+}

arithmetization/src/main/java/net/consensys/linea/zktracer/container/stacked/StackedSet.java

@@ -16,7 +16,6 @@
 package net.consensys.linea.zktracer.container.stacked;
 
 import java.util.Collection;
-import java.util.HashSet;
 import java.util.Set;
 
 import com.google.common.base.Preconditions;
@@ -29,52 +28,48 @@
 
 /**
  * Implements a system of pseudo-stacked squashed sets where {@link
- * operationsCommitedToTheConflation} represents the set of all operations since the beginning of
- * the conflation and {@link operationsInTransaction} represents the operations added by the last
- * transaction. We can pop only the operations added by last transaction. The line counting is done
- * by a separate {@link CountOnlyOperation}.
+ * StackedSet#operationsCommitedToTheConflation()} represents the set of all operations since the
+ * beginning of the conflation and {@link StackedSet#operationsInTransaction()} represents the
+ * operations added by the last transaction. We can pop only the operations added by last
+ * transaction. The line counting is done by a separate {@link CountOnlyOperation}.
  *
  * @param <E> the type of elements stored in the set
  */
 @Accessors(fluent = true)
-public class StackedSet<E extends ModuleOperation> {
+public class StackedSet<E extends ModuleOperation> extends NoLineCountStackedSet<E> {
   private static final Logger log = LoggerFactory.getLogger(StackedSet.class);
-  @Getter private final Set<E> operationsCommitedToTheConflation;
-  @Getter private final Set<E> operationsInTransaction;
   private final CountOnlyOperation lineCounter = new CountOnlyOperation();
   @Getter private boolean conflationFinished = false;
 
   public StackedSet() {
-    operationsCommitedToTheConflation = new HashSet<>();
-    operationsInTransaction = new HashSet<>();
+    super();
   }
 
   /** Prefer this constructor as we preallocate more needed memory */
   public StackedSet(
       final int expectedConflationNumberOperations, final int expectedTransactionNumberOperations) {
-    operationsCommitedToTheConflation = new HashSet<>(expectedConflationNumberOperations);
-    operationsInTransaction = new HashSet<>(expectedTransactionNumberOperations);
+    super(expectedConflationNumberOperations, expectedTransactionNumberOperations);
   }
 
   /**
-   * When we enter transaction n, the previous transaction n-1 {@link operationsInTransaction} (or
-   * an empty set if n == 0) is definitely added to the set of all transactions {@link
-   * operationsCommitedToTheConflation}. We reset {@link operationsInTransaction} to an empty set as
-   * it now represents the operations added at transaction n.
+   * Upon entering a new transaction, the set of operations generated by the previous transaction
+   * {@link StackedSet#operationsInTransaction()} (which is empty if no such transaction exists) is
+   * added to the set of committed operations {@link
+   * StackedSet#operationsCommitedToTheConflation()}. {@link StackedSet#operationsInTransaction()}
+   * is further reset to be empty.
    */
   public void enter() {
-    operationsCommitedToTheConflation.addAll(operationsInTransaction);
-    operationsInTransaction.clear();
+    super.enter();
     lineCounter.enter();
   }
 
   public void pop() {
-    operationsInTransaction.clear();
+    super.pop();
     lineCounter.pop();
   }
 
   public int size() {
-    return operationsInTransaction.size() + operationsCommitedToTheConflation.size();
+    return operationsInTransaction().size() + operationsCommitedToTheConflation().size();
   }
 
   public int lineCount() {
@@ -83,20 +78,20 @@ public int lineCount() {
 
   public Set<E> getAll() {
     Preconditions.checkState(conflationFinished, "Conflation not finished");
-    return operationsCommitedToTheConflation;
+    return operationsCommitedToTheConflation();
   }
 
   public boolean isEmpty() {
     return size() == 0;
   }
 
   public boolean contains(Object o) {
-    return operationsInTransaction.contains(o) || operationsCommitedToTheConflation.contains(o);
+    return operationsInTransaction().contains(o) || operationsCommitedToTheConflation().contains(o);
   }
 
   public boolean add(E e) {
-    if (!operationsCommitedToTheConflation.contains(e)) {
-      final boolean isNew = operationsInTransaction.add(e);
+    if (!operationsCommitedToTheConflation().contains(e)) {
+      final boolean isNew = operationsInTransaction().add(e);
       if (isNew) {
         lineCounter.add(e.lineCount());
       }
@@ -128,15 +123,15 @@ public boolean addAll(@NotNull Collection<? extends E> c) {
   }
 
   public void clear() {
-    operationsCommitedToTheConflation.clear();
-    operationsInTransaction.clear();
+    operationsCommitedToTheConflation().clear();
+    operationsInTransaction().clear();
     lineCounter.clear();
   }
 
   public void finishConflation() {
     conflationFinished = true;
-    operationsCommitedToTheConflation.addAll(operationsInTransaction);
-    operationsInTransaction.clear();
+    operationsCommitedToTheConflation().addAll(operationsInTransaction());
+    operationsInTransaction().clear();
     lineCounter.enter(); // this is not mandatory but it is more consistent
   }
 }

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

@@ -444,6 +444,8 @@ public Hub(final Address l2l1ContractAddress, final Bytes l2l1Topic) {
 
   @Override
   public void enterTransaction() {
+    transients.conflation().stackHeightChecksForStackUnderflows().enter();
+    transients.conflation().stackHeightChecksForStackOverflows().enter();
     for (Module m : modules) {
       m.enterTransaction();
     }
@@ -453,6 +455,8 @@ public void enterTransaction() {
   public void popTransaction() {
     txStack.pop();
     state.pop();
+    transients.conflation().stackHeightChecksForStackUnderflows().pop();
+    transients.conflation().stackHeightChecksForStackOverflows().pop();
     for (Module m : modules) {
       m.popTransaction();
     }

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

@@ -15,11 +15,11 @@
 
 package net.consensys.linea.zktracer.module.hub.transients;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
 
 import lombok.Getter;
 import lombok.experimental.Accessors;
+import net.consensys.linea.zktracer.container.stacked.NoLineCountStackedSet;
 import net.consensys.linea.zktracer.runtime.LogData;
 
 /** Stores data relative to the conflation. */
@@ -28,6 +28,10 @@
 public class Conflation {
   private final DeploymentInfo deploymentInfo = new DeploymentInfo();
   private final List<LogData> logs = new ArrayList<>(100);
+  private final NoLineCountStackedSet<StackHeightCheck> stackHeightChecksForStackUnderflows =
+      new NoLineCountStackedSet<>(256, 32);
+  private final NoLineCountStackedSet<StackHeightCheck> stackHeightChecksForStackOverflows =
+      new NoLineCountStackedSet<>(256, 32);
 
   public int log(LogData logData) {
     this.logs.add(logData);

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright Consensys Software 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.transients;
+
+import static com.google.common.base.Preconditions.*;
+import static net.consensys.linea.zktracer.runtime.stack.Stack.MAX_STACK_SIZE;
+
+import lombok.EqualsAndHashCode;
+
+@EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false)
+public class StackHeightCheck {
+  private static final Integer SHIFT_FACTOR =
+      8; // 5 would suffice but 8 makes it a byte shift, not sure if it matters
+
+  @EqualsAndHashCode.Include final int comparison;
+
+  /**
+   * This constructor creates a {@link StackHeightCheck} for stack underflow detection.
+   *
+   * @param height stack height pre opcode execution
+   * @param delta greatest depth at which touched stack items
+   */
+  public StackHeightCheck(int height, int delta) {
+    checkArgument(0 <= height && height <= MAX_STACK_SIZE && 0 <= delta && delta <= 17);
+    comparison = height << SHIFT_FACTOR | delta;
+  }
+
+  /**
+   * This constructor creates a {@link StackHeightCheck} for stack overflow detection.
+   *
+   * @param heightNew stack height post opcode execution
+   */
+  public StackHeightCheck(int heightNew) {
+    checkArgument(0 <= heightNew && heightNew <= MAX_STACK_SIZE + 1);
+    comparison = heightNew;
+  }
+}