Skip to content

Commit

Permalink
Improve StackItem memory footprint (#1390)
Browse files Browse the repository at this point in the history 
* Improve StackOperation memory footprint by changing height from int to short
Signed-off-by: Ameziane H <[email protected]>
  • Loading branch information
@ahamlat
ahamlat authored Oct 11, 2024
1 parent 6370425 commit 39077d6
Show file tree
Hide file tree
Showing 4 changed files with 93 additions and 58 deletions.
4 changes: 2 additions & 2 deletions ...in/java/net/consensys/linea/zktracer/module/hub/fragment/common/CommonFragmentValues.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,8 +81,8 @@ public CommonFragmentValues(Hub hub) {
// this.contextNumberNew = hub.contextNumberNew(callFrame);
this.pc = hubProcessingPhase == TX_EXEC ? hub.currentFrame().pc() : 0;
this.pcNew = computePcNew(hub, pc, noStackException, hub.state.getProcessingPhase() == TX_EXEC);
this.height = (short) callFrame.stack().getHeight();
this.heightNew = (short) callFrame.stack().getHeightNew();
this.height = callFrame.stack().getHeight();
this.heightNew = callFrame.stack().getHeightNew();

// TODO: partial solution, will not work in general
this.gasExpected = hub.expectedGas();
Expand Down
6 changes: 4 additions & 2 deletions ...etization/src/main/java/net/consensys/linea/zktracer/module/hub/section/TraceSection.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import net.consensys.linea.zktracer.module.hub.fragment.common.CommonFragment;
import net.consensys.linea.zktracer.module.hub.fragment.common.CommonFragmentValues;
import net.consensys.linea.zktracer.runtime.callstack.CallFrame;
import net.consensys.linea.zktracer.runtime.stack.Stack;
import net.consensys.linea.zktracer.runtime.stack.StackLine;
import org.apache.tuweni.bytes.Bytes;
import org.hyperledger.besu.evm.internal.Words;
Expand Down Expand Up @@ -158,12 +159,13 @@ private int computeContextNumberNew() {

private List<TraceFragment> makeStackFragments(final Hub hub, CallFrame f) {
final List<TraceFragment> r = new ArrayList<>(2);
Stack snapshot = f.stack().snapshot();
if (f.pending().lines().isEmpty()) {
for (int i = 0; i < (f.opCodeData().stackSettings().twoLineInstruction() ? 2 : 1); i++) {
r.add(
StackFragment.prepare(
hub,
f.stack().snapshot(),
snapshot,
new StackLine().asStackItems(),
hub.pch().exceptions(),
hub.pch().abortingConditions().snapshot(),
Expand All @@ -177,7 +179,7 @@ private List<TraceFragment> makeStackFragments(final Hub hub, CallFrame f) {
r.add(
StackFragment.prepare(
hub,
f.stack().snapshot(),
snapshot,
line.asStackItems(),
hub.pch().exceptions(),
hub.pch().abortingConditions().snapshot(),
Expand Down
131 changes: 82 additions & 49 deletions arithmetization/src/main/java/net/consensys/linea/zktracer/runtime/stack/Stack.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,8 @@ public class Stack {
public static final byte PUSH = 1;
public static final byte POP = 2;

@Getter int height;
@Getter int heightNew;
@Getter short height;
@Getter short heightNew;
@Getter OpCodeData currentOpcodeData;
Status status;
int stamp;
Expand Down Expand Up @@ -78,7 +78,8 @@ private void twoZero(MessageFrame frame, StackContext pending) {

private void zeroOne(MessageFrame ignoredFrame, StackContext pending) {
pending.addArmingLine(
new IndexedStackOperation(4, StackItem.push(height + 1, stackStampWithOffset(0))));
new IndexedStackOperation(
4, StackItem.push((short) (height + 1), stackStampWithOffset(0))));
}

private void oneOne(MessageFrame frame, StackContext pending) {
Expand All @@ -95,8 +96,10 @@ private void twoOne(MessageFrame frame, StackContext pending) {

pending.addArmingLine(
new IndexedStackOperation(1, StackItem.pop(height, val1, stackStampWithOffset(0))),
new IndexedStackOperation(2, StackItem.pop(height - 1, val2, stackStampWithOffset(1))),
new IndexedStackOperation(4, StackItem.push(height - 1, stackStampWithOffset(2))));
new IndexedStackOperation(
2, StackItem.pop((short) (height - 1), val2, stackStampWithOffset(1))),
new IndexedStackOperation(
4, StackItem.push((short) (height - 1), stackStampWithOffset(2))));
}

private void threeOne(MessageFrame frame, StackContext pending) {
Expand All @@ -106,9 +109,12 @@ private void threeOne(MessageFrame frame, StackContext pending) {

pending.addArmingLine(
new IndexedStackOperation(1, StackItem.pop(height, val1, stackStampWithOffset(0))),
new IndexedStackOperation(2, StackItem.pop(height - 1, val2, stackStampWithOffset(1))),
new IndexedStackOperation(3, StackItem.pop(height - 2, val3, stackStampWithOffset(2))),
new IndexedStackOperation(4, StackItem.push(height - 2, stackStampWithOffset(3))));
new IndexedStackOperation(
2, StackItem.pop((short) (height - 1), val2, stackStampWithOffset(1))),
new IndexedStackOperation(
3, StackItem.pop((short) (height - 2), val3, stackStampWithOffset(2))),
new IndexedStackOperation(
4, StackItem.push((short) (height - 2), stackStampWithOffset(3))));
}

private void loadStore(MessageFrame frame, StackContext pending) {
Expand All @@ -118,7 +124,8 @@ private void loadStore(MessageFrame frame, StackContext pending) {

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height, val1, stackStampWithOffset(0))),
new IndexedStackOperation(4, StackItem.pop(height - 1, val2, stackStampWithOffset(1))));
new IndexedStackOperation(
4, StackItem.pop((short) (height - 1), val2, stackStampWithOffset(1))));
} else {
Bytes val = getStack(frame, 0);

Expand All @@ -133,11 +140,12 @@ private void dup(MessageFrame frame, StackContext pending) {
Bytes val = getStack(frame, depth);

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height - depth, val, stackStampWithOffset(0))),
new IndexedStackOperation(
2, StackItem.pushImmediate(height - depth, val, stackStampWithOffset(1))),
1, StackItem.pop((short) (height - depth), val, stackStampWithOffset(0))),
new IndexedStackOperation(
4, StackItem.pushImmediate(height + 1, val, stackStampWithOffset(2))));
2, StackItem.pushImmediate((short) (height - depth), val, stackStampWithOffset(1))),
new IndexedStackOperation(
4, StackItem.pushImmediate((short) (height + 1), val, stackStampWithOffset(2))));
}

private void swap(MessageFrame frame, StackContext pending) {
Expand All @@ -146,10 +154,11 @@ private void swap(MessageFrame frame, StackContext pending) {
Bytes val2 = getStack(frame, depth);

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height - depth, val1, stackStampWithOffset(0))),
new IndexedStackOperation(
1, StackItem.pop((short) (height - depth), val1, stackStampWithOffset(0))),
new IndexedStackOperation(2, StackItem.pop(height, val2, stackStampWithOffset(1))),
new IndexedStackOperation(
3, StackItem.pushImmediate(height - depth, val2, stackStampWithOffset(2))),
3, StackItem.pushImmediate((short) (height - depth), val2, stackStampWithOffset(2))),
new IndexedStackOperation(
4, StackItem.pushImmediate(height, val1, stackStampWithOffset(3))));
}
Expand All @@ -161,7 +170,8 @@ private void log(MessageFrame frame, StackContext pending) {
// Stack line 1
pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height, offset, stackStampWithOffset(0))),
new IndexedStackOperation(2, StackItem.pop(height - 1, size, stackStampWithOffset(1))));
new IndexedStackOperation(
2, StackItem.pop((short) (height - 1), size, stackStampWithOffset(1))));

