From 3abd205f947d3845e64cb0340decf5cf4b94f8ff Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Fri, 6 Jan 2023 22:55:13 -0500
Subject: [PATCH] Lazily compute ranges for class and function bindings (#1708)

This has a fairly significant performance impact (~320ms to ~307ms on
the CPython benchmark).
---
 src/ast/helpers.rs  | 18 +++++++++++++++++-
 src/checkers/ast.rs | 10 +++++-----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/ast/helpers.rs b/src/ast/helpers.rs
index f3cb58964e..3c46b7c419 100644
--- a/src/ast/helpers.rs
+++ b/src/ast/helpers.rs
@@ -11,7 +11,7 @@ use rustpython_parser::lexer;
 use rustpython_parser::lexer::Tok;
 use rustpython_parser::token::StringKind;
 
-use crate::ast::types::Range;
+use crate::ast::types::{Binding, BindingKind, Range};
 use crate::source_code_generator::SourceCodeGenerator;
 use crate::source_code_style::SourceCodeStyleDetector;
 use crate::SourceCodeLocator;
@@ -406,6 +406,22 @@ pub fn identifier_range(stmt: &Stmt, locator: &SourceCodeLocator) -> Range {
     Range::from_located(stmt)
 }
 
+/// Like `identifier_range`, but accepts a `Binding`.
+pub fn binding_range(binding: &Binding, locator: &SourceCodeLocator) -> Range {
+    if matches!(
+        binding.kind,
+        BindingKind::ClassDefinition | BindingKind::FunctionDefinition
+    ) {
+        if let Some(source) = &binding.source {
+            identifier_range(source, locator)
+        } else {
+            binding.range
+        }
+    } else {
+        binding.range
+    }
+}
+
 // Return the ranges of `Name` tokens within a specified node.
 pub fn find_names<T>(located: &Located<T>, locator: &SourceCodeLocator) -> Vec<Range> {
     let contents = locator.slice_source_code_range(&Range::from_located(located));
diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs
index 0d1865dbab..2d07234f91 100644
--- a/src/checkers/ast.rs
+++ b/src/checkers/ast.rs
@@ -15,7 +15,7 @@ use rustpython_parser::ast::{
 use rustpython_parser::parser;
 
 use crate::ast::helpers::{
-    collect_call_paths, dealias_call_path, extract_handler_names, match_call_path,
+    binding_range, collect_call_paths, dealias_call_path, extract_handler_names, match_call_path,
 };
 use crate::ast::operations::extract_all_names;
 use crate::ast::relocate::relocate_expr;
@@ -560,7 +560,7 @@ where
                     Binding {
                         kind: BindingKind::FunctionDefinition,
                         used: None,
-                        range: helpers::identifier_range(stmt, self.locator),
+                        range: Range::from_located(stmt),
                         source: Some(self.current_stmt().clone()),
                     },
                 );
@@ -1544,7 +1544,7 @@ where
                     Binding {
                         kind: BindingKind::ClassDefinition,
                         used: None,
-                        range: helpers::identifier_range(stmt, self.locator),
+                        range: Range::from_located(stmt),
                         source: Some(self.current_stmt().clone()),
                     },
                 );
@@ -3373,7 +3373,7 @@ impl<'a> Checker<'a> {
                                     name.to_string(),
                                     existing.range.location.row(),
                                 ),
-                                binding.range,
+                                binding_range(&binding, self.locator),
                             ));
                         }
                     }
@@ -3935,7 +3935,7 @@ impl<'a> Checker<'a> {
                                         (*name).to_string(),
                                         binding.range.location.row(),
                                     ),
-                                    self.bindings[*index].range,
+                                    binding_range(&self.bindings[*index], self.locator),
                                 ));
                             }
                         }