Skip to content

Commit 1ae9bb3

Browse files
Merge pull request #6613 from yanghang8612/optimize/energy-calculation
refactor(vm): optimize opcode energy cost calculation with `BigInteger`
2 parents 3f7ff93 + 4971754 commit 1ae9bb3

3 files changed

Lines changed: 281 additions & 0 deletions

File tree

actuator/src/main/java/org/tron/core/vm/EnergyCost.java

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -387,6 +387,27 @@ public static long getVoteWitnessCost2(Program program) {
387387
? amountArrayMemoryNeeded : witnessArrayMemoryNeeded), 0, Op.VOTEWITNESS);
388388
}
389389

390+
public static long getVoteWitnessCost3(Program program) {
391+
Stack stack = program.getStack();
392+
long oldMemSize = program.getMemSize();
393+
BigInteger amountArrayLength = stack.get(stack.size() - 1).value();
394+
BigInteger amountArrayOffset = stack.get(stack.size() - 2).value();
395+
BigInteger witnessArrayLength = stack.get(stack.size() - 3).value();
396+
BigInteger witnessArrayOffset = stack.get(stack.size() - 4).value();
397+
398+
BigInteger wordSize = BigInteger.valueOf(DataWord.WORD_SIZE);
399+
400+
BigInteger amountArraySize = amountArrayLength.multiply(wordSize).add(wordSize);
401+
BigInteger amountArrayMemoryNeeded = memNeeded(amountArrayOffset, amountArraySize);
402+
403+
BigInteger witnessArraySize = witnessArrayLength.multiply(wordSize).add(wordSize);
404+
BigInteger witnessArrayMemoryNeeded = memNeeded(witnessArrayOffset, witnessArraySize);
405+
406+
return VOTE_WITNESS + calcMemEnergy(oldMemSize,
407+
(amountArrayMemoryNeeded.compareTo(witnessArrayMemoryNeeded) > 0
408+
? amountArrayMemoryNeeded : witnessArrayMemoryNeeded), 0, Op.VOTEWITNESS);
409+
}
410+
390411
public static long getWithdrawRewardCost(Program ignored) {
391412
return WITHDRAW_REWARD;
392413
}
@@ -550,6 +571,10 @@ private static BigInteger memNeeded(DataWord offset, DataWord size) {
550571
return size.isZero() ? BigInteger.ZERO : offset.value().add(size.value());
551572
}
552573

574+
private static BigInteger memNeeded(BigInteger offset, BigInteger size) {
575+
return size.equals(BigInteger.ZERO) ? BigInteger.ZERO : offset.add(size);
576+
}
577+
553578
private static boolean isDeadAccount(Program program, DataWord address) {
554579
return program.getContractState().getAccount(address.toTronAddress()) == null;
555580
}

actuator/src/main/java/org/tron/core/vm/OperationRegistry.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,10 @@ public static JumpTable getTable() {
8383
adjustSelfdestruct(table);
8484
}
8585

86+
if (VMConfig.allowTvmOsaka()) {
87+
adjustVoteWitnessCost(table);
88+
}
89+
8690
return table;
8791
}
8892