// Stack line 2
IndexedStackOperation[] line2 = new IndexedStackOperation[] {};
Expand All @@ -173,7 +183,7 @@ private void log(MessageFrame frame, StackContext pending) {
line2 =
new IndexedStackOperation[] {
new IndexedStackOperation(
1, StackItem.pop(height - 2, topic1, stackStampWithOffset(0))),
1, StackItem.pop((short) (height - 2), topic1, stackStampWithOffset(0))),
};
}
case LOG2 -> {
Expand All @@ -183,9 +193,9 @@ private void log(MessageFrame frame, StackContext pending) {
line2 =
new IndexedStackOperation[] {
new IndexedStackOperation(
1, StackItem.pop(height - 2, topic1, stackStampWithOffset(2))),
1, StackItem.pop((short) (height - 2), topic1, stackStampWithOffset(2))),
new IndexedStackOperation(
2, StackItem.pop(height - 3, topic2, stackStampWithOffset(3))),
2, StackItem.pop((short) (height - 3), topic2, stackStampWithOffset(3))),
};
}
case LOG3 -> {
Expand All @@ -196,11 +206,11 @@ private void log(MessageFrame frame, StackContext pending) {
line2 =
new IndexedStackOperation[] {
new IndexedStackOperation(
1, StackItem.pop(height - 2, topic1, stackStampWithOffset(2))),
1, StackItem.pop((short) (height - 2), topic1, stackStampWithOffset(2))),
new IndexedStackOperation(
2, StackItem.pop(height - 3, topic2, stackStampWithOffset(3))),
2, StackItem.pop((short) (height - 3), topic2, stackStampWithOffset(3))),
new IndexedStackOperation(
3, StackItem.pop(height - 4, topic3, stackStampWithOffset(4))),
3, StackItem.pop((short) (height - 4), topic3, stackStampWithOffset(4))),
};
}
case LOG4 -> {
Expand All @@ -212,13 +222,13 @@ private void log(MessageFrame frame, StackContext pending) {
line2 =
new IndexedStackOperation[] {
new IndexedStackOperation(
1, StackItem.pop(height - 2, topic1, stackStampWithOffset(2))),
1, StackItem.pop((short) (height - 2), topic1, stackStampWithOffset(2))),
new IndexedStackOperation(
2, StackItem.pop(height - 3, topic2, stackStampWithOffset(3))),
2, StackItem.pop((short) (height - 3), topic2, stackStampWithOffset(3))),
new IndexedStackOperation(
3, StackItem.pop(height - 4, topic3, stackStampWithOffset(4))),
3, StackItem.pop((short) (height - 4), topic3, stackStampWithOffset(4))),
new IndexedStackOperation(
4, StackItem.pop(height - 5, topic4, stackStampWithOffset(5))),
4, StackItem.pop((short) (height - 5), topic4, stackStampWithOffset(5))),
};
}
default -> throw new RuntimeException("not a LOGx");
Expand All @@ -234,9 +244,12 @@ private void copy(MessageFrame frame, StackContext pending) {
Bytes val3 = getStack(frame, 3);

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height - 1, val1, stackStampWithOffset(1))),
new IndexedStackOperation(2, StackItem.pop(height - 3, val3, stackStampWithOffset(2))),
new IndexedStackOperation(3, StackItem.pop(height - 2, val2, stackStampWithOffset(3))),
new IndexedStackOperation(
1, StackItem.pop((short) (height - 1), val1, stackStampWithOffset(1))),
new IndexedStackOperation(
2, StackItem.pop((short) (height - 3), val3, stackStampWithOffset(2))),
new IndexedStackOperation(
3, StackItem.pop((short) (height - 2), val2, stackStampWithOffset(3))),
new IndexedStackOperation(4, StackItem.pop(height, val0, stamp)));
} else {
Bytes val1 = getStack(frame, 0);
Expand All @@ -245,8 +258,10 @@ private void copy(MessageFrame frame, StackContext pending) {

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height, val1, stackStampWithOffset(1))),
new IndexedStackOperation(2, StackItem.pop(height - 2, val2, stackStampWithOffset(2))),
new IndexedStackOperation(3, StackItem.pop(height - 1, val3, stackStampWithOffset(3))));
new IndexedStackOperation(
2, StackItem.pop((short) (height - 2), val2, stackStampWithOffset(2))),
new IndexedStackOperation(
3, StackItem.pop((short) (height - 1), val3, stackStampWithOffset(3))));
}
}

