Skip to content

Commit d12473e

Browse files
mtshibasharkdp
mtshiba
and
sharkdp
authored
[ty] Support basic narrowing with aliased conditional expressions (#24302)
Co-authored-by: David Peter <[email protected]>
1 parent 76109a9 commit d12473e

5 files changed

Lines changed: 1018 additions & 18 deletions

File tree

‎crates/ty_python_core/src/builder.rs‎

Lines changed: 209 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,8 @@ use crate::use_def::{
5858
};
5959
use crate::{Db, Statement, StatementNodeKey};
6060
use crate::{
61-
EvaluationMode, ExpressionsScopeMap, LoopHeader, LoopToken, PossiblyNarrowedPlaces,
62-
SemanticIndex, VisibleAncestorsIter, get_loop_header,
61+
EvaluationMode, ExpressionsScopeMap, LoopHeader, LoopToken, NarrowingAliasPredicate,
62+
PossiblyNarrowedPlaces, SemanticIndex, VisibleAncestorsIter, get_loop_header,
6363
};
6464

6565
use super::place::PlaceExprRef;
@@ -85,10 +85,24 @@ impl Loop {
8585
}
8686
}
8787

88-
struct ScopeInfo {
88+
/// A narrowing alias: a variable whose RHS is a narrowing expression
89+
/// (e.g., `is_none = x is None`).
90+
#[derive(Clone, Debug)]
91+
struct NarrowingAlias<'ast> {
92+
/// The RHS expression (e.g., `x is None`).
93+
expression: &'ast ast::Expr,
94+
/// The scope whose place table should be used to resolve the aliased expression.
95+
expression_scope: FileScopeId,
96+
/// Places that, if reassigned, should invalidate this alias.
97+
narrowed_places: PossiblyNarrowedPlaces,
98+
}
99+
100+
struct ScopeInfo<'ast> {
89101
file_scope_id: FileScopeId,
90102
/// Current loop state; None if we are not currently visiting a loop
91103
current_loop: Option<Loop>,
104+
/// Saved narrowing aliases from the enclosing scope, restored on `pop_scope`.
105+
narrowing_aliases: FxHashMap<Name, NarrowingAlias<'ast>>,
92106
}
93107

94108
pub(super) struct SemanticIndexBuilder<'db, 'ast> {
@@ -97,7 +111,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
97111
file: File,
98112
source_type: PySourceType,
99113
module: &'ast ParsedModuleRef,
100-
scope_stack: Vec<ScopeInfo>,
114+
scope_stack: Vec<ScopeInfo<'ast>>,
101115
/// The assignments we're currently visiting, with
102116
/// the most recent visit at the end of the Vec.
103117
current_assignments: Vec<CurrentAssignment<'ast, 'db>>,
@@ -145,6 +159,13 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
145159
enclosing_snapshots: FxHashMap<EnclosingSnapshotKey, ScopedEnclosingSnapshotId>,
146160
/// Errors collected by the `semantic_checker`.
147161
semantic_syntax_errors: RefCell<Vec<SemanticSyntaxError>>,
162+
163+
/// Maps alias variable names to their narrowing expressions (same-scope only).
164+
/// TODO: cross-scope alias narrowing support
165+
narrowing_aliases: FxHashMap<Name, NarrowingAlias<'ast>>,
166+
167+
/// Alias metadata for predicate leaf names in the current file.
168+
alias_predicates: FxHashMap<ExpressionNodeKey, NarrowingAliasPredicate<'db>>,
148169
}
149170

150171
impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
@@ -188,19 +209,21 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
188209
semantic_checker: SemanticSyntaxChecker::default(),
189210
in_try: false,
190211
semantic_syntax_errors: RefCell::default(),
212+
narrowing_aliases: FxHashMap::default(),
213+
alias_predicates: FxHashMap::default(),
191214
};
192215

193216
builder.push_scope_with_parent(NodeWithScopeRef::Module, None);
194217
builder
195218
}
196219

197-
fn current_scope_info(&self) -> &ScopeInfo {
220+
fn current_scope_info(&self) -> &ScopeInfo<'ast> {
198221
self.scope_stack
199222
.last()
200223
.expect("SemanticIndexBuilder should have created a root scope")
201224
}
202225

