Skip to content

Commit 5cf9be9

Browse files
committed
feat(tools): Implement While code opcode generator
1 parent 05cf0d8 commit 5cf9be9

File tree

3 files changed

+32
-2
lines changed

3 files changed

+32
-2
lines changed

src/ethereum_test_tools/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@
7979
Conditional,
8080
Initcode,
8181
Switch,
82+
While,
8283
Yul,
8384
YulCompiler,
8485
)
@@ -149,6 +150,7 @@
149150
"TransactionTest",
150151
"TransactionTestFiller",
151152
"UndefinedOpcodes",
153+
"While",
152154
"Withdrawal",
153155
"WithdrawalRequest",
154156
"Yul",
Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
"""Code related utilities and classes."""
22

3-
from .generators import CalldataCase, Case, CodeGasMeasure, Conditional, Initcode, Switch
3+
from .generators import CalldataCase, Case, CodeGasMeasure, Conditional, Initcode, Switch, While
44
from .yul import Solc, Yul, YulCompiler
55

66
__all__ = (
7-
"Case",
87
"CalldataCase",
8+
"Case",
99
"CodeGasMeasure",
1010
"Conditional",
1111
"Initcode",
1212
"Solc",
1313
"Switch",
14+
"While",
1415
"Yul",
1516
"YulCompiler",
1617
)

src/ethereum_test_tools/code/generators.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,33 @@ def __new__(
235235
return super().__new__(cls, bytecode)
236236

237237

238+
class While(Bytecode):
239+
"""Helper class used to generate while-loop bytecode."""
240+
241+
def __new__(
242+
cls,
243+
*,
244+
body: Bytecode | Op,
245+
condition: Bytecode | Op,
246+
evm_code_type: EVMCodeType = EVMCodeType.LEGACY,
247+
):
248+
"""
249+
Assemble the loop bytecode.
250+
251+
The condition nor the body can leave a stack item on the stack.
252+
"""
253+
bytecode = Bytecode()
254+
if evm_code_type == EVMCodeType.LEGACY:
255+
bytecode += Op.JUMPDEST
256+
bytecode += body
257+
bytecode += Op.JUMPI(
258+
Op.SUB(Op.PC, Op.PUSH4[len(body) + len(condition) + 6]), condition
259+
)
260+
elif evm_code_type == EVMCodeType.EOF_V1:
261+
raise NotImplementedError("EOF while loops are not implemented")
262+
return super().__new__(cls, bytecode)
263+
264+
238265
@dataclass(kw_only=True)
239266
class Case:
240267
"""

0 commit comments

Comments
 (0)