Expand All @@ -264,27 +279,40 @@ private void call(MessageFrame frame, StackContext pending) {
Bytes val7 = getStack(frame, 6);

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height - 3, val4, stackStampWithOffset(3))),
new IndexedStackOperation(2, StackItem.pop(height - 4, val5, stackStampWithOffset(4))),
new IndexedStackOperation(3, StackItem.pop(height - 5, val6, stackStampWithOffset(5))),
new IndexedStackOperation(4, StackItem.pop(height - 6, val7, stackStampWithOffset(6))));
new IndexedStackOperation(
1, StackItem.pop((short) (height - 3), val4, stackStampWithOffset(3))),
new IndexedStackOperation(
2, StackItem.pop((short) (height - 4), val5, stackStampWithOffset(4))),
new IndexedStackOperation(
3, StackItem.pop((short) (height - 5), val6, stackStampWithOffset(5))),
new IndexedStackOperation(
4, StackItem.pop((short) (height - 6), val7, stackStampWithOffset(6))));
pending.addArmingLine(
new IndexedStackOperation(1, StackItem.pop(height, val1, stackStampWithOffset(0))),
new IndexedStackOperation(2, StackItem.pop(height - 1, val2, stackStampWithOffset(1))),
new IndexedStackOperation(3, StackItem.pop(height - 2, val3, stackStampWithOffset(2))),
new IndexedStackOperation(4, StackItem.push(height - 6, stackStampWithOffset(7))));
new IndexedStackOperation(
2, StackItem.pop((short) (height - 1), val2, stackStampWithOffset(1))),
new IndexedStackOperation(
3, StackItem.pop((short) (height - 2), val3, stackStampWithOffset(2))),
new IndexedStackOperation(
4, StackItem.push((short) (height - 6), stackStampWithOffset(7))));
} else {

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height - 2, val3, stackStampWithOffset(3))),
new IndexedStackOperation(2, StackItem.pop(height - 3, val4, stackStampWithOffset(4))),
new IndexedStackOperation(3, StackItem.pop(height - 4, val5, stackStampWithOffset(5))),
new IndexedStackOperation(4, StackItem.pop(height - 5, val6, stackStampWithOffset(6))));
new IndexedStackOperation(
1, StackItem.pop((short) (height - 2), val3, stackStampWithOffset(3))),
new IndexedStackOperation(
2, StackItem.pop((short) (height - 3), val4, stackStampWithOffset(4))),
new IndexedStackOperation(
3, StackItem.pop((short) (height - 4), val5, stackStampWithOffset(5))),
new IndexedStackOperation(
4, StackItem.pop((short) (height - 5), val6, stackStampWithOffset(6))));

