Repository: astral-sh/ruff
File: crates/ty_ide/src/folding_range.rs (Rust)

Code before:
            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);
    }

    /// 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);
    }
}

impl SourceOrderVisitor<'_> for FoldingRangeVisitor<'_> {
    fn enter_node(&mut self, node: AnyNodeRef<'_>) -> TraversalSignal {
        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
                // the LSP protocol specifies that this is a problem.
                // If we do need to de-dupe, then we'll want to keep
                // this one since it attaches a "comment" folding range
                // kind to the range. So we'll need to skip over the
                // corresponding range for the literal.
                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
                // string literals.
                self.add_docstring_range(&class.body);
            }
            AnyNodeRef::StmtIf(if_stmt) => {
                // Fold each branch individually rather than the entire if block.
                // The if clause range is from the start of the if to the end of its body.
                if let Some(last_stmt) = if_stmt.body.last() {
                    self.add_range(TextRange::new(if_stmt.start(), last_stmt.end()));
                }
                // Each elif/else clause has its own range.

Review comment (suggestion):
I think it would be a little clearer if `self.add_decorator_range(&func.decorator_list);` were moved into `add_function_def_range`. And similarly for below for classes.

Code after:
                TokenKind::Def
            };
            let search_range = TextRange::new(decorator.end(), func.name.start());
            self.add_def_range(target, search_range, func.end());
        } else {
            self.add_range(func.range());
        }
    }

    /// Add a folding range for class definitions, excluding decorators.
    fn add_class_def_range(&mut self, class: &StmtClassDef) {
        if let Some(decorator) = class.decorator_list.last() {
            let search_range = TextRange::new(decorator.end(), class.name.start());
            self.add_def_range(TokenKind::Class, search_range, class.end());
        } else {
            self.add_range(class.range());
        }
    }
}

impl SourceOrderVisitor<'_> for FoldingRangeVisitor<'_> {
    fn enter_node(&mut self, node: AnyNodeRef<'_>) -> TraversalSignal {
        match node {
            // Compound statements that create folding regions
            AnyNodeRef::StmtFunctionDef(func) => {
                self.add_function_def_range(func);
                // Note that this may be duplicative with folding
                // ranges added for string literals. But I don't think
                // the LSP protocol specifies that this is a problem.
                // If we do need to de-dupe, then we'll want to keep
                // this one since it attaches a "comment" folding range
                // kind to the range. So we'll need to skip over the
                // corresponding range for the literal.
                self.add_docstring_range(&func.body);
            }
            AnyNodeRef::StmtClassDef(class) => {
                self.add_class_def_range(class);
                // See comment above for class docstrings about this
                // being duplicative with adding folding ranges for
                // string literals.
                self.add_docstring_range(&class.body);
            }
            AnyNodeRef::StmtIf(if_stmt) => {
                // Fold each branch individually rather than the entire if block.
                // The if clause range is from the start of the if to the end of its body.
                if let Some(last_stmt) = if_stmt.body.last() {
                    self.add_range(TextRange::new(if_stmt.start(), last_stmt.end()));
                }
                // Each elif/else clause has its own range.
                for clause in &if_stmt.elif_else_clauses {
                    self.add_range(clause.range());