From 050ae6251e1103486f4798fd4d8135fe2472470d Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Wed, 1 Jul 2026 14:39:31 +0900 Subject: [PATCH] fix fix fix flow? --- pyrefly/benches/micro.rs | 2 +- pyrefly/lib/binding/bindings.rs | 1 + pyrefly/lib/binding/narrow.rs | 158 ++++++++++++++++++++++++++++- pyrefly/lib/binding/scope.rs | 134 ++++++++++++++++++++++++ pyrefly/lib/binding/stmt.rs | 20 +++- pyrefly/lib/test/flow_branching.rs | 37 +++++-- 6 files changed, 339 insertions(+), 13 deletions(-) diff --git a/pyrefly/benches/micro.rs b/pyrefly/benches/micro.rs index 5e76cfac0e..f066313e88 100644 --- a/pyrefly/benches/micro.rs +++ b/pyrefly/benches/micro.rs @@ -319,7 +319,7 @@ fn overload_resolution(overloads: usize, calls: usize) -> String { } src.push_str("def choose(x): return x\n"); for i in 0..calls { - let _ = write!(src, "r{i} = choose({})\n", ARGS[i % ARGS.len()]); + let _ = writeln!(src, "r{i} = choose({})", ARGS[i % ARGS.len()]); } src } diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index a39712e96a..fbc214d675 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -2646,6 +2646,7 @@ impl<'a> BindingsBuilder<'a> { self.scopes.narrow_in_current_flow(name, narrowed_idx); } } + self.scopes.apply_conditional_flow_facts(narrow_ops); } pub fn bind_lambda_param(&mut self, name: &Identifier, kind: LambdaKind, usage: &Usage) { diff --git a/pyrefly/lib/binding/narrow.rs b/pyrefly/lib/binding/narrow.rs index 108e670bf0..7b97bf16cf 100644 --- a/pyrefly/lib/binding/narrow.rs +++ b/pyrefly/lib/binding/narrow.rs @@ -31,6 +31,8 @@ use ruff_python_ast::ExprUnaryOp; use ruff_python_ast::Identifier; use ruff_python_ast::Number; use ruff_python_ast::UnaryOp; +use ruff_python_ast::comparable::ComparableArguments; +use ruff_python_ast::comparable::ComparableExpr; use ruff_python_ast::name::Name; use ruff_text_size::Ranged; use ruff_text_size::TextRange; @@ -54,7 +56,7 @@ assert_words!(NarrowOp, 11); /// Indicates where an isinstance-style narrow operation originated from. /// This determines whether validation needs to happen during narrowing. -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum NarrowSource { /// From an isinstance() call - validation already happened in special_calls.rs. Call, @@ -290,6 +292,84 @@ impl DisplayWith for NarrowOp { } impl AtomicNarrowOp { + /// Compares narrowing semantics while ignoring AST source locations and trivia. + fn semantically_eq(&self, other: &Self) -> bool { + use AtomicNarrowOp::*; + + let expr_eq = + |left: &Expr, right: &Expr| ComparableExpr::from(left) == ComparableExpr::from(right); + let optional_expr_eq = |left: &Option>, right: &Option>| { + left.as_deref().map(ComparableExpr::from) == right.as_deref().map(ComparableExpr::from) + }; + let arguments_eq = |left: &Arguments, right: &Arguments| { + ComparableArguments::from(left) == ComparableArguments::from(right) + }; + + match (self, other) { + (Is(left), Is(right)) + | (IsNot(left), IsNot(right)) + | (Eq(left), Eq(right)) + | (NotEq(left), NotEq(right)) + | (IsSubclass(left), IsSubclass(right)) + | (IsNotSubclass(left), IsNotSubclass(right)) + | (TypeEq(left), TypeEq(right)) + | (TypeNotEq(left), TypeNotEq(right)) + | (In(left), In(right)) + | (NotIn(left), NotIn(right)) + | (LenEq(left), LenEq(right)) + | (LenNotEq(left), LenNotEq(right)) + | (LenGt(left), LenGt(right)) + | (LenGte(left), LenGte(right)) + | (LenLt(left), LenLt(right)) + | (LenLte(left), LenLte(right)) => expr_eq(left, right), + (IsInstance(left, left_source), IsInstance(right, right_source)) + | (IsNotInstance(left, left_source), IsNotInstance(right, right_source)) => { + left_source == right_source && expr_eq(left, right) + } + (HasAttr(left), HasAttr(right)) + | (NotHasAttr(left), NotHasAttr(right)) + | (HasKey(left), HasKey(right)) + | (NotHasKey(left), NotHasKey(right)) => left == right, + (GetAttr(left_name, left_default), GetAttr(right_name, right_default)) + | (NotGetAttr(left_name, left_default), NotGetAttr(right_name, right_default)) => { + left_name == right_name && optional_expr_eq(left_default, right_default) + } + (TypeGuard(left_type, left_args), TypeGuard(right_type, right_args)) + | (NotTypeGuard(left_type, left_args), NotTypeGuard(right_type, right_args)) + | (TypeIs(left_type, left_args), TypeIs(right_type, right_args)) + | (NotTypeIs(left_type, left_args), NotTypeIs(right_type, right_args)) => { + left_type == right_type && arguments_eq(left_args, right_args) + } + (Call(left_expr, left_args), Call(right_expr, right_args)) + | (NotCall(left_expr, left_args), NotCall(right_expr, right_args)) => { + expr_eq(left_expr, right_expr) && arguments_eq(left_args, right_args) + } + (PolarsColumnMutation(left), PolarsColumnMutation(right)) => match (left, right) { + (PolarsMutationKind::Add, PolarsMutationKind::Add) + | (PolarsMutationKind::Replace, PolarsMutationKind::Replace) => true, + ( + PolarsMutationKind::Insert(left_name, left_index, left_expr), + PolarsMutationKind::Insert(right_name, right_index, right_expr), + ) => { + left_name == right_name + && left_index == right_index + && expr_eq(left_expr, right_expr) + } + _ => false, + }, + (ClassCoverageGate(left), ClassCoverageGate(right)) + | (ClassCoverageGateNeg(left), ClassCoverageGateNeg(right)) => left == right, + (IsSequence, IsSequence) + | (IsNotSequence, IsNotSequence) + | (IsMapping, IsMapping) + | (IsNotMapping, IsNotMapping) + | (IsTruthy, IsTruthy) + | (IsFalsy, IsFalsy) + | (Placeholder, Placeholder) => true, + _ => false, + } + } + /// Produce a Python-like snippet for hover display. /// /// `subject` is the name (possibly with facet chain) being narrowed. @@ -452,6 +532,49 @@ pub struct FacetSubject { pub allow_never_collapse: bool, } +impl FacetSubject { + /// Compares facet semantics while ignoring locations in embedded AST nodes. + fn semantically_eq(&self, other: &Self) -> bool { + self.origin == other.origin + && self.chain.facets().len() == other.chain.facets().len() + && self + .chain + .facets() + .iter() + .zip(other.chain.facets()) + .all(|(left, right)| match (left, right) { + ( + UnresolvedFacetKind::Attribute(left), + UnresolvedFacetKind::Attribute(right), + ) => left == right, + (UnresolvedFacetKind::Index(left), UnresolvedFacetKind::Index(right)) => { + left == right + } + (UnresolvedFacetKind::Key(left), UnresolvedFacetKind::Key(right)) => { + left == right + } + ( + UnresolvedFacetKind::VariableSubscript(left), + UnresolvedFacetKind::VariableSubscript(right), + ) => left.id == right.id, + ( + UnresolvedFacetKind::MatchArg { + class: left_class, + index: left_index, + }, + UnresolvedFacetKind::MatchArg { + class: right_class, + index: right_index, + }, + ) => { + left_index == right_index + && ComparableExpr::from(left_class) == ComparableExpr::from(right_class) + } + _ => false, + }) + } +} + #[derive(Clone, Debug)] pub enum NarrowingSubject { Name(Name), @@ -504,6 +627,28 @@ impl NarrowOp { } } + /// Compares narrowing semantics while ignoring AST source locations and trivia. + pub(crate) fn semantically_eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Atomic(left_facet, left_op), Self::Atomic(right_facet, right_op)) => { + match (left_facet, right_facet) { + (Some(left), Some(right)) if left.semantically_eq(right) => {} + (None, None) => {} + _ => return false, + } + left_op.semantically_eq(right_op) + } + (Self::And(left), Self::And(right)) | (Self::Or(left), Self::Or(right)) => { + left.len() == right.len() + && left + .iter() + .zip(right) + .all(|(left, right)| left.semantically_eq(right)) + } + _ => false, + } + } + /// Produce a Python-like snippet for hover display. /// /// `base_name` is the variable being narrowed. `snippet` converts a @@ -874,6 +1019,17 @@ impl NarrowOps { Self(SmallMap::new()) } + /// Compares narrowing semantics without considering the source range of each operation. + pub(crate) fn semantically_eq(&self, other: &Self) -> bool { + self.0.len() == other.0.len() + && self.0.iter().all(|(name, (op, _))| { + other + .0 + .get(name) + .is_some_and(|(other_op, _)| op.semantically_eq(other_op)) + }) + } + /// Adds or merges a narrowing operation for the given subject. /// /// If the subject's name already has a narrowing operation, combines them with AND. diff --git a/pyrefly/lib/binding/scope.rs b/pyrefly/lib/binding/scope.rs index 084cf0a01c..9465983e07 100644 --- a/pyrefly/lib/binding/scope.rs +++ b/pyrefly/lib/binding/scope.rs @@ -7,6 +7,7 @@ use std::cmp::max; use std::fmt::Debug; +use std::iter; use std::mem; use itertools::Either; @@ -651,6 +652,7 @@ pub struct Flow { info: SmallMap, /// Pairs for which the current flow excludes both values being `None`. excluded_none_pairs: ExcludedNonePairs, + conditional_facts: Vec, // Have we seen control flow terminate? // // We continue to analyze the rest of the code after a flow terminates, but @@ -684,6 +686,7 @@ impl Flow { in_loop: bool, ) { self.excluded_none_pairs.invalidate(name.key()); + self.invalidate_conditional_facts_for_assignment(name); match self.info.entry_hashed(name.cloned()) { Entry::Vacant(e) => { e.insert(FlowInfo::new_value(idx, style)); @@ -713,6 +716,12 @@ impl Flow { fn get_value_mut(&mut self, name: &Name) -> Option<&mut FlowValue> { self.info.get_mut(name)?.value_mut() } + + fn invalidate_conditional_facts_for_assignment(&mut self, name: Hashed<&Name>) { + self.conditional_facts.retain(|fact| { + fact.name != **name.key() && fact.conditions.0.get_hashed(name).is_none() + }); + } } /// Copy new or newly-initialized entries from `source` into `target`. @@ -762,6 +771,36 @@ struct FlowInfo { loop_prior: Idx, } +/// A value known to hold when a previously-analyzed predicate is active again. +/// +/// For `if b == "bar": x = 3`, the branch records that under `b == "bar"`, +/// `x` has the branch-end flow info. The fact is discarded if `x` or `b` is +/// rebound before the predicate is used again. +#[derive(Debug, Clone)] +pub(crate) struct ConditionalFlowFact { + conditions: NarrowOps, + name: Name, + info: FlowInfo, +} + +impl ConditionalFlowFact { + fn matches_active_conditions(&self, active: &NarrowOps) -> bool { + !self.conditions.0.is_empty() + && self.conditions.0.iter().all(|(name, (op, _))| { + active + .0 + .get(name) + .is_some_and(|(active_op, _)| active_op.semantically_eq(op)) + }) + } + + fn same_fact(&self, other: &Self) -> bool { + self.name == other.name + && self.conditions.semantically_eq(&other.conditions) + && self.info.idx() == other.info.idx() + } +} + /// The most recent value for a name. Used in several cases: /// - Actual runtime assignments /// - Certain cases where we track a type for unbound locals, such as after a bare @@ -2677,6 +2716,9 @@ impl Scopes { pub fn mark_as_deleted(&mut self, name: &Name) { let scope = self.current_mut(); scope.flow.excluded_none_pairs.invalidate(name); + scope + .flow + .invalidate_conditional_facts_for_assignment(Hashed::new(name)); if let Some(value) = scope.flow.get_value_mut(name) { value.style = FlowStyle::Uninitialized; } @@ -2721,6 +2763,73 @@ impl Scopes { Some(self.current().forks.last()?.base.get_info(name)?.idx()) } + pub(crate) fn conditional_flow_facts_for_current_branch( + &self, + conditions: &NarrowOps, + ) -> Vec { + if conditions.0.is_empty() { + return Vec::new(); + } + let scope = self.current(); + let fork = scope + .forks + .last() + .expect("conditional flow facts requested outside of a fork"); + + for name in conditions.0.keys() { + let base_value = fork.base.get_info(name).and_then(|info| info.value()); + let current_value = scope.flow.get_info(name).and_then(|info| info.value()); + if base_value.map(|value| value.idx) != current_value.map(|value| value.idx) { + return Vec::new(); + } + } + + scope + .flow + .info + .iter() + .filter_map(|(name, info)| { + let value = info.value()?; + let base_value = fork.base.get_info(name).and_then(|info| info.value()); + if base_value.is_some_and(|base_value| base_value.idx == value.idx) { + None + } else { + Some(ConditionalFlowFact { + conditions: conditions.clone(), + name: name.clone(), + info: info.clone(), + }) + } + }) + .collect() + } + + pub(crate) fn add_conditional_flow_facts(&mut self, facts: Vec) { + let flow = &mut self.current_mut().flow; + for fact in facts { + if !flow + .conditional_facts + .iter() + .any(|existing| existing.same_fact(&fact)) + { + flow.conditional_facts.push(fact); + } + } + } + + pub(crate) fn apply_conditional_flow_facts(&mut self, active: &NarrowOps) { + let flow = &mut self.current_mut().flow; + let matching_facts = flow + .conditional_facts + .iter() + .filter(|fact| fact.matches_active_conditions(active)) + .cloned() + .collect::>(); + for fact in matching_facts { + flow.info.insert(fact.name, fact.info); + } + } + /// Copy walrus-defined names from the current branch flow into the fork's /// base flow. Walrus names must be in the base flow so the merge can /// distinguish paths where the walrus was vs. wasn't evaluated. This works @@ -4093,6 +4202,24 @@ fn determine_definition_status( } } +fn common_conditional_facts(flows: &[&Flow]) -> Vec { + let Some((first, rest)) = flows.split_first() else { + return Vec::new(); + }; + first + .conditional_facts + .iter() + .filter(|fact| { + rest.iter().all(|flow| { + flow.conditional_facts + .iter() + .any(|other| fact.same_fact(other)) + }) + }) + .cloned() + .collect() +} + /// Information about a single branch being merged, including both the flow info /// for a specific name (if present) and the termination key from the flow this branch came from. struct MergeBranchEntry { @@ -4522,6 +4649,12 @@ impl<'a> BindingsBuilder<'a> { } else { live_branches }; + let fact_flows: Vec<&Flow> = if merge_style.is_loop() { + iter::once(&base).chain(flows.iter()).collect() + } else { + flows.iter().collect() + }; + let conditional_facts = common_conditional_facts(&fact_flows); // A plain `Loop` may skip its body entirely, so the pre-loop flow stays a possible // path past it. `LoopDefinitelyRuns` means the body runs at least once, so it does // not. @@ -4606,6 +4739,7 @@ impl<'a> BindingsBuilder<'a> { let flow = Flow { info: merged_flow_infos, excluded_none_pairs, + conditional_facts, has_terminated, terminated_by_raise: has_terminated && any_terminated_by_raise, is_definitely_unreachable: all_are_unreachable, diff --git a/pyrefly/lib/binding/stmt.rs b/pyrefly/lib/binding/stmt.rs index 3f29b59025..a90020af2a 100644 --- a/pyrefly/lib/binding/stmt.rs +++ b/pyrefly/lib/binding/stmt.rs @@ -1572,6 +1572,7 @@ impl<'a> BindingsBuilder<'a> { // The `BranchTest` of the latest clause bound so far, which decides whether // control reaches the next clause. let mut previous_test = None; + let mut conditional_flow_facts = Vec::new(); let mut contains_environment_test_with_no_else = false; let mut is_first_branch = true; let mut following_runtime_only_branch = false; @@ -1667,7 +1668,18 @@ impl<'a> BindingsBuilder<'a> { NarrowUseLocation::Span(range), &Usage::NonPinningValue(None), ); - negated_prev_ops.and_all(new_narrow_ops.negate()); + let mut active_narrow_ops = negated_prev_ops.clone(); + if active_narrow_ops.0.is_empty() { + active_narrow_ops = new_narrow_ops.clone(); + } else if !new_narrow_ops.0.is_empty() { + active_narrow_ops.and_all(new_narrow_ops.clone()); + } + let negated_new_narrow_ops = new_narrow_ops.negate(); + if negated_prev_ops.0.is_empty() { + negated_prev_ops = negated_new_narrow_ops; + } else { + negated_prev_ops.and_all(negated_new_narrow_ops); + } if is_type_checking_branch { self.type_checking_depth += 1; self.stmts(body, parent); @@ -1692,6 +1704,10 @@ impl<'a> BindingsBuilder<'a> { { contains_environment_test_with_no_else = true; } + conditional_flow_facts.extend( + self.scopes + .conditional_flow_facts_for_current_branch(&active_narrow_ops), + ); self.finish_branch(); if let Some(pair) = excluded_none_pair { self.scopes @@ -1766,6 +1782,8 @@ impl<'a> BindingsBuilder<'a> { }; self.finish_non_exhaustive_fork(&negated_prev_ops, fall_through_key); } + self.scopes + .add_conditional_flow_facts(conditional_flow_facts); // Preserve the configured-environment termination without treating it as // universally unreachable. This keeps later bindings out of a dead branch // while suppressing diagnostics that only apply to code live in this config. diff --git a/pyrefly/lib/test/flow_branching.rs b/pyrefly/lib/test/flow_branching.rs index e4363480ff..8923cdb379 100644 --- a/pyrefly/lib/test/flow_branching.rs +++ b/pyrefly/lib/test/flow_branching.rs @@ -249,6 +249,28 @@ print("not reported for parity") if False else None "#, ); +testcase!( + test_repeated_predicate_tracks_branch_assignment, + r#" +from typing import Literal, assert_type + +def f(b: str): + x: int | None = None + if b == "bar": + x = 3 + if b == "bar": + assert_type(x, Literal[3]) + +def invalidated_predicate(b: str): + x: int | None = None + if b == "bar": + x = 3 + b = "baz" + if b == "bar": + assert_type(x, int | None) +"#, +); + testcase!( test_listcomp_simple, r#" @@ -4040,7 +4062,6 @@ for x in range(10): // `if a:`, the variable is guaranteed to be initialized because the same // condition guards both the definition and the use. testcase!( - bug = "false positive: b is always initialized when a is truthy", test_guarded_initialization_basic, r#" def f(a: bool) -> int: @@ -4048,7 +4069,7 @@ def f(a: bool) -> int: b = 3 c = 5 if a: - return b # E: `b` may be uninitialized + return b return 9 "#, ); @@ -4078,7 +4099,6 @@ def f(a: bool, c: bool) -> int: ); testcase!( - bug = "false positive: b and c are always initialized when a is truthy", test_guarded_initialization_multiple_variables, r#" def f(a: bool) -> int: @@ -4086,13 +4106,12 @@ def f(a: bool) -> int: b = 3 c = 4 if a: - return b + c # E: `b` may be uninitialized # E: `c` may be uninitialized + return b + c return 0 "#, ); testcase!( - bug = "false positive: b is always initialized when a is truthy", test_guarded_initialization_with_intermediate_statements, r#" def f(a: bool) -> int: @@ -4101,13 +4120,12 @@ def f(a: bool) -> int: x = 5 y = x + 1 if a: - return b # E: `b` may be uninitialized + return b return 9 "#, ); testcase!( - bug = "false positive: b is always initialized when a is truthy", test_guarded_initialization_annotation_then_guarded_assign, r#" def f(a: bool) -> int: @@ -4115,20 +4133,19 @@ def f(a: bool) -> int: if a: b = 3 if a: - return b # E: `b` may be uninitialized + return b return 9 "#, ); testcase!( - bug = "false positive: b is always initialized when a is truthy", test_guarded_initialization_repeated_use, r#" def f(a: bool) -> None: if a: b = 3 if a: - print(b) # E: `b` may be uninitialized + print(b) if a: print(b) "#,