pending.addArmingLine(
new IndexedStackOperation(1, StackItem.pop(height, val1, stackStampWithOffset(0))),
new IndexedStackOperation(2, StackItem.pop(height - 1, val2, stackStampWithOffset(1))),
new IndexedStackOperation(4, StackItem.push(height - 5, stackStampWithOffset(7))));
new IndexedStackOperation(
2, StackItem.pop((short) (height - 1), val2, stackStampWithOffset(1))),
new IndexedStackOperation(
4, StackItem.push((short) (height - 5), stackStampWithOffset(7))));
}
}

Expand All @@ -293,25 +321,30 @@ private void create(MessageFrame frame, StackContext pending) {
final Bytes val2 = getStack(frame, 2);

pending.addLine(
new IndexedStackOperation(1, StackItem.pop(height - 1, val1, stackStampWithOffset(1))),
new IndexedStackOperation(2, StackItem.pop(height - 2, val2, stackStampWithOffset(2))));
new IndexedStackOperation(
1, StackItem.pop((short) (height - 1), val1, stackStampWithOffset(1))),
new IndexedStackOperation(
2, StackItem.pop((short) (height - 2), val2, stackStampWithOffset(2))));
// case CREATE2
if (currentOpcodeData.stackSettings().flag2()) {
final Bytes val3 = getStack(frame, 3);
final Bytes val4 = getStack(frame, 0);

pending.addArmingLine(
new IndexedStackOperation(2, StackItem.pop(height - 3, val3, stackStampWithOffset(3))),
new IndexedStackOperation(
2, StackItem.pop((short) (height - 3), val3, stackStampWithOffset(3))),
new IndexedStackOperation(3, StackItem.pop(height, val4, stackStampWithOffset(0))),
new IndexedStackOperation(4, StackItem.push(height - 3, stackStampWithOffset(4))));
new IndexedStackOperation(
4, StackItem.push((short) (height - 3), stackStampWithOffset(4))));
} else
// case CREATE
{
final Bytes val4 = getStack(frame, 0);

pending.addArmingLine(
new IndexedStackOperation(3, StackItem.pop(height, val4, stackStampWithOffset(0))),
new IndexedStackOperation(4, StackItem.push(height - 2, stackStampWithOffset(4))));
new IndexedStackOperation(
4, StackItem.push((short) (height - 2), stackStampWithOffset(4))));
}
}

