use line index over tokenizer

Author: iksuddle

crates/ty_ide/src/folding_range.rs

@@ -3,7 +3,6 @@ use ruff_db::parsed::parsed_module;
 use ruff_db::source::source_text;
 use ruff_python_ast::visitor::source_order::{SourceOrderVisitor, TraversalSignal, walk_body};
 use ruff_python_ast::{AnyNodeRef, Decorator, Stmt, StmtClassDef, StmtFunctionDef};
-use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_source_file::{Line, UniversalNewlines};
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
@@ -49,9 +48,12 @@ pub fn folding_ranges(db: &dyn Db, file: File) -> Vec<FoldingRange> {
     let parsed = parsed_module(db, file).load(db);
     let source = source_text(db, file);
 
+    let line_index = ruff_db::source::line_index(db, file);
+
     let mut visitor = FoldingRangeVisitor {
         source: source.as_str(),
         ranges: vec![],
+        line_index: &line_index,
     };
     visitor.visit_body(parsed.suite());
 
@@ -68,6 +70,7 @@ pub fn folding_ranges(db: &dyn Db, file: File) -> Vec<FoldingRange> {
 struct FoldingRangeVisitor<'a> {
     source: &'a str,
     ranges: Vec<FoldingRange>,
+    line_index: &'a ruff_source_file::LineIndex,
 }
 
 impl<'a> FoldingRangeVisitor<'a> {
@@ -233,38 +236,33 @@ impl<'a> FoldingRangeVisitor<'a> {
     }
 
     /// Add a folding range for the function or class definition.
-    ///
-    /// - target is `async` or `def` for functions, and `class` for classes
     fn add_def_range(
         &mut self,
         range: TextRange,
         decorator_list: &[Decorator],
-        target: SimpleTokenKind,
+        name_start: TextSize,
     ) {
-        if let Some(last) = decorator_list.last() {
-            let tokenizer = SimpleTokenizer::starts_at(last.end(), self.source);
-            if let Some(token) = tokenizer.skip_trivia().find(|token| token.kind == target) {
-                let range = TextRange::new(token.start(), range.end());
-                self.add_range(range);
-            }
+        if decorator_list.last().is_some() {
+            // get the the beginning of the line containing the function or class name
+            let line = self.line_index.line_index(name_start);
+            let start = self.line_index.line_start(line, self.source);
+            let range = TextRange::new(start, range.end());
+            self.add_range(range);
         } else {
             self.add_range(range);
         }
     }
 
     /// Add a folding range for function definitions, excluding decorators.
     fn add_function_def_range(&mut self, func: &StmtFunctionDef) {
-        let target = if func.is_async {
-            SimpleTokenKind::Async
-        } else {
-            SimpleTokenKind::Def
-        };
-        self.add_def_range(func.range(), &func.decorator_list, target);
+        self.add_decorator_range(&func.decorator_list);
+        self.add_def_range(func.range(), &func.decorator_list, func.name.start());
     }
 
     /// Add a folding range for class definitions, excluding decorators.
     fn add_class_def_range(&mut self, class: &StmtClassDef) {
-        self.add_def_range(class.range(), &class.decorator_list, SimpleTokenKind::Class);
+        self.add_decorator_range(&class.decorator_list);
+        self.add_def_range(class.range(), &class.decorator_list, class.name.start());
     }
 }
 
@@ -273,7 +271,6 @@ impl SourceOrderVisitor<'_> for FoldingRangeVisitor<'_> {
         match node {
             // Compound statements that create folding regions
             AnyNodeRef::StmtFunctionDef(func) => {
-                self.add_decorator_range(&func.decorator_list);
                 self.add_function_def_range(func);
                 // Note that this may be duplicative with folding
                 // ranges added for string literals. But I don't think
@@ -285,7 +282,6 @@ impl SourceOrderVisitor<'_> for FoldingRangeVisitor<'_> {
                 self.add_docstring_range(&func.body);
             }
             AnyNodeRef::StmtClassDef(class) => {
-                self.add_decorator_range(&class.decorator_list);
                 self.add_class_def_range(class);
                 // See comment above for class docstrings about this
                 // being duplicative with adding folding ranges for
@@ -1729,7 +1725,7 @@ with open("file.txt") as f:
     }
 
     #[test]
-    fn test_folding_range_decorated_function1() {
+    fn test_folding_range_decorated_function_single() {
         let test = CursorTest::builder()
             .source(
                 "main.py",
@@ -1755,7 +1751,7 @@ def my_function():
     }
 
     #[test]
-    fn test_folding_range_decorated_function2() {
+    fn test_folding_range_decorated_function_multiple() {
         let test = CursorTest::builder()
             .source(
                 "main.py",
@@ -1795,7 +1791,7 @@ def my_function():
     }
 
     #[test]
-    fn test_folding_range_decorated_class1() {
+    fn test_folding_range_decorated_class_single() {
         let test = CursorTest::builder()
             .source(
                 "main.py",
@@ -1824,7 +1820,7 @@ class MyClass:
     }
 
     #[test]
-    fn test_folding_range_decorated_class2() {
+    fn test_folding_range_decorated_class_multiple() {
         let test = CursorTest::builder()
             .source(
                 "main.py",
@@ -1861,6 +1857,108 @@ class MyClass:
         ");
     }
 
+    #[test]
+    fn test_folding_range_decorated_async_function() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+@decorator
+async def my_async_function():
+    pass
+<CURSOR>
+"#,
+            )
+            .build();
+
+        assert_snapshot!(test.folding_ranges(), @r"
+        info[folding-range]: Folding Range
+         --> main.py:3:1
+          |
+        2 |   @decorator
+        3 | / async def my_async_function():
+        4 | |     pass
+          | |________^
+          |
+        ");
+    }
+
+    #[test]
+    fn test_folding_range_decorated_nested_function() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+def outer_function():
+    @decorator
+    def inner_function():
+        pass
+<CURSOR>
+"#,
+            )
+            .build();
+
+        assert_snapshot!(test.folding_ranges(), @"
+        info[folding-range]: Folding Range
+         --> main.py:2:1
+          |
+        2 | / def outer_function():
+        3 | |     @decorator
+        4 | |     def inner_function():
+        5 | |         pass
+          | |____________^
+          |
+
+        info[folding-range]: Folding Range
+         --> main.py:4:1
+          |
+        2 |   def outer_function():
+        3 |       @decorator
+        4 | /     def inner_function():
+        5 | |         pass
+          | |____________^
+          |
+        ");
+    }
+
+    #[test]
+    fn test_folding_range_decorated_async_method() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+class MyClass:
+    @decorator
+    async def my_async_method(self):
+        pass
+<CURSOR>
+"#,
+            )
+            .build();
+
+        assert_snapshot!(test.folding_ranges(), @"
+        info[folding-range]: Folding Range
+         --> main.py:2:1
+          |
+        2 | / class MyClass:
+        3 | |     @decorator
+        4 | |     async def my_async_method(self):
+        5 | |         pass
+          | |____________^
+          |
+
+        info[folding-range]: Folding Range
+         --> main.py:4:1
+          |
+        2 |   class MyClass:
+        3 |       @decorator
+        4 | /     async def my_async_method(self):
+        5 | |         pass
+          | |____________^
+          |
+        ");
+    }
+
     struct FoldingRangeDiagnostic {
         file: File,
         folding_range: FoldingRange,