feat: TRM redesign

.github/actions/setup-go-corset/action.yml

@@ -9,4 +9,4 @@ runs:
 
     - name: Install Go Corset
       shell: bash
-      run: go install github.com/consensys/go-corset/cmd/[email protected].2
+      run: go install github.com/consensys/go-corset/cmd/[email protected].6

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

@@ -200,7 +200,7 @@ public int spillage() {
   private final Mul mul = new Mul(this);
   private final Mod mod = new Mod();
   private final Shf shf = new Shf();
-  @Getter private final Trm trm = new Trm();
+  @Getter private final Trm trm = new Trm(wcp);
 
   // other
   private final Blockdata blockdata;

arithmetization/src/main/java/net/consensys/linea/zktracer/module/trm/Trm.java

@@ -15,40 +15,39 @@
 
 package net.consensys.linea.zktracer.module.trm;
 
-import static net.consensys.linea.zktracer.Trace.LLARGE;
-
 import java.util.List;
 
 import lombok.Getter;
+import lombok.RequiredArgsConstructor;
 import lombok.experimental.Accessors;
 import net.consensys.linea.zktracer.Trace;
 import net.consensys.linea.zktracer.container.module.OperationSetModule;
 import net.consensys.linea.zktracer.container.stacked.ModuleOperationStackedSet;
+import net.consensys.linea.zktracer.module.wcp.Wcp;
 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;
 
 @Getter
 @Accessors(fluent = true)
+@RequiredArgsConstructor
 public class Trm implements OperationSetModule<TrmOperation> {
+  private final Wcp wcp;
   private final ModuleOperationStackedSet<TrmOperation> operations =
       new ModuleOperationStackedSet<>();
 
-  static final int MAX_CT = LLARGE;
-  static final int PIVOT_BIT_FLIPS_TO_TRUE = 12;
-
   @Override
   public String moduleKey() {
     return "TRM";
   }
 
-  public Address callTrimming(Bytes32 rawAddress) {
-    operations.add(new TrmOperation(EWord.of(rawAddress)));
+  public Address callTrimming(final Bytes32 rawAddress) {
+    operations.add(new TrmOperation(EWord.of(rawAddress), wcp));
     return Address.extract(rawAddress);
   }
 
-  public Address callTrimming(Bytes rawAddress) {
+  public Address callTrimming(final Bytes rawAddress) {
     return callTrimming(Bytes32.leftPad(rawAddress));
   }
 
@@ -64,9 +63,8 @@ public int spillage() {
 
   @Override
   public void commit(Trace trace) {
-    int stamp = 0;
     for (TrmOperation operation : operations.sortOperations(new TrmOperationComparator())) {
-      operation.trace(trace.trm, ++stamp);
+      operation.trace(trace.trm);
     }
   }
 }

arithmetization/src/main/java/net/consensys/linea/zktracer/module/trm/TrmOperation.java

@@ -15,62 +15,73 @@
 
 package net.consensys.linea.zktracer.module.trm;
 
-import static net.consensys.linea.zktracer.Trace.LLARGE;
-import static net.consensys.linea.zktracer.module.trm.Trm.MAX_CT;
-import static net.consensys.linea.zktracer.module.trm.Trm.PIVOT_BIT_FLIPS_TO_TRUE;
+import static net.consensys.linea.zktracer.Trace.*;
+import static net.consensys.linea.zktracer.Trace.Trm.TRM_CT_MAX;
+import static net.consensys.linea.zktracer.Trace.Trm.TRM_NB_ROWS;
+import static net.consensys.linea.zktracer.module.wcp.WcpCall.*;
 import static net.consensys.linea.zktracer.types.AddressUtils.isPrecompile;
-import static net.consensys.linea.zktracer.types.Utils.bitDecomposition;
-import static net.consensys.linea.zktracer.types.Utils.leftPadTo;
+import static net.consensys.linea.zktracer.types.Conversions.bigIntegerToBytes;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import lombok.RequiredArgsConstructor;
 import lombok.experimental.Accessors;
 import net.consensys.linea.zktracer.Trace;
 import net.consensys.linea.zktracer.container.ModuleOperation;
+import net.consensys.linea.zktracer.module.wcp.Wcp;
+import net.consensys.linea.zktracer.module.wcp.WcpCall;
 import net.consensys.linea.zktracer.types.EWord;
-import net.consensys.linea.zktracer.types.UnsignedByte;
 import org.apache.tuweni.bytes.Bytes;
 import org.hyperledger.besu.datatypes.Address;
 
 @Accessors(fluent = true)
-@RequiredArgsConstructor
 @EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false)
 public class TrmOperation extends ModuleOperation {
   @EqualsAndHashCode.Include @Getter private final EWord rawAddress;
+  private final List<WcpCall> wcpCalls = new ArrayList<>(TRM_NB_ROWS);
+  private static final Bytes TWOFIFTYSIX_TO_THE_TWENTY_BYTES =
+      bigIntegerToBytes(TWOFIFTYSIX_TO_THE_TWENTY);
+  private static final Bytes TWOFIFTYSIX_TO_THE_TWELVE_MO_BYTES =
+      bigIntegerToBytes(TWOFIFTYSIX_TO_THE_TWELVE_MO);
 
-  void trace(Trace.Trm trace, final int stamp) {
-    final Bytes trmHiBytes =
-        leftPadTo(this.rawAddress.hi().slice(PIVOT_BIT_FLIPS_TO_TRUE, 4), LLARGE);
-    final long trmHi = trmHiBytes.slice(PIVOT_BIT_FLIPS_TO_TRUE, 4).toLong();
-    final Boolean isPrec = isPrecompile(Address.extract(this.rawAddress));
-    final int accLastByte =
-        isPrec ? 9 - (0xff & this.rawAddress.get(31)) : (0xff & this.rawAddress.get(31)) - 10;
-    List<Boolean> ones = bitDecomposition(accLastByte, MAX_CT).bitDecList();
+  public TrmOperation(EWord rawAddress, Wcp wcp) {
+    this.rawAddress = rawAddress;
+    final Bytes trmAddress = rawAddress.toAddress();
 
-    for (int ct = 0; ct < MAX_CT; ct++) {
+    wcpCalls.add(0, ltCall(wcp, trmAddress, TWOFIFTYSIX_TO_THE_TWENTY_BYTES));
+    wcpCalls.add(1, leqCall(wcp, rawAddress.slice(0, 12), TWOFIFTYSIX_TO_THE_TWELVE_MO_BYTES));
+    wcpCalls.add(2, isZeroCall(wcp, trmAddress));
+    wcpCalls.add(3, leqCall(wcp, trmAddress, Bytes.ofUnsignedShort(MAX_PRC_ADDRESS)));
+  }
+
+  void trace(Trace.Trm trace) {
+    final Address trmAddress = rawAddress.toAddress();
+    final boolean isPrec = isPrecompile(trmAddress);
+    final long trmAddrHi = trmAddress.slice(0, 4).toLong();
+
+    for (int ct = 0; ct <= TRM_CT_MAX; ct++) {
       trace
+          .iomf(true)
+          .first(ct == 0)
           .ct(ct)
-          .stamp(stamp)
           .isPrecompile(isPrec)
-          .plateauBit(ct >= PIVOT_BIT_FLIPS_TO_TRUE)
-          .rawAddressHi(this.rawAddress.hi())
-          .rawAddressLo(this.rawAddress.lo())
-          .trmAddressHi(trmHi)
-          .accHi(this.rawAddress.hi().slice(0, ct + 1))
-          .accLo(this.rawAddress.lo().slice(0, ct + 1))
-          .accT(trmHiBytes.slice(0, ct + 1).trimLeadingZeros().toLong())
-          .byteHi(UnsignedByte.of(this.rawAddress.hi().get(ct)))
-          .byteLo(UnsignedByte.of(this.rawAddress.lo().get(ct)))
-          .one(ones.get(ct))
+          .rawAddressHi(rawAddress.hi())
+          .rawAddressLo(rawAddress.lo())
+          .trmAddressHi(trmAddrHi)
+          .inst(wcpCalls.get(ct).instruction())
+          .arg1Hi(wcpCalls.get(ct).arg1Hi())
+          .arg1Lo(wcpCalls.get(ct).arg1Lo())
+          .arg2Hi(wcpCalls.get(ct).arg2Hi())
+          .arg2Lo(wcpCalls.get(ct).arg2Lo())
+          .res(wcpCalls.get(ct).result())
           .validateRow();
     }
   }
 
   @Override
   protected int computeLineCount() {
-    return MAX_CT;
+    return TRM_NB_ROWS;
   }
 }

arithmetization/src/main/java/net/consensys/linea/zktracer/module/wcp/WcpCall.java

@@ -0,0 +1,68 @@
+/*
+ * 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.wcp;
+
+import static net.consensys.linea.zktracer.Trace.*;
+
+import lombok.Getter;
+import lombok.experimental.Accessors;
+import net.consensys.linea.zktracer.types.UnsignedByte;
+import org.apache.tuweni.bytes.Bytes;
+import org.apache.tuweni.bytes.Bytes32;
+
+@Getter()
+@Accessors(fluent = true)
+public class WcpCall {
+  private final UnsignedByte instruction;
+  private final Bytes arg1Hi;
+  private final Bytes arg1Lo;
+  private final Bytes arg2Hi;
+  private final Bytes arg2Lo;
+  private final boolean result;
+
+  public WcpCall(Wcp wcp, byte instruction, Bytes arg1, Bytes arg2) {
+    final Bytes32 arg1Bytes32 = Bytes32.leftPad(arg1);
+    final Bytes32 arg2Bytes32 = Bytes32.leftPad(arg2);
+    this.instruction = UnsignedByte.of(instruction);
+    this.arg1Hi = arg1Bytes32.slice(0, LLARGE);
+    this.arg1Lo = arg1Bytes32.slice(LLARGE, LLARGE);
+    this.arg2Hi = arg2Bytes32.slice(0, LLARGE);
+    this.arg2Lo = arg2Bytes32.slice(LLARGE, LLARGE);
+    this.result =
+        switch (instruction) {
+          case EVM_INST_LT -> wcp.callLT(arg1Bytes32, arg2Bytes32);
+          case EVM_INST_EQ -> wcp.callEQ(arg1Bytes32, arg2Bytes32);
+          case EVM_INST_ISZERO -> wcp.callISZERO(arg1Bytes32);
+          case EVM_INST_GT -> wcp.callGT(arg1Bytes32, arg2Bytes32);
+          case WCP_INST_LEQ -> wcp.callLEQ(arg1Bytes32, arg2Bytes32);
+          case WCP_INST_GEQ -> wcp.callGEQ(arg1Bytes32, arg2Bytes32);
+          default -> throw new IllegalArgumentException(
+              "Unexpected wcp instruction: " + instruction);
+        };
+  }
+
+  public static WcpCall ltCall(Wcp wcp, Bytes arg1, Bytes arg2) {
+    return new WcpCall(wcp, (byte) EVM_INST_LT, arg1, arg2);
+  }
+
+  public static WcpCall leqCall(Wcp wcp, Bytes arg1, Bytes arg2) {
+    return new WcpCall(wcp, (byte) WCP_INST_LEQ, arg1, arg2);
+  }
+
+  public static WcpCall isZeroCall(Wcp wcp, Bytes arg1) {
+    return new WcpCall(wcp, (byte) EVM_INST_ISZERO, arg1, Bytes.EMPTY);
+  }
+}