diff --git a/crates/ark/src/lsp/snapshots/ark__lsp__symbols__tests__symbol_assignment_function_nested_section.snap b/crates/ark/src/lsp/snapshots/ark__lsp__symbols__tests__symbol_assignment_function_nested_section.snap index ddef8b970..fe530c7e5 100644 --- a/crates/ark/src/lsp/snapshots/ark__lsp__symbols__tests__symbol_assignment_function_nested_section.snap +++ b/crates/ark/src/lsp/snapshots/ark__lsp__symbols__tests__symbol_assignment_function_nested_section.snap @@ -1,5 +1,6 @@ --- source: crates/ark/src/lsp/symbols.rs +assertion_line: 542 expression: "test_symbol(\"\n## title0 ----\nfoo <- function() {\n # title1 ----\n ### title2 ----\n ## title3 ----\n # title4 ----\n}\n# title5 ----\")" --- [ @@ -15,8 +16,8 @@ expression: "test_symbol(\"\n## title0 ----\nfoo <- function() {\n # title1 --- character: 0, }, end: Position { - line: 1, - character: 14, + line: 7, + character: 0, }, }, selection_range: Range { @@ -73,8 +74,8 @@ expression: "test_symbol(\"\n## title0 ----\nfoo <- function() {\n # title1 --- character: 2, }, end: Position { - line: 3, - character: 15, + line: 5, + character: 0, }, }, selection_range: Range { @@ -102,7 +103,7 @@ expression: "test_symbol(\"\n## title0 ----\nfoo <- function() {\n # title1 --- }, end: Position { line: 4, - character: 17, + character: 0, }, }, selection_range: Range { @@ -132,7 +133,7 @@ expression: "test_symbol(\"\n## title0 ----\nfoo <- function() {\n # title1 --- }, end: Position { line: 5, - character: 16, + character: 0, }, }, selection_range: Range { @@ -164,8 +165,8 @@ expression: "test_symbol(\"\n## title0 ----\nfoo <- function() {\n # title1 --- character: 2, }, end: Position { - line: 6, - character: 15, + line: 7, + character: 0, }, }, selection_range: Range { diff --git a/crates/ark/src/lsp/symbols.rs b/crates/ark/src/lsp/symbols.rs index bdd31f13b..fd2ebe446 100644 --- a/crates/ark/src/lsp/symbols.rs +++ b/crates/ark/src/lsp/symbols.rs @@ -135,6 +135,61 @@ pub(crate) fn document_symbols( } } +fn fix_section_ranges(symbols: &mut Vec, contents: &Rope) { + fix_section_ranges_recursive(symbols, contents, None); +} + + +fn fix_section_ranges_recursive( + symbols: &mut Vec, + contents: &Rope, + parent_end: Option, +) { + for i in 0..symbols.len() { + // Extract the values we need before creating mutable borrows + let symbol_kind = symbols[i].kind; + + // Only fix ranges for section symbols (STRING kind) + if symbol_kind != SymbolKind::STRING { + // Still need to recurse into children for non-section symbols + if let Some(ref mut children) = symbols[i].children { + fix_section_ranges_recursive(children, contents, parent_end); + } + continue; + } + + // Default end position + let mut end_pos = parent_end.unwrap_or_else(|| { + let last_line_idx = contents.len_lines().saturating_sub(1); + let last_line_len = contents.line(last_line_idx).len_chars(); + tower_lsp::lsp_types::Position { + line: last_line_idx as u32, + character: last_line_len as u32, + } + }); + + // Look for the next sibling section + for j in (i + 1)..symbols.len() { + if symbols[j].kind == SymbolKind::STRING { + // This section ends just before the next sibling section + end_pos = tower_lsp::lsp_types::Position { + line: symbols[j].range.start.line.saturating_sub(1), + character: 0, + }; + break; + } + } + + // Update the range + symbols[i].range.end = end_pos; + + // Now recursively fix children with the calculated end position + if let Some(ref mut children) = symbols[i].children { + fix_section_ranges_recursive(children, contents, Some(end_pos)); + } + } +} + fn index_node( node: &Node, store: Vec, @@ -192,7 +247,9 @@ fn index_expression_list( // Pop all sections from the stack, assigning their childrens and their // parents along the way while store_stack.len() > 0 { - if let Some(store) = store_stack_pop(&mut store_stack)? { + if let Some(mut store) = store_stack_pop(&mut store_stack)? { + // Fix section ranges to extend until the next section of same or higher level + fix_section_ranges(&mut store, contents); return Ok(store); } } @@ -511,4 +568,32 @@ foo <- function() { assert_eq!(test_symbol("{ foo <- 1 }"), vec![foo]); } + + #[test] + fn test_extended_section_ranges() { + let symbols = test_symbol( + "# h1 ---- +library(tidyverse) + +## h2 ---- +my_func <- function(x) { + x + 1 +} + +# Another h1 ---- +x <- 1" + ); + + // The first h1 section should extend to just before "# Another h1 ----" + // The h2 section should extend to just before "# Another h1 ----" + // The second h1 section should extend to end of file + + assert_eq!(symbols.len(), 2); // Two h1 sections + assert_eq!(symbols[0].name, "h1"); + assert_eq!(symbols[0].children.as_ref().unwrap()[0].name, "h2"); + assert_eq!(symbols[0].children.as_ref().unwrap()[0].children.as_ref().unwrap()[0].name, "my_func"); + assert_eq!(symbols[1].name, "Another h1"); + assert!(symbols[0].range.end.line == 7); + assert!(symbols[1].range.end.line == 9); + } }