Improve performance of import matching code (#744)

src/ast/helpers.rs

@@ -3,13 +3,14 @@ use once_cell::sync::Lazy;
 use regex::Regex;
 use rustpython_ast::{Excepthandler, ExcepthandlerKind, Expr, ExprKind, Location, StmtKind};
 
-fn compose_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) {
+#[inline(always)]
+fn collect_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) {
     match &expr.node {
         ExprKind::Call { func, .. } => {
-            compose_call_path_inner(func, parts);
+            collect_call_path_inner(func, parts);
         }
         ExprKind::Attribute { value, attr, .. } => {
-            compose_call_path_inner(value, parts);
+            collect_call_path_inner(value, parts);
             parts.push(attr);
         }
         ExprKind::Name { id, .. } => {
@@ -20,9 +21,9 @@ fn compose_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) {
 }
 
 /// Convert an `Expr` to its call path (like `List`, or `typing.List`).
+#[inline(always)]
 pub fn compose_call_path(expr: &Expr) -> Option<String> {
-    let mut segments = vec![];
-    compose_call_path_inner(expr, &mut segments);
+    let segments = collect_call_paths(expr);
     if segments.is_empty() {
         None
     } else {
@@ -30,6 +31,14 @@ pub fn compose_call_path(expr: &Expr) -> Option<String> {
     }
 }
 
+/// Convert an `Expr` to its call path segments (like ["typing", "List"]).
+#[inline(always)]
+pub fn collect_call_paths(expr: &Expr) -> Vec<&str> {
+    let mut segments = vec![];
+    collect_call_path_inner(expr, &mut segments);
+    segments
+}
+
 /// Return `true` if the `Expr` is a name or attribute reference to `${target}`.
 pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
     match &expr.node {
@@ -72,7 +81,7 @@ pub fn is_assignment_to_a_dunder(node: &StmtKind) -> bool {
 }
 
 /// Extract the names of all handled exceptions.
-pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec<String> {
+pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec<Vec<&str>> {
     let mut handler_names = vec![];
     for handler in handlers {
         match &handler.node {
@@ -80,12 +89,16 @@ pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec<String> {
                 if let Some(type_) = type_ {
                     if let ExprKind::Tuple { elts, .. } = &type_.node {
                         for type_ in elts {
-                            if let Some(name) = compose_call_path(type_) {
-                                handler_names.push(name);
+                            let call_path = collect_call_paths(type_);
+                            if !call_path.is_empty() {
+                                handler_names.push(call_path);
                             }
                         }
-                    } else if let Some(name) = compose_call_path(type_) {
-                        handler_names.push(name);
+                    } else {
+                        let call_path = collect_call_paths(type_);
+                        if !call_path.is_empty() {
+                            handler_names.push(call_path);
+                        }
                     }
                 }
             }
@@ -123,66 +136,88 @@ pub fn to_absolute(relative: &Location, base: &Location) -> Location {
 /// `typing.Union`.
 pub fn match_module_member(
     expr: &Expr,
-    target: &str,
+    module: &str,
+    member: &str,
     from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
 ) -> bool {
-    compose_call_path(expr)
-        .map(|expr| match_call_path(&expr, target, from_imports))
-        .unwrap_or(false)
+    match_call_path(&collect_call_paths(expr), module, member, from_imports)
 }
 
 /// Return `true` if the `call_path` is a reference to `${module}.${target}`.
 ///
 /// Optimized version of `match_module_member` for pre-computed call paths.
 pub fn match_call_path(
-    call_path: &str,
-    target: &str,
+    call_path: &[&str],
+    module: &str,
+    member: &str,
     from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
 ) -> bool {
-    // Case (1a): it's the same call path (`import typing`, `typing.re.Match`).
-    // Case (1b): it's the same call path (`import typing.re`, `typing.re.Match`).
-    if call_path == target {
+    // If we have no segments, we can't ever match.
+    let num_segments = call_path.len();
+    if num_segments == 0 {
+        return false;
+    }
+
+    // If the last segment doesn't match the member, we can't ever match.
+    if call_path[num_segments - 1] != member {
+        return false;
+    }
+
+    // We now only need the module path, so throw out the member name.
+    let call_path = &call_path[..num_segments - 1];
+    let num_segments = call_path.len();
+
+    // Case (1a): We imported star from the parent (`from typing.re import *`,
+    // `Match`).
+    // Case (1b): We imported from the parent (`from typing.re import Match`,
+    // `Match`).
+    // Case (2): It's a builtin (like `list`).
+    if num_segments == 0
+        && (module.is_empty()
+            || from_imports
+                .get(module)
+                .map(|imports| imports.contains(member) || imports.contains("*"))
+                .unwrap_or(false))
+    {
         return true;
     }
 
-    if let Some((parent, member)) = target.rsplit_once('.') {
-        // Case (2): We imported star from the parent (`from typing.re import *`,
-        // `Match`).
-        if call_path == member
-            && from_imports
-                .get(parent)
-                .map(|imports| imports.contains("*"))
-                .unwrap_or(false)
-        {
-            return true;
-        }
-
-        // Case (3): We imported from the parent (`from typing.re import Match`,
-        // `Match`)
-        if call_path == member
-            && from_imports
-                .get(parent)
-                .map(|imports| imports.contains(member))
-                .unwrap_or(false)
-        {
-            return true;
-        }
+    // Case (3a): it's a fully qualified call path (`import typing`,
+    // `typing.re.Match`). Case (3b): it's a fully qualified call path (`import
+    // typing.re`, `typing.re.Match`).
+    if num_segments > 0
+        && module
+            .split('.')
+            .enumerate()
+            .all(|(index, segment)| index < num_segments && call_path[index] == segment)
+    {
+        return true;
     }
 
     // Case (4): We imported from the grandparent (`from typing import re`,
     // `re.Match`)
-    let mut parts = target.rsplitn(3, '.');
-    let member = parts.next();
-    let parent = parts.next();
-    let grandparent = parts.next();
-    if let (Some(member), Some(parent), Some(grandparent)) = (member, parent, grandparent) {
-        if call_path == format!("{parent}.{member}")
-            && from_imports
-                .get(grandparent)
-                .map(|imports| imports.contains(parent))
+    if num_segments > 0 {
+        // Find the number of common segments.
+        // TODO(charlie): Rewrite to avoid this allocation.
+        let parts: Vec<&str> = module.split('.').collect();
+        let num_matches = (0..parts.len())
+            .take(num_segments)
+            .take_while(|i| parts[parts.len() - 1 - i] == call_path[num_segments - 1 - i])
+            .count();
+        if num_matches > 0 {
+            // Verify that we have an import of the final common segment, from the remaining
+            // parent.
+            let cut = parts.len() - num_matches;
+            // TODO(charlie): Rewrite to avoid this allocation.
+            let module = parts[..cut].join(".");
+            let member = parts[cut];
+            if from_imports
+                .get(&module.as_str())
+                .map(|imports| imports.contains(member))
                 .unwrap_or(false)
-        {
-            return true;
+            {
+                return true;
+            }
         }
     }
 
@@ -197,12 +232,25 @@ mod tests {
 
     use crate::ast::helpers::match_module_member;
 
+    #[test]
+    fn builtin() -> Result<()> {
+        let expr = parser::parse_expression("list", "<filename>")?;
+        assert!(match_module_member(
+            &expr,
+            "",
+            "list",
+            &FnvHashMap::default()
+        ));
+        Ok(())
+    }
+
     #[test]
     fn fully_qualified() -> Result<()> {
         let expr = parser::parse_expression("typing.re.Match", "<filename>")?;
         assert!(match_module_member(
             &expr,
-            "typing.re.Match",
+            "typing.re",
+            "Match",
             &FnvHashMap::default()
         ));
         Ok(())
@@ -213,13 +261,15 @@ mod tests {
         let expr = parser::parse_expression("Match", "<filename>")?;
         assert!(!match_module_member(
             &expr,
-            "typing.re.Match",
+            "typing.re",
+            "Match",
             &FnvHashMap::default(),
         ));
         let expr = parser::parse_expression("re.Match", "<filename>")?;
         assert!(!match_module_member(
             &expr,
-            "typing.re.Match",
+            "typing.re",
+            "Match",
             &FnvHashMap::default(),
         ));
         Ok(())
@@ -230,7 +280,8 @@ mod tests {
         let expr = parser::parse_expression("Match", "<filename>")?;
         assert!(match_module_member(
             &expr,
-            "typing.re.Match",
+            "typing.re",
+            "Match",
             &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["*"]))])
         ));
         Ok(())
@@ -241,7 +292,8 @@ mod tests {
         let expr = parser::parse_expression("Match", "<filename>")?;
         assert!(match_module_member(
             &expr,
-            "typing.re.Match",
+            "typing.re",
+            "Match",
             &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["Match"]))])
         ));
         Ok(())
@@ -252,7 +304,24 @@ mod tests {
         let expr = parser::parse_expression("re.Match", "<filename>")?;
         assert!(match_module_member(
             &expr,
-            "typing.re.Match",
+            "typing.re",
+            "Match",
+            &FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))])
+        ));
+
+        let expr = parser::parse_expression("match.Match", "<filename>")?;
+        assert!(match_module_member(
+            &expr,
+            "typing.re.match",
+            "Match",
+            &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["match"]))])
+        ));
+
+        let expr = parser::parse_expression("re.match.Match", "<filename>")?;
+        assert!(match_module_member(
+            &expr,
+            "typing.re.match",
+            "Match",
             &FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))])
         ));
         Ok(())