1use crate::expression::Expression;
2use crate::member::{
3 Member, MemberExpr, MemberExprBuilder, MemberExprRef, MemberTable, MemberTableBuilder,
4 ScopedMemberId,
5};
6use crate::predicate::PatternPredicate;
7use crate::symbol::{ScopedSymbolId, Symbol, SymbolTable, SymbolTableBuilder};
8use crate::{Db, PossiblyNarrowedPlaces};
9use ruff_db::parsed::ParsedModuleRef;
10use ruff_index::IndexVec;
11use ruff_python_ast as ast;
12use smallvec::SmallVec;
13use std::hash::Hash;
14use std::iter::FusedIterator;
15
16pub(crate) fn match_subject_place_expressions(subject: &ast::Expr) -> SmallVec<[&ast::Expr; 2]> {
27 let mut expressions: SmallVec<[&ast::Expr; 2]> = SmallVec::new();
28 expressions.push(subject);
29 match subject {
30 ast::Expr::Subscript(subscript) => expressions.push(&subscript.value),
31 ast::Expr::Attribute(attribute) => expressions.push(&attribute.value),
32 _ => {}
33 }
34 expressions
35}
36
37#[derive(Eq, PartialEq, Debug, get_size2::GetSize)]
39pub enum PlaceExpr {
40 Symbol(Symbol),
42
43 Member(Member),
45}
46
47impl PlaceExpr {
48 pub fn from_expr_name(name: &ast::ExprName) -> Self {
52 PlaceExpr::Symbol(Symbol::new(name.id.clone()))
53 }
54
55 pub fn try_from_expr<'e>(expr: impl Into<ast::ExprRef<'e>>) -> Option<Self> {
64 let expr = expr.into();
65
66 let expr = match expr {
69 ast::ExprRef::Named(named) if named.target.is_name_expr() => {
70 named.target.as_ref().into()
71 }
72 ast::ExprRef::Named(_) => return None,
73 _ => expr,
74 };
75
76 if let ast::ExprRef::Name(name) = expr {
77 return Some(PlaceExpr::Symbol(Symbol::new(name.id.clone())));
78 }
79
80 MemberExprBuilder::visit_expr(expr).and_then(Self::try_from_member_expr)
81 }
82
83 pub(super) fn try_from_member_expr(builder: MemberExprBuilder) -> Option<Self> {
87 let member_expression = MemberExpr::try_from_builder(builder)?;
88 Some(Self::Member(Member::new(member_expression)))
89 }
90}
91
92impl std::fmt::Display for PlaceExpr {
93 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94 match self {
95 Self::Symbol(symbol) => std::fmt::Display::fmt(symbol, f),
96 Self::Member(member) => std::fmt::Display::fmt(member, f),
97 }
98 }
99}
100
101#[derive(Eq, PartialEq, Debug, Copy, Clone)]
105pub enum PlaceExprRef<'a> {
106 Symbol(&'a Symbol),
107 Member(&'a Member),
108}
109
110impl<'a> PlaceExprRef<'a> {
111 pub const fn as_symbol(self) -> Option<&'a Symbol> {
113 if let PlaceExprRef::Symbol(symbol) = self {
114 Some(symbol)
115 } else {
116 None
117 }
118 }
119
120 pub const fn is_symbol(self) -> bool {
122 matches!(self, PlaceExprRef::Symbol(_))
123 }
124
125 pub fn is_declared(self) -> bool {
126 match self {
127 Self::Symbol(symbol) => symbol.is_declared(),
128 Self::Member(member) => member.is_declared(),
129 }
130 }
131
132 pub const fn is_bound(self) -> bool {
133 match self {
134 PlaceExprRef::Symbol(symbol) => symbol.is_bound(),
135 PlaceExprRef::Member(member) => member.is_bound(),
136 }
137 }
138
139 pub fn num_member_segments(self) -> usize {
140 match self {
141 PlaceExprRef::Symbol(_) => 0,
142 PlaceExprRef::Member(member) => member.expression().num_segments(),
143 }
144 }
145}
146
147impl<'a> From<&'a Symbol> for PlaceExprRef<'a> {
148 fn from(value: &'a Symbol) -> Self {
149 Self::Symbol(value)
150 }
151}
152
153impl<'a> From<&'a Member> for PlaceExprRef<'a> {
154 fn from(value: &'a Member) -> Self {
155 Self::Member(value)
156 }
157}
158
159impl<'a> From<&'a PlaceExpr> for PlaceExprRef<'a> {
160 fn from(value: &'a PlaceExpr) -> Self {
161 match value {
162 PlaceExpr::Symbol(symbol) => PlaceExprRef::Symbol(symbol),
163 PlaceExpr::Member(member) => PlaceExprRef::Member(member),
164 }
165 }
166}
167
168impl std::fmt::Display for PlaceExprRef<'_> {
169 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
170 match self {
171 Self::Symbol(symbol) => std::fmt::Display::fmt(symbol, f),
172 Self::Member(member) => std::fmt::Display::fmt(member, f),
173 }
174 }
175}
176
177#[derive(
179 Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, get_size2::GetSize, salsa::SalsaValue,
180)]
181pub enum ScopedPlaceId {
182 Symbol(ScopedSymbolId),
183 Member(ScopedMemberId),
184}
185
186#[derive(Debug, Eq, PartialEq, get_size2::GetSize)]
187pub struct PlaceTable {
188 symbols: SymbolTable,
189 members: MemberTable,
190}
191
192impl PlaceTable {
193 pub fn parents<'a>(&'a self, place_expr: impl Into<PlaceExprRef<'a>>) -> ParentPlaceIter<'a> {
197 match place_expr.into() {
198 PlaceExprRef::Symbol(_) => ParentPlaceIter::for_symbol(),
199 PlaceExprRef::Member(member) => {
200 ParentPlaceIter::for_member(member.expression(), &self.symbols, &self.members)
201 }
202 }
203 }
204
205 pub fn symbols(&self) -> std::slice::Iter<'_, Symbol> {
207 self.symbols.iter()
208 }
209
210 pub fn members(&self) -> std::slice::Iter<'_, Member> {
212 self.members.iter()
213 }
214
215 #[track_caller]
220 pub fn symbol(&self, id: ScopedSymbolId) -> &Symbol {
221 self.symbols.symbol(id)
222 }
223
224 pub fn symbol_by_name(&self, name: &str) -> Option<&Symbol> {
228 self.symbols.symbol_id(name).map(|id| self.symbol(id))
229 }
230
231 #[track_caller]
236 pub fn member(&self, id: ScopedMemberId) -> &Member {
237 self.members.member(id)
238 }
239
240 pub fn symbol_id(&self, name: &str) -> Option<ScopedSymbolId> {
242 self.symbols.symbol_id(name)
243 }
244
245 pub fn place_id<'e>(&self, place_expr: impl Into<PlaceExprRef<'e>>) -> Option<ScopedPlaceId> {
247 let place_expr = place_expr.into();
248
249 match place_expr {
250 PlaceExprRef::Symbol(symbol) => self.symbols.symbol_id(symbol.name()).map(Into::into),
251 PlaceExprRef::Member(member) => {
252 self.members.member_id(member.expression()).map(Into::into)
253 }
254 }
255 }
256
257 #[track_caller]
262 pub fn place(&self, place_id: impl Into<ScopedPlaceId>) -> PlaceExprRef<'_> {
263 match place_id.into() {
264 ScopedPlaceId::Symbol(symbol) => self.symbol(symbol).into(),
265 ScopedPlaceId::Member(member) => self.member(member).into(),
266 }
267 }
268
269 pub fn member_id_by_instance_attribute_name(&self, name: &str) -> Option<ScopedMemberId> {
270 self.members.place_id_by_instance_attribute_name(name)
271 }
272}
273
274#[derive(Default)]
275pub struct PlaceTableBuilder {
276 symbols: SymbolTableBuilder,
277 member: MemberTableBuilder,
278
279 associated_symbol_members: IndexVec<ScopedSymbolId, SmallVec<[ScopedMemberId; 4]>>,
280 associated_sub_members: IndexVec<ScopedMemberId, SmallVec<[ScopedMemberId; 4]>>,
281}
282
283impl PlaceTableBuilder {
284 pub(crate) fn place_id(&self, expression: PlaceExprRef) -> Option<ScopedPlaceId> {
286 match expression {
287 PlaceExprRef::Symbol(symbol) => self.symbols.symbol_id(symbol.name()).map(Into::into),
288 PlaceExprRef::Member(member) => {
289 self.member.member_id(member.expression()).map(Into::into)
290 }
291 }
292 }
293
294 #[track_caller]
295 pub(super) fn symbol(&self, id: ScopedSymbolId) -> &Symbol {
296 self.symbols.symbol(id)
297 }
298
299 pub(super) fn member(&self, id: ScopedMemberId) -> &Member {
300 self.member.member(id)
301 }
302
303 pub(super) fn symbol_id(&self, name: &str) -> Option<ScopedSymbolId> {
304 self.symbols.symbol_id(name)
305 }
306
307 #[track_caller]
308 pub(super) fn symbol_mut(&mut self, id: ScopedSymbolId) -> &mut Symbol {
309 self.symbols.symbol_mut(id)
310 }
311
312 #[track_caller]
313 fn member_mut(&mut self, id: ScopedMemberId) -> &mut Member {
314 self.member.member_mut(id)
315 }
316
317 #[track_caller]
318 pub(crate) fn place(&self, place_id: impl Into<ScopedPlaceId>) -> PlaceExprRef<'_> {
319 match place_id.into() {
320 ScopedPlaceId::Symbol(id) => PlaceExprRef::Symbol(self.symbols.symbol(id)),
321 ScopedPlaceId::Member(id) => PlaceExprRef::Member(self.member.member(id)),
322 }
323 }
324
325 pub(crate) fn associated_place_ids(&self, place: ScopedPlaceId) -> &[ScopedMemberId] {
326 match place {
327 ScopedPlaceId::Symbol(symbol) => &self.associated_symbol_members[symbol],
328 ScopedPlaceId::Member(member) => &self.associated_sub_members[member],
329 }
330 }
331
332 pub(crate) fn iter(&self) -> impl Iterator<Item = PlaceExprRef<'_>> {
333 self.symbols
334 .iter()
335 .map(Into::into)
336 .chain(self.member.iter().map(PlaceExprRef::Member))
337 }
338
339 pub(crate) fn symbols(&self) -> impl Iterator<Item = &Symbol> {
340 self.symbols.iter()
341 }
342
343 pub(crate) fn add_symbol(&mut self, symbol: Symbol) -> (ScopedSymbolId, bool) {
344 let (id, is_new) = self.symbols.add(symbol);
345
346 if is_new {
347 let new_id = self.associated_symbol_members.push(SmallVec::new_const());
348 debug_assert_eq!(new_id, id);
349 }
350
351 (id, is_new)
352 }
353
354 fn add_member(&mut self, member: Member) -> (ScopedMemberId, bool) {
355 let (id, is_new) = self.member.add(member);
356
357 if is_new {
358 let new_id = self.associated_sub_members.push(SmallVec::new_const());
359 debug_assert_eq!(new_id, id);
360
361 let member = self.member.member(id);
362
363 for parent_id in
365 ParentPlaceIter::for_member(member.expression(), &self.symbols, &self.member)
366 {
367 match parent_id {
368 ScopedPlaceId::Symbol(scoped_symbol_id) => {
369 self.associated_symbol_members[scoped_symbol_id].push(id);
370 }
371 ScopedPlaceId::Member(scoped_member_id) => {
372 self.associated_sub_members[scoped_member_id].push(id);
373 }
374 }
375 }
376 }
377
378 (id, is_new)
379 }
380
381 pub(crate) fn add_place(&mut self, place: PlaceExpr) -> (ScopedPlaceId, bool) {
382 match place {
383 PlaceExpr::Symbol(symbol) => {
384 let (id, is_new) = self.add_symbol(symbol);
385 (ScopedPlaceId::Symbol(id), is_new)
386 }
387 PlaceExpr::Member(member) => {
388 let (id, is_new) = self.add_member(member);
389 (ScopedPlaceId::Member(id), is_new)
390 }
391 }
392 }
393
394 #[track_caller]
395 pub(super) fn mark_bound(&mut self, id: ScopedPlaceId) {
396 match id {
397 ScopedPlaceId::Symbol(symbol_id) => {
398 self.symbol_mut(symbol_id).mark_bound();
399 }
400 ScopedPlaceId::Member(member_id) => {
401 self.member_mut(member_id).mark_bound();
402 }
403 }
404 }
405
406 #[track_caller]
407 pub(super) fn mark_declared(&mut self, id: ScopedPlaceId) {
408 match id {
409 ScopedPlaceId::Symbol(symbol_id) => {
410 self.symbol_mut(symbol_id).mark_declared();
411 }
412 ScopedPlaceId::Member(member_id) => {
413 self.member_mut(member_id).mark_declared();
414 }
415 }
416 }
417
418 pub(crate) fn finish(self) -> PlaceTable {
419 PlaceTable {
420 symbols: self.symbols.build(),
421 members: self.member.build(),
422 }
423 }
424}
425
426impl ScopedPlaceId {
427 pub const fn is_symbol(self) -> bool {
428 matches!(self, ScopedPlaceId::Symbol(_))
429 }
430
431 pub const fn is_member(self) -> bool {
432 matches!(self, ScopedPlaceId::Member(_))
433 }
434
435 pub const fn as_symbol(self) -> Option<ScopedSymbolId> {
436 if let ScopedPlaceId::Symbol(id) = self {
437 Some(id)
438 } else {
439 None
440 }
441 }
442
443 pub const fn expect_symbol(self) -> ScopedSymbolId {
444 match self {
445 ScopedPlaceId::Symbol(symbol) => symbol,
446 ScopedPlaceId::Member(_) => {
447 panic!("Expected ScopedPlaceId::Symbol, found ScopedPlaceId::Member")
448 }
449 }
450 }
451}
452
453impl<T> std::ops::Index<ScopedPlaceId> for Vec<T> {
454 type Output = T;
455
456 fn index(&self, index: ScopedPlaceId) -> &Self::Output {
457 match index {
458 ScopedPlaceId::Symbol(id) => &self[id.index()],
459 ScopedPlaceId::Member(id) => &self[id.index()],
460 }
461 }
462}
463
464impl From<ScopedMemberId> for ScopedPlaceId {
465 fn from(value: ScopedMemberId) -> Self {
466 Self::Member(value)
467 }
468}
469
470impl From<ScopedSymbolId> for ScopedPlaceId {
471 fn from(value: ScopedSymbolId) -> Self {
472 Self::Symbol(value)
473 }
474}
475
476pub struct ParentPlaceIter<'a> {
477 state: Option<ParentPlaceIterState<'a>>,
478}
479
480enum ParentPlaceIterState<'a> {
481 Symbol {
482 symbol_name: &'a str,
483 symbols: &'a SymbolTable,
484 },
485 Member {
486 symbols: &'a SymbolTable,
487 members: &'a MemberTable,
488 next_member: MemberExprRef<'a>,
489 },
490}
491
492impl<'a> ParentPlaceIterState<'a> {
493 fn parent_state(
494 expression: &MemberExprRef<'a>,
495 symbols: &'a SymbolTable,
496 members: &'a MemberTable,
497 ) -> Self {
498 match expression.parent() {
499 Some(parent) => Self::Member {
500 next_member: parent,
501 symbols,
502 members,
503 },
504 None => Self::Symbol {
505 symbol_name: expression.symbol_name(),
506 symbols,
507 },
508 }
509 }
510}
511
512impl<'a> ParentPlaceIter<'a> {
513 fn for_symbol() -> Self {
514 ParentPlaceIter { state: None }
515 }
516
517 fn for_member(
518 expression: &'a MemberExpr,
519 symbol_table: &'a SymbolTable,
520 member_table: &'a MemberTable,
521 ) -> Self {
522 let expr_ref = expression.as_ref();
523 ParentPlaceIter {
524 state: Some(ParentPlaceIterState::parent_state(
525 &expr_ref,
526 symbol_table,
527 member_table,
528 )),
529 }
530 }
531}
532
533impl Iterator for ParentPlaceIter<'_> {
534 type Item = ScopedPlaceId;
535
536 fn next(&mut self) -> Option<Self::Item> {
537 loop {
538 match self.state.take()? {
539 ParentPlaceIterState::Symbol {
540 symbol_name,
541 symbols,
542 } => {
543 let id = symbols.symbol_id(symbol_name)?;
544 break Some(id.into());
545 }
546 ParentPlaceIterState::Member {
547 symbols,
548 members,
549 next_member,
550 } => {
551 self.state = Some(ParentPlaceIterState::parent_state(
552 &next_member,
553 symbols,
554 members,
555 ));
556
557 if let Some(id) = members.member_id(next_member) {
558 break Some(id.into());
559 }
560 }
561 }
562 }
563 }
564}
565
566impl FusedIterator for ParentPlaceIter<'_> {}
567
568pub(crate) struct PossiblyNarrowedPlacesBuilder<'db, 'a> {
573 db: &'db dyn Db,
574 places: &'a PlaceTableBuilder,
575}
576
577impl<'db, 'a> PossiblyNarrowedPlacesBuilder<'db, 'a> {
578 pub(crate) fn new(db: &'db dyn Db, places: &'a PlaceTableBuilder) -> Self {
579 Self { db, places }
580 }
581
582 pub(crate) fn expression(self, expr: &ast::Expr) -> PossiblyNarrowedPlaces {
584 self.expression_node(expr)
585 }
586
587 pub(crate) fn pattern(
589 self,
590 pattern: PatternPredicate<'db>,
591 module: &ParsedModuleRef,
592 ) -> PossiblyNarrowedPlaces {
593 self.pattern_kind(pattern.subject(self.db), module)
594 }
595
596 fn expression_node(&self, expr: &ast::Expr) -> PossiblyNarrowedPlaces {
597 match expr {
598 ast::Expr::Name(_) => self.simple_expr(expr),
600 ast::Expr::Attribute(attribute) => {
602 let mut places = self.simple_expr(expr);
603 places.extend(self.simple_expr(&attribute.value));
604 places
605 }
606 ast::Expr::Subscript(subscript) => {
608 let mut places = self.simple_expr(expr);
609 places.extend(self.simple_expr(&subscript.value));
610 places
611 }
612 ast::Expr::Compare(expr_compare) => self.expr_compare(expr_compare),
614 ast::Expr::Call(expr_call) => self.expr_call(expr_call),
616 ast::Expr::UnaryOp(unary_op) if unary_op.op == ast::UnaryOp::Not => {
618 self.expression_node(&unary_op.operand)
619 }
620 ast::Expr::BoolOp(bool_op) => self.expr_bool_op(bool_op),
622 ast::Expr::If(expr_if) => self.expr_if(expr_if),
624 ast::Expr::Named(expr_named) => {
626 let mut places = self.simple_expr(&expr_named.target);
627 places.extend(self.expression_node(&expr_named.value));
628 places
629 }
630 _ => PossiblyNarrowedPlaces::default(),
631 }
632 }
633
634 fn simple_expr(&self, expr: &ast::Expr) -> PossiblyNarrowedPlaces {
636 let mut places = PossiblyNarrowedPlaces::default();
637 if let Some(place_expr) = PlaceExpr::try_from_expr(expr) {
638 if let Some(place) = self.places.place_id((&place_expr).into()) {
639 places.insert(place);
640 }
641 }
642 places
643 }
644
645 fn expr_compare(&self, expr_compare: &ast::ExprCompare) -> PossiblyNarrowedPlaces {
648 let mut places = PossiblyNarrowedPlaces::default();
649
650 self.add_narrowing_target(&expr_compare.left, &mut places);
652
653 for comparator in &expr_compare.comparators {
655 self.add_narrowing_target(comparator, &mut places);
656 }
657
658 let can_narrow_tagged_union_base = matches!(
659 &*expr_compare.ops,
660 [ast::CmpOp::Eq | ast::CmpOp::NotEq | ast::CmpOp::Is | ast::CmpOp::IsNot]
661 );
662
663 for expr in std::iter::once(&*expr_compare.left).chain(&expr_compare.comparators) {
665 if can_narrow_tagged_union_base
666 && let ast::Expr::Subscript(subscript) = expr.expression_value()
667 && let Some(place_expr) = PlaceExpr::try_from_expr(&subscript.value)
668 && let Some(place) = self.places.place_id((&place_expr).into())
669 {
670 places.insert(place);
671 }
672 if can_narrow_tagged_union_base
673 && let ast::Expr::Attribute(attribute) = expr
674 && let Some(place_expr) = PlaceExpr::try_from_expr(&attribute.value)
675 && let Some(place) = self.places.place_id((&place_expr).into())
676 {
677 places.insert(place);
678 }
679 }
680
681 places
682 }
683
684 fn expr_call(&self, expr_call: &ast::ExprCall) -> PossiblyNarrowedPlaces {
687 let mut places = PossiblyNarrowedPlaces::default();
688
689 for argument in expr_call.arguments.args.iter().take(2).chain(
696 expr_call
697 .arguments
698 .keywords
699 .iter()
700 .filter(|keyword| keyword.arg.is_some())
701 .map(|keyword| &keyword.value),
702 ) {
703 if let Some(place_expr) = PlaceExpr::try_from_expr(argument) {
704 if let Some(place) = self.places.place_id((&place_expr).into()) {
705 places.insert(place);
706 }
707 }
708 }
709
710 if let Some(first_arg) = expr_call.arguments.args.first() {
712 if expr_call.arguments.args.len() == 1 && expr_call.arguments.keywords.is_empty() {
713 places.extend(self.expression_node(first_arg));
714 }
715 }
716
717 places
718 }
719
720 fn expr_bool_op(&self, bool_op: &ast::ExprBoolOp) -> PossiblyNarrowedPlaces {
722 let mut places = PossiblyNarrowedPlaces::default();
723 for value in &bool_op.values {
724 places.extend(self.expression_node(value));
725 }
726 places
727 }
728
729 fn expr_if(&self, expr_if: &ast::ExprIf) -> PossiblyNarrowedPlaces {
730 let mut places = self.expression_node(&expr_if.test);
731 places.extend(self.expression_node(&expr_if.body));
732 places.extend(self.expression_node(&expr_if.orelse));
733 places
734 }
735
736 fn add_narrowing_target(&self, expr: &ast::Expr, places: &mut PossiblyNarrowedPlaces) {
738 if let Some(place_expr) = PlaceExpr::try_from_expr(expr)
739 && let Some(place) = self.places.place_id((&place_expr).into())
740 {
741 places.insert(place);
742 }
743
744 match expr.expression_value() {
745 ast::Expr::Call(call) if call.arguments.args.len() == 1 => {
747 if let Some(first_arg) = call.arguments.args.first()
748 && let Some(place_expr) = PlaceExpr::try_from_expr(first_arg)
749 && let Some(place) = self.places.place_id((&place_expr).into())
750 {
751 places.insert(place);
752 }
753 }
754 ast::Expr::Attribute(attribute) if attribute.attr.as_str() == "__class__" => {
756 if let Some(place_expr) = PlaceExpr::try_from_expr(&attribute.value)
757 && let Some(place) = self.places.place_id((&place_expr).into())
758 {
759 places.insert(place);
760 }
761 }
762 _ => {}
763 }
764 }
765
766 fn pattern_kind(
768 &self,
769 subject: Expression<'db>,
770 module: &ParsedModuleRef,
771 ) -> PossiblyNarrowedPlaces {
772 let mut places = PossiblyNarrowedPlaces::default();
773
774 let subject_node = subject.node_ref(self.db).node(module);
775 for expression in match_subject_place_expressions(subject_node) {
776 if let Some(place) = PlaceExpr::try_from_expr(expression)
777 .and_then(|place| self.places.place_id((&place).into()))
778 {
779 places.insert(place);
780 }
781 }
782 places
783 }
784}