203-
fn current_scope_info_mut(&mut self) -> &mut ScopeInfo {
226+
fn current_scope_info_mut(&mut self) -> &mut ScopeInfo<'ast> {
204227
self.scope_stack
205228
.last_mut()
206229
.expect("SemanticIndexBuilder should have created a root scope")
@@ -349,9 +372,13 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
349372

350373
debug_assert_eq!(ast_id_scope, file_scope_id);
351374

375+
// Save narrowing aliases. They will be restored with `pop_scope` after returning from inspecting the inner scope.
376+
// TODO: Cross-scope alias narrowing is not supported yet.
377+
let saved_aliases = std::mem::take(&mut self.narrowing_aliases);
352378
self.scope_stack.push(ScopeInfo {
353379
file_scope_id,
354380
current_loop: None,
381+
narrowing_aliases: saved_aliases,
355382
});
356383
}
357384

@@ -655,11 +682,13 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
655682

656683
let ScopeInfo {
657684
file_scope_id: popped_scope_id,
685+
narrowing_aliases,
658686
..
659687
} = self
660688
.scope_stack
661689
.pop()
662690
.expect("Root scope should be present");
691+
self.narrowing_aliases = narrowing_aliases;
663692

664693
let children_end = self.scopes.next_index();
665694

@@ -705,6 +734,165 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
705734
&mut self.ast_ids[scope_id]
706735
}
707736

