1010
1111import ast
1212from collections .abc import Iterable
13+ from dataclasses import dataclass
1314from typing import Protocol , cast
1415
1516from .cfg_model import CFG , Block
17+ from .meta_markers import CFG_META_PREFIX
1618
1719__all__ = ["CFG" , "CFGBuilder" ]
1820
@@ -26,17 +28,28 @@ class _TryLike(Protocol):
2628 finalbody : list [ast .stmt ]
2729
2830
31+ @dataclass (slots = True )
32+ class _LoopContext :
33+ continue_target : Block
34+ break_target : Block
35+
36+
37+ def _meta_expr (value : str ) -> ast .Expr :
38+ return ast .Expr (value = ast .Name (id = f"{ CFG_META_PREFIX } { value } " , ctx = ast .Load ()))
39+
40+
2941# =========================
3042# CFG Builder
3143# =========================
3244
3345
3446class CFGBuilder :
35- __slots__ = ("cfg" , "current" )
47+ __slots__ = ("_loop_stack" , " cfg" , "current" )
3648
3749 def __init__ (self ) -> None :
3850 self .cfg : CFG
3951 self .current : Block
52+ self ._loop_stack : list [_LoopContext ] = []
4053
4154 def build (
4255 self ,
@@ -73,6 +86,12 @@ def _visit(self, stmt: ast.stmt) -> None:
7386 self .current .is_terminated = True
7487 self .current .add_successor (self .cfg .exit )
7588
89+ case ast .Break ():
90+ self ._visit_break (stmt )
91+
92+ case ast .Continue ():
93+ self ._visit_continue (stmt )
94+
7695 case ast .If ():
7796 self ._visit_if (stmt )
7897
@@ -123,36 +142,61 @@ def _visit_if(self, stmt: ast.If) -> None:
123142 def _visit_while (self , stmt : ast .While ) -> None :
124143 cond_block = self .cfg .create_block ()
125144 body_block = self .cfg .create_block ()
145+ else_block = self .cfg .create_block () if stmt .orelse else None
126146 after_block = self .cfg .create_block ()
127147
128148 self .current .add_successor (cond_block )
129149
130150 self .current = cond_block
131- self ._emit_condition (stmt .test , body_block , after_block )
151+ false_target = else_block if else_block is not None else after_block
152+ self ._emit_condition (stmt .test , body_block , false_target )
132153
154+ self ._loop_stack .append (
155+ _LoopContext (continue_target = cond_block , break_target = after_block )
156+ )
133157 self .current = body_block
134158 self ._visit_statements (stmt .body )
135159 if not self .current .is_terminated :
136160 self .current .add_successor (cond_block )
161+ self ._loop_stack .pop ()
162+
163+ if else_block is not None :
164+ self .current = else_block
165+ self ._visit_statements (stmt .orelse )
166+ if not self .current .is_terminated :
167+ self .current .add_successor (after_block )
137168
138169 self .current = after_block
139170
140171 def _visit_for (self , stmt : ast .For | ast .AsyncFor ) -> None :
141172 iter_block = self .cfg .create_block ()
142173 body_block = self .cfg .create_block ()
174+ else_block = self .cfg .create_block () if stmt .orelse else None
143175 after_block = self .cfg .create_block ()
144176
145177 self .current .add_successor (iter_block )
146178
147179 self .current = iter_block
148180 self .current .statements .append (ast .Expr (value = stmt .iter ))
149181 self .current .add_successor (body_block )
150- self .current .add_successor (after_block )
182+ self .current .add_successor (
183+ else_block if else_block is not None else after_block
184+ )
151185
186+ self ._loop_stack .append (
187+ _LoopContext (continue_target = iter_block , break_target = after_block )
188+ )
152189 self .current = body_block
153190 self ._visit_statements (stmt .body )
154191 if not self .current .is_terminated :
155192 self .current .add_successor (iter_block )
193+ self ._loop_stack .pop ()
194+
195+ if else_block is not None :
196+ self .current = else_block
197+ self ._visit_statements (stmt .orelse )
198+ if not self .current .is_terminated :
199+ self .current .add_successor (after_block )
156200
157201 self .current = after_block
158202
@@ -188,19 +232,36 @@ def _visit_try(self, stmt: _TryLike) -> None:
188232 self .current .add_successor (try_entry )
189233 self .current = try_entry
190234
191- handlers_blocks = [self .cfg .create_block () for _ in stmt .handlers ]
235+ handler_test_blocks = [self .cfg .create_block () for _ in stmt .handlers ]
236+ handler_body_blocks = [self .cfg .create_block () for _ in stmt .handlers ]
192237 else_block = self .cfg .create_block () if stmt .orelse else None
193238 final_block = self .cfg .create_block ()
194239
240+ for idx , (handler , test_block , body_block ) in enumerate (
241+ zip (stmt .handlers , handler_test_blocks , handler_body_blocks , strict = True )
242+ ):
243+ test_block .statements .append (_meta_expr (f"TRY_HANDLER_INDEX:{ idx } " ))
244+ if handler .type is not None :
245+ type_repr = ast .dump (handler .type , annotate_fields = False )
246+ test_block .statements .append (
247+ _meta_expr (f"TRY_HANDLER_TYPE:{ type_repr } " )
248+ )
249+ else :
250+ test_block .statements .append (_meta_expr ("TRY_HANDLER_TYPE:BARE" ))
251+ test_block .add_successor (body_block )
252+ if idx + 1 < len (handler_test_blocks ):
253+ test_block .add_successor (handler_test_blocks [idx + 1 ])
254+ else :
255+ test_block .add_successor (final_block )
256+
195257 # Process each statement in try body
196258 # Link only statements that can raise to exception handlers
197259 for stmt_node in stmt .body :
198260 if self .current .is_terminated :
199261 break
200262
201- if _stmt_can_raise (stmt_node ):
202- for h_block in handlers_blocks :
203- self .current .add_successor (h_block )
263+ if _stmt_can_raise (stmt_node ) and handler_test_blocks :
264+ self .current .add_successor (handler_test_blocks [0 ])
204265
205266 self ._visit (stmt_node )
206267
@@ -212,11 +273,8 @@ def _visit_try(self, stmt: _TryLike) -> None:
212273 self .current .add_successor (final_block )
213274
214275 # Process handlers
215- for handler , h_block in zip (stmt .handlers , handlers_blocks , strict = True ):
216- self .current = h_block
217- if handler .type :
218- self .current .statements .append (ast .Expr (value = handler .type ))
219-
276+ for handler , body_block in zip (stmt .handlers , handler_body_blocks , strict = True ):
277+ self .current = body_block
220278 self ._visit_statements (handler .body )
221279 if not self .current .is_terminated :
222280 self .current .add_successor (final_block )
@@ -236,25 +294,40 @@ def _visit_try(self, stmt: _TryLike) -> None:
236294 def _visit_match (self , stmt : ast .Match ) -> None :
237295 self .current .statements .append (ast .Expr (value = stmt .subject ))
238296
239- subject_block = self . current
297+ previous_test_block : Block | None = None
240298 after_block = self .cfg .create_block ()
241299
242- for case_ in stmt .cases :
243- case_block = self .cfg .create_block ()
244- subject_block .add_successor (case_block )
300+ for idx , case_ in enumerate (stmt .cases ):
301+ case_test_block = self .cfg .create_block ()
302+ case_body_block = self .cfg .create_block ()
303+
304+ if previous_test_block is None :
305+ self .current .add_successor (case_test_block )
306+ else :
307+ previous_test_block .add_successor (case_test_block )
245308
246- self . current = case_block
309+ case_test_block . statements . append ( _meta_expr ( f"MATCH_CASE_INDEX: { idx } " ))
247310
248311 # Record pattern structure
249312 pattern_repr = ast .dump (case_ .pattern , annotate_fields = False )
250- self . current .statements .append (
251- ast . Expr ( value = ast . Constant ( value = f"PATTERN :{ pattern_repr } ") )
313+ case_test_block .statements .append (
314+ _meta_expr ( f"MATCH_PATTERN :{ pattern_repr } " )
252315 )
316+ if case_ .guard is not None :
317+ case_test_block .statements .append (ast .Expr (value = case_ .guard ))
318+
319+ case_test_block .add_successor (case_body_block )
253320
321+ self .current = case_body_block
254322 self ._visit_statements (case_ .body )
255323 if not self .current .is_terminated :
256324 self .current .add_successor (after_block )
257325
326+ previous_test_block = case_test_block
327+
328+ if previous_test_block is not None :
329+ previous_test_block .add_successor (after_block )
330+
258331 self .current = after_block
259332
260333 def _emit_condition (
@@ -300,6 +373,22 @@ def _emit_boolop(
300373
301374 self .current = current
302375
376+ def _visit_break (self , stmt : ast .Break ) -> None :
377+ self .current .statements .append (stmt )
378+ self .current .is_terminated = True
379+ if self ._loop_stack :
380+ self .current .add_successor (self ._loop_stack [- 1 ].break_target )
381+ return
382+ self .current .add_successor (self .cfg .exit )
383+
384+ def _visit_continue (self , stmt : ast .Continue ) -> None :
385+ self .current .statements .append (stmt )
386+ self .current .is_terminated = True
387+ if self ._loop_stack :
388+ self .current .add_successor (self ._loop_stack [- 1 ].continue_target )
389+ return
390+ self .current .add_successor (self .cfg .exit )
391+
303392
304393def _stmt_can_raise (stmt : ast .stmt ) -> bool :
305394 if isinstance (stmt , ast .Raise ):
0 commit comments