@@ -236,6 +236,7 @@ pub struct SymbolTable {
236
236
scopes_by_node : FxHashMap < NodeKey , ScopeId > ,
237
237
/// dependencies of this module
238
238
dependencies : Vec < Dependency > ,
239
+ flow_graph : FlowGraph ,
239
240
}
240
241
241
242
impl SymbolTable {
@@ -245,6 +246,7 @@ impl SymbolTable {
245
246
table : SymbolTable :: new ( ) ,
246
247
scopes : vec ! [ root_scope_id] ,
247
248
current_definition : None ,
249
+ current_flow_node : FlowGraph :: start ( ) ,
248
250
} ;
249
251
builder. visit_body ( & module. body ) ;
250
252
builder. table
@@ -257,6 +259,7 @@ impl SymbolTable {
257
259
defs : FxHashMap :: default ( ) ,
258
260
scopes_by_node : FxHashMap :: default ( ) ,
259
261
dependencies : Vec :: new ( ) ,
262
+ flow_graph : FlowGraph :: new ( ) ,
260
263
} ;
261
264
table. scopes_by_id . push ( Scope {
262
265
name : Name :: new ( "<module>" ) ,
@@ -274,6 +277,23 @@ impl SymbolTable {
274
277
& self . dependencies
275
278
}
276
279
280
+ /// Return an iterator over all definitions of `symbol_id` reachable from `use_expr`. The value
281
+ /// of `symbol_id` in `use_expr` must originate from one of the iterated definitions (or from
282
+ /// an external reassignment of the name outside of this scope).
283
+ pub ( crate ) fn reachable_definitions (
284
+ & self ,
285
+ symbol_id : SymbolId ,
286
+ use_expr : & ast:: Expr ,
287
+ ) -> ReachableDefinitionsIterator {
288
+ let node_key = NodeKey :: from_node ( use_expr. into ( ) ) ;
289
+ let flow_node_id = self . flow_graph . ast_to_flow [ & node_key] ;
290
+ ReachableDefinitionsIterator {
291
+ table : self ,
292
+ flow_node_id,
293
+ symbol_id,
294
+ }
295
+ }
296
+
277
297
pub ( crate ) const fn root_scope_id ( ) -> ScopeId {
278
298
ScopeId :: from_usize ( 0 )
279
299
}
@@ -523,11 +543,72 @@ where
523
543
}
524
544
}
525
545
546
+ pub ( crate ) struct ReachableDefinitionsIterator < ' a > {
547
+ table : & ' a SymbolTable ,
548
+ flow_node_id : FlowNodeId ,
549
+ symbol_id : SymbolId ,
550
+ }
551
+
552
+ impl < ' a > Iterator for ReachableDefinitionsIterator < ' a > {
553
+ type Item = Definition ;
554
+
555
+ fn next ( & mut self ) -> Option < Self :: Item > {
556
+ loop {
557
+ match & self . table . flow_graph . flow_nodes_by_id [ self . flow_node_id ] {
558
+ FlowNode :: Start => return None ,
559
+ FlowNode :: Definition ( def_node) => {
560
+ self . flow_node_id = def_node. predecessor ;
561
+ if def_node. symbol_id == self . symbol_id {
562
+ return Some ( def_node. definition . clone ( ) ) ;
563
+ }
564
+ }
565
+ }
566
+ }
567
+ }
568
+ }
569
+
570
+ impl < ' a > FusedIterator for ReachableDefinitionsIterator < ' a > { }
571
+
572
+ #[ newtype_index]
573
+ struct FlowNodeId ;
574
+
575
+ #[ derive( Debug ) ]
576
+ enum FlowNode {
577
+ Start ,
578
+ Definition ( DefinitionFlowNode ) ,
579
+ }
580
+
581
+ #[ derive( Debug ) ]
582
+ struct DefinitionFlowNode {
583
+ symbol_id : SymbolId ,
584
+ definition : Definition ,
585
+ predecessor : FlowNodeId ,
586
+ }
587
+
588
+ #[ derive( Debug , Default ) ]
589
+ struct FlowGraph {
590
+ flow_nodes_by_id : IndexVec < FlowNodeId , FlowNode > ,
591
+ ast_to_flow : FxHashMap < NodeKey , FlowNodeId > ,
592
+ }
593
+
594
+ impl FlowGraph {
595
+ fn new ( ) -> Self {
596
+ let mut graph = FlowGraph :: default ( ) ;
597
+ graph. flow_nodes_by_id . push ( FlowNode :: Start ) ;
598
+ graph
599
+ }
600
+
601
+ fn start ( ) -> FlowNodeId {
602
+ FlowNodeId :: from_usize ( 0 )
603
+ }
604
+ }
605
+
526
606
struct SymbolTableBuilder {
527
607
table : SymbolTable ,
528
608
scopes : Vec < ScopeId > ,
529
609
/// the definition whose target(s) we are currently walking
530
610
current_definition : Option < Definition > ,
611
+ current_flow_node : FlowNodeId ,
531
612
}
532
613
533
614
impl SymbolTableBuilder {
@@ -546,7 +627,16 @@ impl SymbolTableBuilder {
546
627
. defs
547
628
. entry ( symbol_id)
548
629
. or_default ( )
549
- . push ( definition) ;
630
+ . push ( definition. clone ( ) ) ;
631
+ self . current_flow_node = self
632
+ . table
633
+ . flow_graph
634
+ . flow_nodes_by_id
635
+ . push ( FlowNode :: Definition ( DefinitionFlowNode {
636
+ definition,
637
+ symbol_id,
638
+ predecessor : self . current_flow_node ,
639
+ } ) ) ;
550
640
symbol_id
551
641
}
552
642
@@ -561,6 +651,7 @@ impl SymbolTableBuilder {
561
651
self . table
562
652
. add_child_scope ( self . cur_scope ( ) , name, kind, definition, defining_symbol) ;
563
653
self . scopes . push ( scope_id) ;
654
+ self . current_flow_node = FlowGraph :: start ( ) ;
564
655
scope_id
565
656
}
566
657
@@ -624,6 +715,10 @@ impl PreorderVisitor<'_> for SymbolTableBuilder {
624
715
}
625
716
}
626
717
}
718
+ self . table
719
+ . flow_graph
720
+ . ast_to_flow
721
+ . insert ( NodeKey :: from_node ( expr. into ( ) ) , self . current_flow_node ) ;
627
722
ast:: visitor:: preorder:: walk_expr ( self , expr) ;
628
723
}
629
724
@@ -766,15 +861,13 @@ impl DerefMut for SymbolTablesStorage {
766
861
767
862
#[ cfg( test) ]
768
863
mod tests {
769
- use textwrap:: dedent;
770
-
771
- use crate :: parse:: Parsed ;
772
- use crate :: symbols:: ScopeKind ;
773
-
774
- use super :: { SymbolFlags , SymbolId , SymbolIterator , SymbolTable } ;
864
+ use crate :: symbols:: { ScopeKind , SymbolFlags , SymbolTable } ;
775
865
776
866
mod from_ast {
777
- use super :: * ;
867
+ use crate :: parse:: Parsed ;
868
+ use crate :: symbols:: { Definition , ScopeKind , SymbolId , SymbolIterator , SymbolTable } ;
869
+ use ruff_python_ast as ast;
870
+ use textwrap:: dedent;
778
871
779
872
fn parse ( code : & str ) -> Parsed {
780
873
Parsed :: from_text ( & dedent ( code) )
@@ -1021,6 +1114,35 @@ mod tests {
1021
1114
assert_eq ! ( func_scope. name( ) , "C" ) ;
1022
1115
assert_eq ! ( names( table. symbols_for_scope( func_scope_id) ) , vec![ "x" ] ) ;
1023
1116
}
1117
+
1118
+ #[ test]
1119
+ fn reachability_trivial ( ) {
1120
+ let parsed = parse ( "x = 1; x" ) ;
1121
+ let ast = parsed. ast ( ) ;
1122
+ let table = SymbolTable :: from_ast ( ast) ;
1123
+ let x_sym = table
1124
+ . root_symbol_id_by_name ( "x" )
1125
+ . expect ( "x symbol should exist" ) ;
1126
+ let ast:: Stmt :: Expr ( ast:: StmtExpr { value : x_use, .. } ) = & ast. body [ 1 ] else {
1127
+ panic ! ( "should be an expr" )
1128
+ } ;
1129
+ let x_defs: Vec < _ > = table. reachable_definitions ( x_sym, x_use) . collect ( ) ;
1130
+ assert_eq ! ( x_defs. len( ) , 1 ) ;
1131
+ let Definition :: Assignment ( node_key) = & x_defs[ 0 ] else {
1132
+ panic ! ( "def should be an assignment" )
1133
+ } ;
1134
+ let Some ( def_node) = node_key. resolve ( ast. into ( ) ) else {
1135
+ panic ! ( "node key should resolve" )
1136
+ } ;
1137
+ let ast:: Expr :: NumberLiteral ( ast:: ExprNumberLiteral {
1138
+ value : ast:: Number :: Int ( num) ,
1139
+ ..
1140
+ } ) = & * def_node. value
1141
+ else {
1142
+ panic ! ( "should be a number literal" )
1143
+ } ;
1144
+ assert_eq ! ( * num, 1 ) ;
1145
+ }
1024
1146
}
1025
1147
1026
1148
#[ test]
0 commit comments