737+
/// Try to register a narrowing alias for a simple name assignment.
738+
///
739+
/// Any pre-existing alias entry for the `target` name has already been removed by
740+
/// [`Self::invalidate_narrowing_aliases_for`] in the binding pathway that ran before
741+
/// this call, so we only need to decide whether to insert a new entry.
742+
fn try_register_narrowing_alias(&mut self, target: &ast::Expr, value: Option<&'ast ast::Expr>) {
743+
let Some(target_name_expr) = target.as_name_expr() else {
744+
return;
745+
};
746+
let Some(value) = value else { return };
747+
let target_name = &target_name_expr.id;
748+
749+
if !Self::can_register_narrowing_alias(value) {
750+
return;
751+
}
752+
753+
let place_table = self.current_place_table();
754+
let narrowed_places =
755+
PossiblyNarrowedPlacesBuilder::new(self.db, place_table).expression(value);
756+
757+
// Don't register if the target itself is one of the narrowed places (e.g. `x = x is None`),
758+
// since the alias would be invalidated immediately by this same assignment.
759+
let target_is_narrowed = place_table
760+
.symbol_id(target_name)
761+
.is_some_and(|symbol| narrowed_places.contains(&symbol.into()));
762+
763+
if !narrowed_places.is_empty() && !target_is_narrowed {
764+
self.narrowing_aliases.insert(
765+
target_name.clone(),
766+
NarrowingAlias {
767+
expression: value,
768+
expression_scope: self.current_scope(),
769+
narrowed_places,
770+
},
771+
);
772+
}
773+
}
774+
775+
/// Invalidate any narrowing aliases affected by a new definition of `place`.
776+
fn invalidate_narrowing_aliases_for(&mut self, place: ScopedPlaceId) {
777+
let place_table = &self.place_tables[self.current_scope()];
778+
let associated_members = place_table.associated_place_ids(place);
779+
let reassigned_alias_name = place
780+
.as_symbol()
781+
.map(|symbol_id| place_table.symbol(symbol_id).name());
782+
783+
self.narrowing_aliases.retain(|name, alias| {
784+
// Drop aliases that narrow the reassigned place or any of its members.
785+
// e.g. `is_none = x is None and ...; x = 1`
786+
!alias.narrowed_places.contains(&place)
787+
// e.g. `is_none = a.x is None; a = A()`
788+
&& !associated_members
789+
.iter()
790+
.any(|m| alias.narrowed_places.contains(&(*m).into()))
791+
// Drop the alias whose own variable is the reassigned place.
792+
// e.g. `is_none = x is None; is_none = False`
793+
&& reassigned_alias_name != Some(name)
794+
});
795+
}
796+
797+
fn can_register_narrowing_alias(value: &ast::Expr) -> bool {
798+
match value {
799+
// Bare names are too common to treat as alias candidates on every assignment,
800+
// and doing so would noticeably degrade performance. Excluding them only means
801+
// we don't infer truthiness narrowing for arbitrary chained aliases.
802+
ast::Expr::Name(_) => false,
803+
ast::Expr::Compare(_) | ast::Expr::Call(_) => true,
804+
ast::Expr::UnaryOp(unary) if unary.op == ast::UnaryOp::Not => {
805+
Self::can_register_narrowing_alias(&unary.operand)
806+
}
807+
ast::Expr::BoolOp(bool_op) => bool_op
808+
.values
809+
.iter()
810+
.any(Self::can_register_narrowing_alias),
811+
ast::Expr::If(expr_if) => {
812+
Self::can_register_narrowing_alias(&expr_if.test)
813+
|| Self::can_register_narrowing_alias(&expr_if.body)
814+
|| Self::can_register_narrowing_alias(&expr_if.orelse)
815+
}
816+
_ => false,
817+
}
818+
}
819+
820+
/// Walk a predicate expression tree, calling `f` on each leaf position
821+
/// where an alias Name could appear.
822+
fn walk_narrowing_alias_predicate<'expr>(
823+
expr: &'expr ast::Expr,
824+
f: &mut impl FnMut(&'expr ast::Expr),
825+
) {
826+
match expr {
827+
ast::Expr::Name(_) => f(expr),
828+
ast::Expr::UnaryOp(unary) if unary.op == ast::UnaryOp::Not => {
829+
Self::walk_narrowing_alias_predicate(&unary.operand, f);
830+
}
831+
ast::Expr::BoolOp(bool_op) => {
832+
for value in &bool_op.values {
833+
Self::walk_narrowing_alias_predicate(value, f);
834+
}
835+
}
836+
ast::Expr::Call(call) => {
837+
for arg in &call.arguments.args {
838+
Self::walk_narrowing_alias_predicate(arg, f);
839+
}
840+
for keyword in &call.arguments.keywords {
841+
Self::walk_narrowing_alias_predicate(&keyword.value, f);
842+
}
843+
}
844+
ast::Expr::If(expr_if) => {
845+
Self::walk_narrowing_alias_predicate(&expr_if.test, f);
846+
Self::walk_narrowing_alias_predicate(&expr_if.body, f);
847+
Self::walk_narrowing_alias_predicate(&expr_if.orelse, f);
848+
}
849+
_ => {}
850+
}
851+
}
852+
853+
/// Register alias predicates for alias Names found in a predicate expression.
854+
fn register_narrowing_alias_predicates(&mut self, expr: &'ast ast::Expr) {
855+
Self::walk_narrowing_alias_predicate(expr, &mut |leaf| {
856+
let Some(name) = leaf.as_name_expr() else {
857+
return;
858+
};
859+
let Some(alias) = self.narrowing_aliases.get(&name.id).cloned() else {
860+
return;
861+
};
862+
let aliased_expression = Expression::new(
863+
self.db,
864+
self.file,
865+
alias.expression_scope,
866+
AstNodeRef::new(self.module, alias.expression),
867+
None,
868+
ExpressionKind::Normal,
869+
);
870+
self.alias_predicates.insert(
871+
ExpressionNodeKey::from(leaf),
872+
NarrowingAliasPredicate {
873+
expression: aliased_expression,
874+
},
875+
);
876+
});
877+
}
878+
879+
/// Add narrowed places from aliased expressions to the possibly-narrowed set.
880+
fn add_alias_narrowed_places(&self, expr: &ast::Expr, places: &mut PossiblyNarrowedPlaces) {
881+
Self::walk_narrowing_alias_predicate(expr, &mut |leaf| {
882+
let key = ExpressionNodeKey::from(leaf);
883+
if let Some(alias_predicate) = self.alias_predicates.get(&key) {
884+
let aliased_node = alias_predicate
885+
.expression
886+
.node_ref(self.db)
887+
.node(self.module);
888+
let aliased_places =
889+
PossiblyNarrowedPlacesBuilder::new(self.db, self.current_place_table())
890+
.expression(aliased_node);
891+
places.extend(aliased_places);
892+
}
893+
});
894+
}
895+
708896
fn flow_snapshot(&self) -> FlowSnapshot {
709897
self.current_use_def_map().snapshot()
710898
}
@@ -941,6 +1129,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
9411129
// ```
9421130
if category.is_binding() && !is_loop_header {
9431131
self.mark_place_bound(place);
1132+
self.invalidate_narrowing_aliases_for(place);
9441133
}
9451134
if category.is_declaration() {
9461135
self.mark_place_declared(place);
@@ -1139,14 +1328,14 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
11391328

11401329
fn record_expression_narrowing_constraint(
11411330
&mut self,
1142-
predicate_node: &ast::Expr,
1331+
predicate_node: &'ast ast::Expr,
11431332
) -> (PredicateOrLiteral<'db>, ScopedPredicateId) {
11441333
let predicate = self.build_predicate(predicate_node);
11451334
let predicate_id = self.record_narrowing_constraint(predicate);
11461335
(predicate, predicate_id)
11471336
}
11481337

1149-
fn build_predicate(&mut self, predicate_node: &ast::Expr) -> PredicateOrLiteral<'db> {
1338+
fn build_predicate(&mut self, predicate_node: &'ast ast::Expr) -> PredicateOrLiteral<'db> {
11501339
// Some commonly used test expressions are eagerly evaluated as `true`
11511340
// or `false` here for performance reasons. This list does not need to
11521341
// be exhaustive. More complex expressions will still evaluate to the
@@ -1170,6 +1359,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
11701359
}
11711360
}
11721361

1362+
self.register_narrowing_alias_predicates(predicate_node);
1363+
11731364
let expression = self.add_standalone_expression(predicate_node);
11741365

11751366
match resolve_to_literal(predicate_node) {
@@ -1236,8 +1427,10 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
12361427
match pred.node {
12371428
PredicateNode::Expression(expression) => {
12381429
let expression_node = expression.node_ref(self.db).node(self.module);
1239-
PossiblyNarrowedPlacesBuilder::new(self.db, place_table)
1240-
.expression(expression_node)
1430+
let mut places = PossiblyNarrowedPlacesBuilder::new(self.db, place_table)
1431+
.expression(expression_node);
1432+
self.add_alias_narrowed_places(expression_node, &mut places);
1433+
places
12411434
}
12421435
PredicateNode::Pattern(pattern) => {
12431436
let module = self.module;
@@ -1951,6 +2144,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
19512144
enclosing_snapshots: self.enclosing_snapshots,
19522145
semantic_syntax_errors: self.semantic_syntax_errors.into_inner(),
19532146
generator_functions: self.generator_functions,
2147+
narrowing_alias_predicates: self.alias_predicates,
19542148
}
19552149
}
19562150

@@ -2373,6 +2567,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
23732567
self.push_assignment(CurrentAssignment::Assign { node, unpack: None });
23742568
self.visit_expr(target);
23752569
self.pop_assignment();
2570+
2571+
self.try_register_narrowing_alias(target, Some(&node.value));
23762572
} else {
23772573
let value = self.add_standalone_assigned_expression(&node.value, node);
23782574

@@ -2426,6 +2622,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
24262622
self.push_assignment(node.into());
24272623
self.visit_expr(&node.target);
24282624
self.pop_assignment();
2625+
2626+
self.try_register_narrowing_alias(&node.target, node.value.as_deref());
24292627
} else {
24302628
self.visit_expr(&node.target);
24312629
}
@@ -3093,6 +3291,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
30933291
}
30943292

