11"""Account structure of ethereum/tests fillers."""
22
3- from typing import Any , Dict , Mapping
3+ from typing import Any , Dict , List , Mapping , Set , Tuple
44
55from pydantic import BaseModel
66
7- from ethereum_test_base_types import EthereumTestRootModel , HexNumber
7+ from ethereum_test_base_types import Bytes , EthereumTestRootModel , HexNumber , Storage
88from ethereum_test_types import Alloc
99
1010from .common import (
@@ -78,7 +78,7 @@ def resolve(self, tags: TagDict) -> Dict[str, Any]:
7878 if compiled_code := self .code .compiled (tags ):
7979 account_properties ["code" ] = compiled_code
8080 if self .nonce is not None :
81- account_properties ["nonce" ] = self .nonce
81+ account_properties ["nonce" ] = self .nonce
8282 if self .storage is not None :
8383 if resolved_storage := self .storage .resolve (tags ):
8484 account_properties ["storage" ] = resolved_storage
@@ -90,51 +90,165 @@ class PreInFiller(EthereumTestRootModel):
9090
9191 root : Dict [AddressOrTagInFiller , AccountInFiller ]
9292
93- def setup (self , pre : Alloc , all_dependencies : Dict [str , Tag ]) -> TagDict :
94- """Resolve the pre-state."""
95- max_tries = len (self .root )
93+ def _build_dependency_graph (
94+ self ,
95+ ) -> Tuple [Dict [str , Set [str ]], Dict [str , AddressOrTagInFiller ]]:
96+ """Build a dependency graph for all tags."""
97+ dep_graph : Dict [str , Set [str ]] = {}
98+ tag_to_address : Dict [str , AddressOrTagInFiller ] = {}
99+
100+ # First pass: identify all tags and their dependencies
101+ for address_or_tag , account in self .root .items ():
102+ if isinstance (address_or_tag , Tag ):
103+ tag_name = address_or_tag .name
104+ tag_to_address [tag_name ] = address_or_tag
105+ dep_graph [tag_name ] = set ()
106+
107+ # Get dependencies from account properties
108+ dependencies = account .tag_dependencies ()
109+ for dep_name in dependencies :
110+ if dep_name != tag_name : # Ignore self-references
111+ dep_graph [tag_name ].add (dep_name )
112+
113+ return dep_graph , tag_to_address
114+
115+ def _topological_sort (self , dep_graph : Dict [str , Set [str ]]) -> List [str ]:
116+ """Perform topological sort on dependency graph."""
117+ # Create a copy to modify
118+ graph = {node : deps .copy () for node , deps in dep_graph .items ()}
119+
120+ # Find nodes with no dependencies
121+ no_deps = [node for node , deps in graph .items () if not deps ]
122+ sorted_nodes = []
123+
124+ while no_deps :
125+ # Process a node with no dependencies
126+ node = no_deps .pop ()
127+ sorted_nodes .append (node )
128+
129+ # Remove this node from other nodes' dependencies
130+ for other_node , deps in graph .items ():
131+ if node in deps :
132+ deps .remove (node )
133+ if not deps and other_node not in sorted_nodes :
134+ no_deps .append (other_node )
135+
136+ # Check for cycles
137+ remaining = [node for node in graph if node not in sorted_nodes ]
138+ if remaining :
139+ # Handle cycles by processing remaining nodes in any order
140+ # This works because self-references are allowed
141+ sorted_nodes .extend (remaining )
142+
143+ return sorted_nodes
96144
97- unresolved_accounts_in_pre = dict (self .root )
145+ def setup (self , pre : Alloc , all_dependencies : Dict [str , Tag ]) -> TagDict :
146+ """Resolve the pre-state with improved tag resolution."""
98147 resolved_accounts : TagDict = {}
99148
100- while max_tries > 0 :
101- # Naive approach to resolve accounts
102- unresolved_accounts_keys = list (unresolved_accounts_in_pre .keys ())
103- for address_or_tag in unresolved_accounts_keys :
104- account = unresolved_accounts_in_pre [address_or_tag ]
105- tag_dependencies = account .tag_dependencies ()
106- # check if all tag dependencies are resolved
107- if all (dependency in resolved_accounts for dependency in tag_dependencies ):
108- account_properties = account .resolve (resolved_accounts )
109- if isinstance (address_or_tag , Tag ):
110- if isinstance (address_or_tag , ContractTag ):
111- resolved_accounts [address_or_tag .name ] = pre .deploy_contract (
112- ** account_properties ,
113- label = address_or_tag .name ,
114- )
115- elif isinstance (address_or_tag , SenderTag ):
116- if "balance" in account_properties :
117- account_properties ["amount" ] = account_properties .pop ("balance" )
118- resolved_accounts [address_or_tag .name ] = pre .fund_eoa (
119- ** account_properties ,
120- label = address_or_tag .name ,
121- )
122- else :
123- raise ValueError (f"Unknown tag type: { address_or_tag } )
124- else :
125- raise ValueError (f"Unknown tag type: { address_or_tag } )
126- del unresolved_accounts_in_pre [address_or_tag ]
127-
128- max_tries -= 1
129- assert len (unresolved_accounts_in_pre ) == 0 , (
130- f"Unresolved accounts: { unresolved_accounts_in_pre }
131- )
149+ # Separate tagged and non-tagged accounts
150+ tagged_accounts = {}
151+ non_tagged_accounts = {}
152+
153+ for address_or_tag , account in self .root .items ():
154+ if isinstance (address_or_tag , Tag ):
155+ tagged_accounts [address_or_tag ] = account
156+ else :
157+ non_tagged_accounts [address_or_tag ] = account
158+
159+ # Step 1: Process non-tagged accounts but don't compile code yet
160+ # We'll compile code later after all tags are resolved
161+ non_tagged_to_process = []
162+ for address , account in non_tagged_accounts .items ():
163+ non_tagged_to_process .append ((address , account ))
164+ resolved_accounts [address .hex ()] = address
165+
166+ # Step 2: Build dependency graph for tagged accounts
167+ dep_graph , tag_to_address = self ._build_dependency_graph ()
168+
169+ # Step 3: Get topological order
170+ resolution_order = self ._topological_sort (dep_graph )
171+
172+ # Step 4: Pre-deploy all contract tags and pre-fund EOAs to get addresses
173+ for tag_name in resolution_order :
174+ if tag_name in tag_to_address :
175+ tag = tag_to_address [tag_name ]
176+ if isinstance (tag , ContractTag ):
177+ # Deploy with placeholder to get address
178+ deployed_address = pre .deploy_contract (
179+ code = b"" , # Temporary placeholder
180+ label = tag_name ,
181+ )
182+ resolved_accounts [tag_name ] = deployed_address
183+ elif isinstance (tag , SenderTag ):
184+ # Create EOA to get address - use amount=1 to ensure account is created
185+ eoa = pre .fund_eoa (amount = 1 , label = tag_name )
186+ # Store the EOA object for SenderKeyTag resolution
187+ resolved_accounts [tag_name ] = eoa
188+
189+ # Step 5: Now resolve all properties with all addresses available
190+ for tag_name in resolution_order :
191+ if tag_name in tag_to_address :
192+ tag = tag_to_address [tag_name ]
193+ assert isinstance (tag , (ContractTag , SenderTag )), (
194+ f"Tag { tag_name }
195+ )
196+ account = tagged_accounts [tag ]
197+
198+ # All addresses are now available, so resolve properties
199+ account_properties = account .resolve (resolved_accounts )
200+
201+ if isinstance (tag , ContractTag ):
202+ # Update the already-deployed contract
203+ deployed_address = resolved_accounts [tag_name ]
204+ deployed_account = pre [deployed_address ]
205+
206+ if deployed_account is not None :
207+ if "code" in account_properties :
208+ deployed_account .code = Bytes (account_properties ["code" ])
209+ if "balance" in account_properties :
210+ deployed_account .balance = account_properties ["balance" ]
211+ if "nonce" in account_properties :
212+ deployed_account .nonce = account_properties ["nonce" ]
213+ if "storage" in account_properties :
214+ deployed_account .storage = Storage (root = account_properties ["storage" ])
215+
216+ elif isinstance (tag , SenderTag ):
217+ eoa_account = pre [resolved_accounts [tag_name ]]
218+
219+ if eoa_account is not None :
220+ if "balance" in account_properties :
221+ eoa_account .balance = account_properties ["balance" ]
222+ if "nonce" in account_properties :
223+ eoa_account .nonce = account_properties ["nonce" ]
224+ if "code" in account_properties :
225+ eoa_account .code = Bytes (account_properties ["code" ])
226+ if "storage" in account_properties :
227+ eoa_account .storage = Storage (root = account_properties ["storage" ])
228+
229+ # Step 6: Now process non-tagged accounts (including code compilation)
230+ for address , account in non_tagged_to_process :
231+ account_properties = account .resolve (resolved_accounts )
232+ if "balance" in account_properties :
233+ pre .fund_address (address , account_properties ["balance" ])
234+
235+ existing_account = pre [address ]
236+ if existing_account is not None :
237+ if "code" in account_properties :
238+ existing_account .code = Bytes (account_properties ["code" ])
239+ if "nonce" in account_properties :
240+ existing_account .nonce = account_properties ["nonce" ]
241+ if "storage" in account_properties :
242+ existing_account .storage = Storage (root = account_properties ["storage" ])
243+
244+ # Step 7: Handle any extra dependencies not in pre
132245 for extra_dependency in all_dependencies :
133246 if extra_dependency not in resolved_accounts :
134247 if all_dependencies [extra_dependency ].type != "eoa" :
135248 raise ValueError (f"Contract dependency { extra_dependency } )
136- # Assume the extra EOA is an empty account
137- resolved_accounts [extra_dependency ] = pre .fund_eoa (
138- amount = 0 , label = extra_dependency
139- )
249+
250+ # Create new EOA - this will have a dynamically generated key and address
251+ eoa = pre .fund_eoa (amount = 0 , label = extra_dependency )
252+ resolved_accounts [extra_dependency ] = eoa
253+
140254 return resolved_accounts
0 commit comments