Expand Down Expand Up @@ -346,7 +379,7 @@ public void processInstruction(final Hub hub, MessageFrame frame, int stackStamp
final int delta = currentOpcodeData.stackSettings().delta();

Preconditions.checkArgument(heightNew == frame.stackSize());
height = frame.stackSize();
height = (short) frame.stackSize();
heightNew += currentOpcodeData.stackSettings().nbAdded();
heightNew -= currentOpcodeData.stackSettings().nbRemoved();

Expand Down
10 changes: 5 additions & 5 deletions arithmetization/src/main/java/net/consensys/linea/zktracer/runtime/stack/StackItem.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ public final class StackItem {
* The relative height of the element with regard to the stack height just before executing the
* linked EVM instruction.
*/
@Getter private final int height;
@Getter private final short height;

/** The value having been popped from/pushed on the stack. */
@Getter @Setter private Bytes value;
Expand Down Expand Up @@ -65,23 +65,23 @@ private StackItem() {
this.stackStamp = 0;
}

StackItem(int height, Bytes value, byte action, int stackStamp) {
StackItem(short height, Bytes value, byte action, int stackStamp) {
this.height = height;
this.value = value;
this.action = action;
this.stackStamp = stackStamp;
}

public static StackItem pop(int height, Bytes value, int stackStamp) {
public static StackItem pop(short height, Bytes value, int stackStamp) {
return new StackItem(height, value, Stack.POP, stackStamp);
}

public static StackItem push(int height, int stackStamp) {
public static StackItem push(short height, int stackStamp) {
return new StackItem(
height, MARKER /* marker value, erased on unlatching */, Stack.PUSH, stackStamp);
}

public static StackItem pushImmediate(int height, Bytes val, int stackStamp) {
public static StackItem pushImmediate(short height, Bytes val, int stackStamp) {
return new StackItem(height, val.copy(), Stack.PUSH, stackStamp);
}
}

0 comments on commit 39077d6

Please sign in to comment.