@@ -706,4 +710,12 @@ public static void adjustSelfdestruct(JumpTable table) {
706710
EnergyCost::getSuicideCost3,
707711
OperationActions::suicideAction2));
708712
}
713+
714+
public static void adjustVoteWitnessCost(JumpTable table) {
715+
table.set(new Operation(
716+
Op.VOTEWITNESS, 4, 1,
717+
EnergyCost::getVoteWitnessCost3,
718+
OperationActions::voteWitnessAction,
719+
VMConfig::allowTvmVote));
720+
}
709721
}
Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
package org.tron.common.runtime.vm;
2+
3+
import static org.junit.Assert.assertEquals;
4+
import static org.junit.Assert.assertTrue;
5+
import static org.mockito.Mockito.mock;
6+
import static org.mockito.Mockito.when;
7+
8+
import java.math.BigInteger;
9+
import lombok.extern.slf4j.Slf4j;
10+
import org.junit.AfterClass;
11+
import org.junit.BeforeClass;
12+
import org.junit.Test;
13+
import org.tron.common.BaseTest;
14+
import org.tron.common.TestConstants;
15+
import org.tron.common.parameter.CommonParameter;
16+
import org.tron.core.config.args.Args;
17+
import org.tron.core.vm.EnergyCost;
18+
import org.tron.core.vm.JumpTable;
19+
import org.tron.core.vm.Op;
20+
import org.tron.core.vm.Operation;
21+
import org.tron.core.vm.OperationRegistry;
22+
import org.tron.core.vm.config.ConfigLoader;
23+
import org.tron.core.vm.config.VMConfig;
24+
import org.tron.core.vm.program.Program;
25+
import org.tron.core.vm.program.Stack;
26+
27+
@Slf4j
28+
public class VoteWitnessCost3Test extends BaseTest {
29+
30+
static {
31+
Args.setParam(new String[]{"--output-directory", dbPath()}, TestConstants.TEST_CONF);
32+
}
33+
34+
@BeforeClass
35+
public static void init() {
36+
CommonParameter.getInstance().setDebug(true);
37+
VMConfig.initAllowTvmVote(1);
38+
VMConfig.initAllowEnergyAdjustment(1);
39+
}
40+
41+
@AfterClass
42+
public static void destroy() {
43+
ConfigLoader.disable = false;
44+
VMConfig.initAllowTvmVote(0);
45+
VMConfig.initAllowEnergyAdjustment(0);
46+
VMConfig.initAllowTvmOsaka(0);
47+
Args.clearParam();
48+
}
49+
50+
private Program mockProgram(long witnessOffset, long witnessLength,
51+
long amountOffset, long amountLength, int memSize) {
52+
Program program = mock(Program.class);
53+
Stack stack = new Stack();
54+
// Stack order: bottom -> top: witnessOffset, witnessLength, amountOffset, amountLength
55+
stack.push(new DataWord(witnessOffset));
56+
stack.push(new DataWord(witnessLength));
57+
stack.push(new DataWord(amountOffset));
58+
stack.push(new DataWord(amountLength));
59+
when(program.getStack()).thenReturn(stack);
60+
when(program.getMemSize()).thenReturn(memSize);
61+
return program;
62+
}
63+
64+
private Program mockProgram(DataWord witnessOffset, DataWord witnessLength,
65+
DataWord amountOffset, DataWord amountLength, int memSize) {
66+
Program program = mock(Program.class);
67+
Stack stack = new Stack();
68+
stack.push(witnessOffset);
69+
stack.push(witnessLength);
70+
stack.push(amountOffset);
71+
stack.push(amountLength);
72+
when(program.getStack()).thenReturn(stack);
73+
when(program.getMemSize()).thenReturn(memSize);
74+
return program;
75+
}
76+
77+
@Test
78+
public void testNormalCase() {
79+
// 2 witnesses at offset 0, 2 amounts at offset 128
80+
Program program = mockProgram(0, 2, 128, 2, 0);
81+
long cost = EnergyCost.getVoteWitnessCost3(program);
82+
// amountArraySize = 2 * 32 + 32 = 96, memNeeded = 128 + 96 = 224
83+
// witnessArraySize = 2 * 32 + 32 = 96, memNeeded = 0 + 96 = 96
84+
// max = 224, memWords = (224 + 31) / 32 * 32 / 32 = 7
85+
// memEnergy = 3 * 7 + 7 * 7 / 512 = 21
86+
// total = 30000 + 21 = 30021
87+
assertEquals(30021, cost);
88+
}
89+
90+
@Test
91+
public void testConsistentWithCost2ForSmallValues() {
92+
// For small values, cost3 should produce the same result as cost2
93+
long[][] testCases = {
94+
{0, 1, 64, 1, 0}, // 1 witness, 1 amount
95+
{0, 3, 128, 3, 0}, // 3 witnesses, 3 amounts
96+
{0, 5, 256, 5, 0}, // 5 witnesses, 5 amounts
97+
{64, 2, 192, 2, 0}, // non-zero offsets
98+
{0, 10, 512, 10, 0}, // 10 witnesses
99+
};
100+
101+
for (long[] tc : testCases) {
102+
Program p2 = mockProgram(tc[0], tc[1], tc[2], tc[3], (int) tc[4]);
103+
Program p3 = mockProgram(tc[0], tc[1], tc[2], tc[3], (int) tc[4]);
104+
long cost2 = EnergyCost.getVoteWitnessCost2(p2);
105+
long cost3 = EnergyCost.getVoteWitnessCost3(p3);
106+
assertEquals("Mismatch for case: witnessOff=" + tc[0] + " witnessLen=" + tc[1]
107+
+ " amountOff=" + tc[2] + " amountLen=" + tc[3], cost2, cost3);
108+
}
109+
}
110+
111+
@Test
112+
public void testZeroLengthArrays() {
113+
// Both arrays have zero length, but cost3 always adds wordSize for dynamic array prefix
114+
Program program = mockProgram(0, 0, 0, 0, 0);
115+
long cost = EnergyCost.getVoteWitnessCost3(program);
116+
// arraySize = 0 * 32 + 32 = 32, memNeeded = 0 + 32 = 32
117+
// memWords = (32 + 31) / 32 * 32 / 32 = 1
118+
// memEnergy = 3 * 1 + 1 * 1 / 512 = 3
119+
assertEquals(30003, cost);
120+
}
121+
122+
@Test
123+
public void testZeroLengthOneArray() {
124+
// witness array zero, amount array non-zero
125+
Program program = mockProgram(0, 0, 64, 1, 0);
126+
long cost = EnergyCost.getVoteWitnessCost3(program);
127+
// witnessArraySize = 0 * 32 + 32 = 32, witnessMemNeeded = 0 + 32 = 32
128+
// amountArraySize = 1 * 32 + 32 = 64, amountMemNeeded = 64 + 64 = 128
129+
// memWords = 128 / 32 = 4
130+
// memEnergy = 3 * 4 + 4 * 4 / 512 = 12
131+
assertEquals(30012, cost);
132+
}
133+
134+
@Test
135+
public void testLargeArrayLengthOverflow() {
136+
// Use a very large value that would overflow in DataWord.mul() in cost2
137+
// DataWord max is 2^256-1, multiplying by 32 would overflow
138+
// In cost3, BigInteger handles this correctly and should trigger memoryOverflow
139+
String maxHex = "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff";
140+
DataWord largeLength = new DataWord(maxHex);
141+
DataWord zeroOffset = new DataWord(0);
142+
143+
Program program = mockProgram(zeroOffset, new DataWord(1),
144+
zeroOffset, largeLength, 0);
145+
146+
boolean overflowCaught = false;
147+
try {
148+
EnergyCost.getVoteWitnessCost3(program);
149+
} catch (Program.OutOfMemoryException e) {
150+
// cost3 should detect memory overflow via checkMemorySize
151+
overflowCaught = true;
152+
}
153+
assertTrue("cost3 should throw memoryOverflow for huge array length", overflowCaught);
154+
}
155+
156+
@Test
157+
public void testLargeOffsetOverflow() {
158+
// Large offset + normal size should trigger memoryOverflow in cost3
159+
String largeHex = "00ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff";
160+
DataWord largeOffset = new DataWord(largeHex);
161+
162+
Program program = mockProgram(largeOffset, new DataWord(1),
163+
new DataWord(0), new DataWord(1), 0);
164+
165+
boolean overflowCaught = false;
166+
try {
167+
EnergyCost.getVoteWitnessCost3(program);
168+
} catch (Program.OutOfMemoryException e) {
169+
overflowCaught = true;
170+
}
171+
assertTrue("cost3 should throw memoryOverflow for huge offset", overflowCaught);
172+
}
173+
174+
@Test
175+
public void testExistingMemorySize() {
176+
// When program already has memory allocated, additional cost is incremental
177+
Program p1 = mockProgram(0, 2, 128, 2, 0);
178+
long costFromZero = EnergyCost.getVoteWitnessCost3(p1);
179+
180+
Program p2 = mockProgram(0, 2, 128, 2, 224);
181+
long costWithExistingMem = EnergyCost.getVoteWitnessCost3(p2);
182+
183+
// With existing memory >= needed, no additional mem cost
184+
assertEquals(30000, costWithExistingMem);
185+
assertTrue(costFromZero > costWithExistingMem);
186+
}
187+
188+
@Test
189+
public void testAmountArrayLargerThanWitnessArray() {
190+
// amount array needs more memory => amount determines cost
191+
Program program = mockProgram(0, 1, 0, 5, 0);
192+
long cost = EnergyCost.getVoteWitnessCost3(program);
193+
// witnessArraySize = 1 * 32 + 32 = 64, memNeeded = 0 + 64 = 64
194+
// amountArraySize = 5 * 32 + 32 = 192, memNeeded = 0 + 192 = 192
195+
// max = 192, memWords = (192 + 31) / 32 * 32 / 32 = 6
196+
// memEnergy = 3 * 6 + 6 * 6 / 512 = 18
197+
assertEquals(30018, cost);
198+
}
199+
200+
@Test
201+
public void testWitnessArrayLargerThanAmountArray() {
202+
// witness array needs more memory => witness determines cost
203+
Program program = mockProgram(0, 5, 0, 1, 0);
204+
long cost = EnergyCost.getVoteWitnessCost3(program);
205+
// witnessArraySize = 5 * 32 + 32 = 192, memNeeded = 0 + 192 = 192
206+
// amountArraySize = 1 * 32 + 32 = 64, memNeeded = 0 + 64 = 64
207+
// max = 192
208+
assertEquals(30018, cost);
209+
}
210+
211+
@Test
212+
public void testOperationRegistryWithoutOsaka() {
213+
VMConfig.initAllowTvmOsaka(0);
214+
JumpTable table = OperationRegistry.getTable();
215+
Operation voteOp = table.get(Op.VOTEWITNESS);
216+
assertTrue(voteOp.isEnabled());
217+
218+
// Without osaka, should use cost2 (from adjustForFairEnergy since allowEnergyAdjustment=1)
219+
Program program = mockProgram(0, 2, 128, 2, 0);
220+
long cost = voteOp.getEnergyCost(program);
221+
long expectedCost2 = EnergyCost.getVoteWitnessCost2(
222+
mockProgram(0, 2, 128, 2, 0));
223+
assertEquals(expectedCost2, cost);
224+
}
225+
226+
@Test
227+
public void testOperationRegistryWithOsaka() {
228+
VMConfig.initAllowTvmOsaka(1);
229+
try {
230+
JumpTable table = OperationRegistry.getTable();
231+
Operation voteOp = table.get(Op.VOTEWITNESS);
232+
assertTrue(voteOp.isEnabled());
233+
234+
// With osaka, should use cost3
235+
Program program = mockProgram(0, 2, 128, 2, 0);
236+
long cost = voteOp.getEnergyCost(program);
237+
long expectedCost3 = EnergyCost.getVoteWitnessCost3(
238+
mockProgram(0, 2, 128, 2, 0));
239+
assertEquals(expectedCost3, cost);
240+
} finally {
241+
VMConfig.initAllowTvmOsaka(0);
242+
}
243+
}
244+
}

0 commit comments

Comments
 (0)