From 68a12f72a417655cb84dcd4284dec12f218eba5f Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 25 Jul 2023 22:53:29 -0400 Subject: [PATCH] Tweaks --- crates/ruff_python_ast/src/call_path.rs | 136 +++++++++++++++++++++++ crates/ruff_python_semantic/src/model.rs | 43 ++++--- 2 files changed, 155 insertions(+), 24 deletions(-) diff --git a/crates/ruff_python_ast/src/call_path.rs b/crates/ruff_python_ast/src/call_path.rs index c11c456fe7..b9f8bb912d 100644 --- a/crates/ruff_python_ast/src/call_path.rs +++ b/crates/ruff_python_ast/src/call_path.rs @@ -4,6 +4,142 @@ use smallvec::{smallvec, SmallVec}; /// A representation of a qualified name, like `typing.List`. pub type CallPath<'a> = SmallVec<[&'a str; 8]>; +/// Convert an `Expr` to its [`CallPath`] segments (like `["typing", "List"]`). +pub fn collect_head_path(expr: &Expr) -> Option<(&ast::ExprName, CallPath)> { + // Unroll the loop up to eight times, to match the maximum number of expected attributes. + // In practice, unrolling appears to give about a 4x speed-up on this hot path. + let attr1 = match expr { + Expr::Attribute(attr1) => attr1, + // Ex) `foo` + Expr::Name(name) => return Some((name, CallPath::new())), + _ => return None, + }; + + let attr2 = match attr1.value.as_ref() { + Expr::Attribute(attr2) => attr2, + // Ex) `foo.bar` + Expr::Name(name) => { + return Some((name, CallPath::from_slice(&[attr1.attr.as_str()]))); + } + _ => return None, + }; + + let attr3 = match attr2.value.as_ref() { + Expr::Attribute(attr3) => attr3, + // Ex) `foo.bar.baz` + Expr::Name(name) => { + return Some(( + name, + CallPath::from_slice(&[attr2.attr.as_str(), attr1.attr.as_str()]), + )); + } + _ => return None, + }; + + let attr4 = match attr3.value.as_ref() { + Expr::Attribute(attr4) => attr4, + // Ex) `foo.bar.baz.bop` + Expr::Name(name) => { + return Some(( + name, + CallPath::from_slice(&[ + attr3.attr.as_str(), + attr2.attr.as_str(), + attr1.attr.as_str(), + ]), + )); + } + _ => return None, + }; + + let attr5 = match attr4.value.as_ref() { + Expr::Attribute(attr5) => attr5, + // Ex) `foo.bar.baz.bop.bap` + Expr::Name(name) => { + return Some(( + name, + CallPath::from_slice(&[ + attr4.attr.as_str(), + attr3.attr.as_str(), + attr2.attr.as_str(), + attr1.attr.as_str(), + ]), + )); + } + _ => return None, + }; + + let attr6 = match attr5.value.as_ref() { + Expr::Attribute(attr6) => attr6, + // Ex) `foo.bar.baz.bop.bap.bab` + Expr::Name(name) => { + return Some(( + name, + CallPath::from_slice(&[ + attr5.attr.as_str(), + attr4.attr.as_str(), + attr3.attr.as_str(), + attr2.attr.as_str(), + attr1.attr.as_str(), + ]), + )); + } + _ => return None, + }; + + let attr7 = match attr6.value.as_ref() { + Expr::Attribute(attr7) => attr7, + // Ex) `foo.bar.baz.bop.bap.bab.bob` + Expr::Name(name) => { + return Some(( + name, + CallPath::from_slice(&[ + attr6.attr.as_str(), + attr5.attr.as_str(), + attr4.attr.as_str(), + attr3.attr.as_str(), + attr2.attr.as_str(), + attr1.attr.as_str(), + ]), + )); + } + _ => return None, + }; + + let attr8 = match attr7.value.as_ref() { + Expr::Attribute(attr8) => attr8, + // Ex) `foo.bar.baz.bop.bap.bab.bob.bib` + Expr::Name(name) => { + return Some(( + name, + CallPath::from_slice(&[ + attr7.attr.as_str(), + attr6.attr.as_str(), + attr5.attr.as_str(), + attr4.attr.as_str(), + attr3.attr.as_str(), + attr2.attr.as_str(), + attr1.attr.as_str(), + ]), + )); + } + _ => return None, + }; + + let (name, mut call_path) = collect_head_path(&attr8.value)?; + call_path.extend([ + attr8.attr.as_str(), + attr7.attr.as_str(), + attr6.attr.as_str(), + attr5.attr.as_str(), + attr4.attr.as_str(), + attr3.attr.as_str(), + attr2.attr.as_str(), + attr1.attr.as_str(), + ]); + Some((name, call_path)) +} + /// Convert an `Expr` to its [`CallPath`] segments (like `["typing", "List"]`). pub fn collect_call_path(expr: &Expr) -> Option { // Unroll the loop up to eight times, to match the maximum number of expected attributes. diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs index 13ac52113a..1f70f87721 100644 --- a/crates/ruff_python_semantic/src/model.rs +++ b/crates/ruff_python_semantic/src/model.rs @@ -2,9 +2,11 @@ use std::path::Path; use bitflags::bitflags; use rustc_hash::FxHashMap; -use smallvec::{smallvec, SmallVec}; +use smallvec::smallvec; -use ruff_python_ast::call_path::{collect_call_path, from_unqualified_name, CallPath}; +use ruff_python_ast::call_path::{ + collect_call_path, collect_head_path, from_unqualified_name, CallPath, +}; use ruff_python_ast::{self as ast, Expr, Ranged, Stmt}; use ruff_python_stdlib::path::is_python_stub_file; use ruff_python_stdlib::typing::is_typing_extension; @@ -623,40 +625,33 @@ impl<'a> SemanticModel<'a> { } } + let (head, tail) = collect_head_path(value)?; + // If the name was already resolved, look it up; otherwise, search for the symbol. - let head = match_head(value)?; - let binding = self - .resolved_names - .get(&head.into()) - .map(|id| self.binding(*id)) - .or_else(|| self.find_binding(&head.id))?; + let binding = if let Some(id) = self.resolved_names.get(&head.into()) { + self.binding(*id) + } else { + self.find_binding(&head.id)? + }; match &binding.kind { BindingKind::Import(Import { call_path, .. }) => { - let x = collect_call_path(value)?; - let (_, tail) = x.split_first()?; - - let mut resolved = SmallVec::with_capacity(call_path.len() + tail.len()); - resolved.extend_from_slice(call_path); - resolved.extend_from_slice(tail); + let resolved: CallPath = call_path.iter().chain(tail.iter()).copied().collect(); + // resolved.extend_from_slice(call_path); + // resolved.extend_from_slice(tail); Some(resolved) } BindingKind::SubmoduleImport(SubmoduleImport { qualified_name, .. }) => { - let x = collect_call_path(value)?; - let (_, tail) = x.split_first()?; - let name = qualified_name.split('.').next().unwrap_or(qualified_name); let mut source_path: CallPath = from_unqualified_name(name); - source_path.extend_from_slice(tail); + source_path.extend_from_slice(tail.as_slice()); Some(source_path) } BindingKind::FromImport(FromImport { call_path, .. }) => { - let x = collect_call_path(value)?; - let (_, tail) = x.split_first()?; - - let mut resolved = SmallVec::with_capacity(call_path.len() + tail.len()); - resolved.extend_from_slice(call_path); - resolved.extend_from_slice(tail); + let resolved: CallPath = call_path.iter().chain(tail.iter()).copied().collect(); + // let mut resolved = SmallVec::with_capacity(call_path.len() + tail.len()); + // resolved.extend_from_slice(call_path); + // resolved.extend_from_slice(tail); Some(resolved) } BindingKind::Builtin => Some(smallvec!["", head.id.as_str()]),