30953293
let place_id = self.add_place(target);
3294+
self.invalidate_narrowing_aliases_for(place_id);
30963295
self.delete_binding(place_id);
30973296
}
30983297
}

‎crates/ty_python_core/src/lib.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,17 @@ pub struct SemanticIndex<'db> {
307307

308308
/// Set of all generator functions in this file.
309309
generator_functions: FxHashSet<FileScopeId>,
310+
311+
/// Narrowing alias metadata for predicate leaf names.
312+
/// When a predicate references an alias variable (e.g., `is_none` from `is_none = x is None`),
313+
/// the alias Name node is mapped to its aliased expression for constraint-generation time.
314+
narrowing_alias_predicates: FxHashMap<ExpressionNodeKey, NarrowingAliasPredicate<'db>>,
315+
}
316+
317+
#[derive(Debug, Clone, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
318+
pub struct NarrowingAliasPredicate<'db> {
319+
/// Aliased expression, e.g., `x is None` in `is_none = x is None`.
320+
pub expression: Expression<'db>,
310321
}
311322

312323
impl<'db> SemanticIndex<'db> {
@@ -319,6 +330,14 @@ impl<'db> SemanticIndex<'db> {
319330
&self.place_tables[scope_id]
320331
}
321332

333+
/// Returns alias metadata for an alias Name node in a predicate, if one exists.
334+
pub fn narrowing_alias_predicate(
335+
&self,
336+
key: impl Into<ExpressionNodeKey>,
337+
) -> Option<&NarrowingAliasPredicate<'db>> {
338+
self.narrowing_alias_predicates.get(&key.into())
339+
}
340+
322341
/// Returns the use-def map for a specific scope.
323342
///
324343
/// Use the Salsa cached [`use_def_map()`] query if you only need the

0 commit comments

Comments
 (0)