Skip to content

Commit 27f6f04

Browse files
authored
[red-knot] initial (very incomplete) flow graph (#11624)
<!-- Thank you for contributing to Ruff! To help us out with reviewing, please consider the following: - Does this pull request include a summary of the change? (See below.) - Does this pull request include a descriptive title? - Does this pull request include references to any relevant issues? --> ## Summary Introduces the skeleton of the flow graph. So far it doesn't actually handle any non-linear control flow :) But it does show how we can go from an expression that references a symbol, backward through the flow graph, to find reachable definitions of that symbol. Adding non-linear control flow will mean adding flow nodes with multiple predecessors, which will introduce more complexity into `ReachableDefinitionsIterator.next()`. But one step at a time. ## Test Plan Added a (very basic) test.
1 parent d62a617 commit 27f6f04

File tree

2 files changed

+137
-12
lines changed

2 files changed

+137
-12
lines changed

crates/red_knot/src/ast_ids.rs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -275,10 +275,7 @@ pub struct TypedNodeKey<N: AstNode> {
275275

276276
impl<N: AstNode> TypedNodeKey<N> {
277277
pub fn from_node(node: &N) -> Self {
278-
let inner = NodeKey {
279-
kind: node.as_any_node_ref().kind(),
280-
range: node.range(),
281-
};
278+
let inner = NodeKey::from_node(node.as_any_node_ref());
282279
Self {
283280
inner,
284281
_marker: PhantomData,
@@ -352,6 +349,12 @@ pub struct NodeKey {
352349
}
353350

354351
impl NodeKey {
352+
pub fn from_node(node: AnyNodeRef) -> Self {
353+
NodeKey {
354+
kind: node.kind(),
355+
range: node.range(),
356+
}
357+
}
355358
pub fn resolve<'a>(&self, root: AnyNodeRef<'a>) -> Option<AnyNodeRef<'a>> {
356359
// We need to do a binary search here. Only traverse into a node if the range is withint the node
357360
let mut visitor = FindNodeKeyVisitor {

crates/red_knot/src/symbols.rs

Lines changed: 130 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,7 @@ pub struct SymbolTable {
236236
scopes_by_node: FxHashMap<NodeKey, ScopeId>,
237237
/// dependencies of this module
238238
dependencies: Vec<Dependency>,
239+
flow_graph: FlowGraph,
239240
}
240241

241242
impl SymbolTable {
@@ -245,6 +246,7 @@ impl SymbolTable {
245246
table: SymbolTable::new(),
246247
scopes: vec![root_scope_id],
247248
current_definition: None,
249+
current_flow_node: FlowGraph::start(),
248250
};
249251
builder.visit_body(&module.body);
250252
builder.table
@@ -257,6 +259,7 @@ impl SymbolTable {
257259
defs: FxHashMap::default(),
258260
scopes_by_node: FxHashMap::default(),
259261
dependencies: Vec::new(),
262+
flow_graph: FlowGraph::new(),
260263
};
261264
table.scopes_by_id.push(Scope {
262265
name: Name::new("<module>"),
@@ -274,6 +277,23 @@ impl SymbolTable {
274277
&self.dependencies
275278
}
276279

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+
277297
pub(crate) const fn root_scope_id() -> ScopeId {
278298
ScopeId::from_usize(0)
279299
}
@@ -523,11 +543,72 @@ where
523543
}
524544
}
525545

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+
526606
struct SymbolTableBuilder {
527607
table: SymbolTable,
528608
scopes: Vec<ScopeId>,
529609
/// the definition whose target(s) we are currently walking
530610
current_definition: Option<Definition>,
611+
current_flow_node: FlowNodeId,
531612
}
532613

533614
impl SymbolTableBuilder {
@@ -546,7 +627,16 @@ impl SymbolTableBuilder {
546627
.defs
547628
.entry(symbol_id)
548629
.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+
}));
550640
symbol_id
551641
}
552642

@@ -561,6 +651,7 @@ impl SymbolTableBuilder {
561651
self.table
562652
.add_child_scope(self.cur_scope(), name, kind, definition, defining_symbol);
563653
self.scopes.push(scope_id);
654+
self.current_flow_node = FlowGraph::start();
564655
scope_id
565656
}
566657

@@ -624,6 +715,10 @@ impl PreorderVisitor<'_> for SymbolTableBuilder {
624715
}
625716
}
626717
}
718+
self.table
719+
.flow_graph
720+
.ast_to_flow
721+
.insert(NodeKey::from_node(expr.into()), self.current_flow_node);
627722
ast::visitor::preorder::walk_expr(self, expr);
628723
}
629724

@@ -766,15 +861,13 @@ impl DerefMut for SymbolTablesStorage {
766861

767862
#[cfg(test)]
768863
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};
775865

776866
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;
778871

779872
fn parse(code: &str) -> Parsed {
780873
Parsed::from_text(&dedent(code))
@@ -1021,6 +1114,35 @@ mod tests {
10211114
assert_eq!(func_scope.name(), "C");
10221115
assert_eq!(names(table.symbols_for_scope(func_scope_id)), vec!["x"]);
10231116
}
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+
}
10241146
}
10251147

10261148
#[test]

0 commit comments

Comments
 (0)