Skip to content

Commit 53ce3ac

Browse files
anvacaruanvacaru
committed
add rest of opcode hooks
1 parent 345b9b9 commit 53ce3ac

File tree

2 files changed

+205
-25
lines changed

2 files changed

+205
-25
lines changed

‎src/kontrol/foundry.py

Lines changed: 203 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,10 @@ def __init__(
109109
CustomStep(self._forget_branch_pattern, self._exec_forget_custom_step),
110110
CustomStep(self._call_fork_pattern, self._exec_fork_call_custom_step),
111111
CustomStep(self._sload_fork_pattern, self._exec_fork_sload_custom_step),
112+
CustomStep(self._balance_fork_pattern, self._exec_balance_fork_custom_step),
113+
CustomStep(self._extcodesize_fork_pattern, self._exec_extcodesize_fork_custom_step),
114+
CustomStep(self._extcodehash_fork_pattern, self._exec_extcodehash_fork_custom_step),
115+
CustomStep(self._extcodecopy_fork_pattern, self._exec_extcodecopy_fork_custom_step),
112116
)
113117

114118
super().__init__(
@@ -138,11 +142,11 @@ def cut_point_rules(
138142
cut_point_rules.extend(
139143
[
140144
'EVM.call.false',
141-
'FOUNDRY.sload.w3provider',
142-
'FOUNDRY.balance.w3provider',
143-
'FOUNDRY.extcodesize.w3provider',
144-
'FOUNDRY.extcodehash.w3provider',
145-
'FOUNDRY.extcodecopy.w3provider',
145+
# 'FOUNDRY.sload.w3provider',
146+
# 'FOUNDRY.balance.w3provider',
147+
# 'FOUNDRY.extcodesize.w3provider',
148+
# 'FOUNDRY.extcodehash.w3provider',
149+
# 'FOUNDRY.extcodecopy.w3provider',
146150
]
147151
)
148152
return cut_point_rules + KEVMSemantics.cut_point_rules(
@@ -189,9 +193,6 @@ def _call_fork_pattern(self) -> KSequence:
189193
KVariable('###STATIC'),
190194
],
191195
),
192-
KApply('#return___EVM_KItem_Int_Int', [KVariable('###RETURN1'), KVariable('###RETURN2')]),
193-
KApply('pc', KVariable('###PC')),
194-
KApply('execute', []),
195196
KVariable('###CONTINUATION'),
196197
]
197198
)
@@ -201,11 +202,6 @@ def _sload_fork_pattern(self) -> KSequence:
201202
return KSequence(
202203
[
203204
KApply('FETCH_ACCOUNT_STORAGE', [KVariable('###ACCTID'), KVariable('###ACCTSLOT')]),
204-
KApply('#push_EVM_InternalOp'),
205-
KApply('pc', KVariable('###PC1')),
206-
KApply('execute', []),
207-
KApply('#return___EVM_KItem_Int_Int', [KVariable('###RETURN1'), KVariable('###RETURN2')]),
208-
KApply('pc', KVariable('###PC2')),
209205
KVariable('###CONTINUATION'),
210206
]
211207
)
@@ -215,7 +211,36 @@ def _balance_fork_pattern(self) -> KSequence:
215211
return KSequence(
216212
[
217213
KApply('FETCH_ACCOUNT_BALANCE', KVariable('###ACCTID')),
218-
KApply('#push_EVM_InternalOp'),
214+
KVariable('###CONTINUATION'),
215+
]
216+
)
217+
218+
@property
219+
def _extcodesize_fork_pattern(self) -> KSequence:
220+
return KSequence(
221+
[
222+
KApply('FETCH_ACCOUNT_CODE_SIZE', KVariable('###ACCTID')),
223+
KVariable('###CONTINUATION'),
224+
]
225+
)
226+
227+
@property
228+
def _extcodehash_fork_pattern(self) -> KSequence:
229+
return KSequence(
230+
[
231+
KApply('FETCH_ACCOUNT_CODE_HASH', KVariable('###ACCTID')),
232+
KVariable('###CONTINUATION'),
233+
]
234+
)
235+
236+
@property
237+
def _extcodecopy_fork_pattern(self) -> KSequence:
238+
return KSequence(
239+
[
240+
KApply(
241+
'FETCH_ACCOUNT_CODE',
242+
[KVariable('###ACCTID'), KVariable('###MEMSTART'), KVariable('###PGMSTART'), KVariable('###WIDTH')],
243+
),
219244
KVariable('###CONTINUATION'),
220245
]
221246
)
@@ -368,7 +393,7 @@ def _exec_fork_call_custom_step(self, subst: Subst, cterm: CTerm, _c: CTermSymbo
368393
raise ValueError(f'Account {address_value} already processed for fork.')
369394

370395
# Read the account from the web3 provider
371-
account_code, account = fetch_account_from_provider(provider=self._provider, target_address=target_address)
396+
account_code, _, account = fetch_account_from_provider(provider=self._provider, target_address=target_address)
372397

373398
# Update the ACCOUNTS_CELL by appending the new account.
374399
accounts_cell = cterm.cell('ACCOUNTS_CELL')
@@ -396,9 +421,6 @@ def _exec_fork_call_custom_step(self, subst: Subst, cterm: CTerm, _c: CTermSymbo
396421
subst['###STATIC'],
397422
],
398423
),
399-
KApply('#return___EVM_KItem_Int_Int', [subst['###RETURN1'], subst['###RETURN2']]),
400-
KApply('pc', subst['###PC']),
401-
KApply('execute', []),
402424
subst['###CONTINUATION'],
403425
]
404426
)
@@ -432,7 +454,7 @@ def _exec_fork_sload_custom_step(self, subst: Subst, cterm: CTerm, _c: CTermSymb
432454
# Fetch the account from the web3 provider if it hasn't been processed yet.
433455

434456
if address_value not in self._external_accounts:
435-
_, account_to_update = fetch_account_from_provider(self._provider, target_address)
457+
_, _, account_to_update = fetch_account_from_provider(self._provider, target_address)
436458
new_account_list.extend(accounts)
437459
else:
438460
# find target_address in <accounts> cell map
@@ -462,11 +484,6 @@ def _exec_fork_sload_custom_step(self, subst: Subst, cterm: CTerm, _c: CTermSymb
462484
continuation = KSequence(
463485
[
464486
value,
465-
KApply('#push_EVM_InternalOp'),
466-
KApply('pc', subst['###PC1']),
467-
KApply('execute'),
468-
KApply('#return___EVM_KItem_Int_Int', [subst['###RETURN1'], subst['###RETURN2']]),
469-
KApply('pc', subst['###PC2']),
470487
subst['###CONTINUATION'],
471488
]
472489
)
@@ -485,6 +502,168 @@ def _exec_fork_sload_custom_step(self, subst: Subst, cterm: CTerm, _c: CTermSymb
485502
)
486503
return Step(CTerm(new_cterm.config, cterm.constraints), 1, (), ['FETCH_ACCOUNT_STORAGE'], cut=True)
487504

505+
def _exec_balance_fork_custom_step(self, subst: Subst, cterm: CTerm, _c: CTermSymbolic) -> KCFGExtendResult | None:
506+
if self._provider is None:
507+
raise ValueError('No web3 provider configured for fork execution.')
508+
509+
target_address = subst['###ACCTID']
510+
if not type(target_address) is KToken:
511+
raise TypeError('Expected target_address to be KToken instance.')
512+
address_value = int(target_address.token)
513+
514+
_, account_balance, account = fetch_account_from_provider(self._provider, target_address)
515+
516+
# Update the ACCOUNTS_CELL by appending the new account.
517+
accounts_cell = cterm.cell('ACCOUNTS_CELL')
518+
all_accounts = flatten_label('_AccountCellMap_', accounts_cell)
519+
all_accounts.append(account)
520+
new_accounts_cell = KEVM.accounts(all_accounts)
521+
522+
# Build the new continuation.
523+
continuation = KSequence(
524+
[
525+
account_balance,
526+
subst['###CONTINUATION'],
527+
]
528+
)
529+
530+
# Update the FORKEDACCOUNTS_CELL via a helper.
531+
new_forked_accounts_cell = self._update_forked_accounts(cterm, target_address)
532+
self._external_accounts.add(address_value)
533+
534+
# Update the Accounts cell with the continuation
535+
new_cterm = CTerm.from_kast(set_cell(cterm.kast, 'FORKEDACCOUNTS_CELL', new_forked_accounts_cell))
536+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'ACCOUNTS_CELL', new_accounts_cell))
537+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'K_CELL', continuation))
538+
539+
return Step(CTerm(new_cterm.config, cterm.constraints), 1, (), ['FETCH_ACCOUNT_BALANCE'], cut=True)
540+
541+
def _exec_extcodesize_fork_custom_step(
542+
self, subst: Subst, cterm: CTerm, _c: CTermSymbolic
543+
) -> KCFGExtendResult | None:
544+
if self._provider is None:
545+
raise ValueError('No web3 provider configured for fork execution.')
546+
547+
target_address = subst['###ACCTID']
548+
if not type(target_address) is KToken:
549+
raise TypeError('Expected target_address to be KToken instance.')
550+
address_value = int(target_address.token)
551+
552+
account_code, _, account = fetch_account_from_provider(self._provider, target_address)
553+
554+
# Update the ACCOUNTS_CELL by appending the new account.
555+
accounts_cell = cterm.cell('ACCOUNTS_CELL')
556+
all_accounts = flatten_label('_AccountCellMap_', accounts_cell)
557+
all_accounts.append(account)
558+
new_accounts_cell = KEVM.accounts(all_accounts)
559+
560+
# Build the new continuation.
561+
continuation = KSequence(
562+
[
563+
KEVM.size_bytes(account_code),
564+
subst['###CONTINUATION'],
565+
]
566+
)
567+
# Update the FORKEDACCOUNTS_CELL via a helper.
568+
new_forked_accounts_cell = self._update_forked_accounts(cterm, target_address)
569+
self._external_accounts.add(address_value)
570+
571+
# Update the Accounts cell with the continuation
572+
new_cterm = CTerm.from_kast(set_cell(cterm.kast, 'FORKEDACCOUNTS_CELL', new_forked_accounts_cell))
573+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'ACCOUNTS_CELL', new_accounts_cell))
574+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'K_CELL', continuation))
575+
576+
return Step(CTerm(new_cterm.config, cterm.constraints), 1, (), ['FETCH_ACCOUNT_CODE_SIZE'], cut=True)
577+
578+
def _exec_extcodehash_fork_custom_step(
579+
self, subst: Subst, cterm: CTerm, _c: CTermSymbolic
580+
) -> KCFGExtendResult | None:
581+
if self._provider is None:
582+
raise ValueError('No web3 provider configured for fork execution.')
583+
584+
target_address = subst['###ACCTID']
585+
if not type(target_address) is KToken:
586+
raise TypeError('Expected target_address to be KToken instance.')
587+
address_value = int(target_address.token)
588+
589+
account_code, _, account = fetch_account_from_provider(self._provider, target_address)
590+
591+
# Update the ACCOUNTS_CELL by appending the new account.
592+
accounts_cell = cterm.cell('ACCOUNTS_CELL')
593+
all_accounts = flatten_label('_AccountCellMap_', accounts_cell)
594+
all_accounts.append(account)
595+
new_accounts_cell = KEVM.accounts(all_accounts)
596+
597+
# Build the new continuation.
598+
continuation = KSequence(
599+
[
600+
KApply('keccak', account_code),
601+
subst['###CONTINUATION'],
602+
]
603+
)
604+
# Update the FORKEDACCOUNTS_CELL via a helper.
605+
new_forked_accounts_cell = self._update_forked_accounts(cterm, target_address)
606+
self._external_accounts.add(address_value)
607+
608+
# Update the Accounts cell with the continuation
609+
new_cterm = CTerm.from_kast(set_cell(cterm.kast, 'FORKEDACCOUNTS_CELL', new_forked_accounts_cell))
610+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'ACCOUNTS_CELL', new_accounts_cell))
611+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'K_CELL', continuation))
612+
613+
return Step(CTerm(new_cterm.config, cterm.constraints), 1, (), ['FETCH_ACCOUNT_CODE_HASH'], cut=True)
614+
615+
def _exec_extcodecopy_fork_custom_step(
616+
self, subst: Subst, cterm: CTerm, _c: CTermSymbolic
617+
) -> KCFGExtendResult | None:
618+
if self._provider is None:
619+
raise ValueError('No web3 provider configured for fork execution.')
620+
621+
target_address = subst['###ACCTID']
622+
memstart = subst['###MEMSTART']
623+
pgmstart = subst['###PGMSTART']
624+
width = subst['###WIDTH']
625+
if not (
626+
type(target_address) is KToken
627+
and type(memstart) is KToken
628+
and type(pgmstart) is KToken
629+
and type(width) is KToken
630+
):
631+
raise TypeError('Expected target_address to be KToken instance.')
632+
address_value = int(target_address.token)
633+
634+
account_code, _, account = fetch_account_from_provider(self._provider, target_address)
635+
636+
# Update the ACCOUNTS_CELL by appending the new account.
637+
accounts_cell = cterm.cell('ACCOUNTS_CELL')
638+
all_accounts = flatten_label('_AccountCellMap_', accounts_cell)
639+
all_accounts.append(account)
640+
new_accounts_cell = KEVM.accounts(all_accounts)
641+
642+
# <localMem> LM => LM [ MEMSTART := #range(PGM, PGMSTART, WIDTH) ] </localMem>
643+
local_mem = cterm.cell('LOCALMEM_CELL')
644+
new_local_mem = KApply(
645+
'mapWriteRange', [local_mem, memstart, KApply('#range', [account_code, pgmstart, width])]
646+
)
647+
cterm = CTerm.from_kast(set_cell(cterm.kast, 'LOCALMEM_CELL', new_local_mem))
648+
649+
# Build the new continuation.
650+
continuation = KSequence(
651+
[
652+
subst['###CONTINUATION'],
653+
]
654+
)
655+
656+
# Update the FORKEDACCOUNTS_CELL via a helper.
657+
new_forked_accounts_cell = self._update_forked_accounts(cterm, target_address)
658+
self._external_accounts.add(address_value)
659+
660+
# Update the Accounts cell with the continuation
661+
new_cterm = CTerm.from_kast(set_cell(cterm.kast, 'FORKEDACCOUNTS_CELL', new_forked_accounts_cell))
662+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'ACCOUNTS_CELL', new_accounts_cell))
663+
new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'K_CELL', continuation))
664+
665+
return Step(CTerm(new_cterm.config, cterm.constraints), 1, (), ['FETCH_ACCOUNT_CODE'], cut=True)
666+
488667
def _update_forked_accounts(self, cterm: CTerm, target_address: KToken) -> KInner:
489668
"""Update the FORKEDACCOUNTS_CELL by adding target_address if it is not already present."""
490669

‎src/kontrol/web3_rpc.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ def get_block_metadata(provider: Web3) -> dict[str, int]:
5757
}
5858

5959

60-
def fetch_account_from_provider(provider: Web3, target_address: KToken) -> tuple[KToken, KApply]:
60+
def fetch_account_from_provider(provider: Web3, target_address: KToken) -> tuple[KToken, KToken, KApply]:
6161
"""Fetch the account's code and balance from the provider, and return a tuple containing:
6262
- A KToken representing the account code.
6363
- A KApply account cell with the fetched data and empty storage.
@@ -68,6 +68,7 @@ def fetch_account_from_provider(provider: Web3, target_address: KToken) -> tuple
6868

6969
return (
7070
code_token,
71+
balance_token,
7172
KEVM.account_cell(
7273
id=target_address,
7374
balance=balance_token,

0 commit comments

Comments
 (0)