Refactor to use mutation instead of passing stacks around

Author: lionel-

crates/ark/src/lsp/symbols.rs

@@ -9,7 +9,6 @@
 
 use std::result::Result::Ok;
 
-use anyhow::anyhow;
 use ropey::Rope;
 use stdext::unwrap::IntoResult;
 use tower_lsp::lsp_types::DocumentSymbol;
@@ -32,8 +31,6 @@ use crate::treesitter::BinaryOperatorType;
 use crate::treesitter::NodeType;
 use crate::treesitter::NodeTypeExt;
 
-type StoreStack = Vec<(usize, Option<DocumentSymbol>, Vec<DocumentSymbol>)>;
-
 fn new_symbol(name: String, kind: SymbolKind, range: Range) -> DocumentSymbol {
     DocumentSymbol {
         name,
@@ -114,6 +111,15 @@ pub fn symbols(params: &WorkspaceSymbolParams) -> anyhow::Result<Vec<SymbolInfor
     Ok(info)
 }
 
+/// Represents a section in the document with its title, level, range, and children
+#[derive(Debug)]
+struct Section {
+    title: String,
+    level: usize,
+    range: Range,
+    children: Vec<DocumentSymbol>,
+}
+
 pub(crate) fn document_symbols(
     state: &WorldState,
     params: &DocumentSymbolParams,
@@ -123,158 +129,123 @@ pub(crate) fn document_symbols(
     let ast = &document.ast;
     let contents = &document.contents;
 
-    let node = ast.root_node();
+    // Start walking from the root node
+    let root_node = ast.root_node();
+    let mut result = Vec::new();
 
-    // Index from the root
-    match index_node(&node, vec![], &contents) {
-        Ok(children) => Ok(children),
-        Err(err) => {
-            log::error!("Error indexing node: {err:?}");
-            return Ok(Vec::new());
-        },
+    // Extract and process all symbols from the AST
+    if let Err(err) = collect_symbols(&root_node, contents, 0, &mut result) {
+        log::error!("Failed to collect symbols: {err:?}");
+        return Ok(Vec::new());
     }
+
+    Ok(result)
 }
 
-fn index_node(
+/// Collect all document symbols from a node recursively
+fn collect_symbols(
     node: &Node,
-    store: Vec<DocumentSymbol>,
     contents: &Rope,
-) -> anyhow::Result<Vec<DocumentSymbol>> {
-    Ok(match node.node_type() {
-        // Handle comment sections in expression lists
+    current_level: usize,
+    symbols: &mut Vec<DocumentSymbol>,
+) -> anyhow::Result<()> {
+    match node.node_type() {
         NodeType::Program | NodeType::BracedExpression => {
-            index_expression_list(&node, store, contents)?
+            collect_sections(node, contents, current_level, symbols)?;
         },
-        // Index assignments as object or function symbols
+
         NodeType::BinaryOperator(BinaryOperatorType::LeftAssignment) |
         NodeType::BinaryOperator(BinaryOperatorType::EqualsAssignment) => {
-            index_assignment(&node, store, contents)?
+            collect_assignment(node, contents, symbols)?;
         },
-        // Nothing to index. FIXME: We should handle argument lists, e.g. to
-        // index inside functions passed as arguments, or inside `test_that()`
-        // blocks.
-        _ => store,
-    })
+
+        // For all other node types, no symbols need to be added
+        _ => {},
+    }
+
+    Ok(())
 }
 
-// Handles root node and braced lists
-fn index_expression_list(
+fn collect_sections(
     node: &Node,
-    store: Vec<DocumentSymbol>,
     contents: &Rope,
-) -> anyhow::Result<Vec<DocumentSymbol>> {
+    current_level: usize,
+    symbols: &mut Vec<DocumentSymbol>,
+) -> anyhow::Result<()> {
+    // In lists of expressions we track and collect section comments, then
+    // collect symbols from children nodes
+
     let mut cursor = node.walk();
 
-    // This is a stack of section levels and associated stores for comments of
-    // the type `# title ----`. It contains all currently active sections.
-    // The top-level section is the current store and has level 0. It should
-    // always be in the stack and popping it before we have finished indexing
-    // the whole expression list is a logic error.
-    let mut store_stack: StoreStack = vec![(0, None, store)];
+    // Track active sections at each level
+    let mut active_sections: Vec<Section> = Vec::new();
 
     for child in node.children(&mut cursor) {
         if let NodeType::Comment = child.node_type() {
-            store_stack = index_comments(&child, store_stack, contents)?;
-            continue;
-        }
-
-        // Get the current store to index the child subtree with.
-        // We restore the store in the stack right after that.
-        let Some((level, symbol, store)) = store_stack.pop() else {
-            return Err(anyhow!(
-                "Internal error: Store stack must have at least one element"
-            ));
-        };
-        let store = index_node(&child, store, contents)?;
-        store_stack.push((level, symbol, store));
-    }
+            let comment_text = contents.node_slice(&child)?.to_string();
+
+            // If we have a section comment, add it to our stack and close any sections if needed
+            if let Some((level, title)) = parse_comment_as_section(&comment_text) {
+                let absolute_level = current_level + level;
+
+                // Close any sections with equal or higher level
+                while !active_sections.is_empty() &&
+                    active_sections.last().unwrap().level >= absolute_level
+                {
+                    finalize_section(&mut active_sections, symbols)?;
+                }
+
+                let range = Range {
+                    start: convert_point_to_position(contents, child.start_position()),
+                    end: convert_point_to_position(contents, child.end_position()),
+                };
+                let section = Section {
+                    title,
+                    level: absolute_level,
+                    range,
+                    children: Vec::new(),
+                };
+                active_sections.push(section);
+            }
 
-    // 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)? {
-            return Ok(store);
+            continue;
         }
-    }
 
-    Err(anyhow!(
-        "Internal error: Store stack must have at least one element"
-    ))
-}
-
-// Pop store from the stack, recording its children and adding it as child to
-// its parent (which becomes the last element in the stack).
-fn store_stack_pop(store_stack: &mut StoreStack) -> anyhow::Result<Option<Vec<DocumentSymbol>>> {
-    let Some((_, symbol, last)) = store_stack.pop() else {
-        return Ok(None);
-    };
-
-    if let Some(mut sym) = symbol {
-        // Assign children to symbol
-        sym.children = Some(last);
-
-        let Some((_, _, ref mut parent_store)) = store_stack.last_mut() else {
-            return Err(anyhow!(
-                "Internal error: Store stack must have at least one element"
-            ));
-        };
-
-        // Store symbol as child of the last symbol on the stack
-        parent_store.push(sym);
-
-        Ok(None)
-    } else {
-        Ok(Some(last))
-    }
-}
-
-fn index_comments(
-    node: &Node,
-    mut store_stack: StoreStack,
-    contents: &Rope,
-) -> anyhow::Result<StoreStack> {
-    let comment_text = contents.node_slice(&node)?.to_string();
-
-    // Check if the comment starts with one or more '#' followed by any text and ends with 4+ punctuations
-    let Some((level, title)) = parse_comment_as_section(&comment_text) else {
-        return Ok(store_stack);
-    };
-
-    // Create a section symbol based on the parsed comment
-    let start = convert_point_to_position(contents, node.start_position());
-    let end = convert_point_to_position(contents, node.end_position());
-    let symbol = new_symbol(title, SymbolKind::STRING, Range { start, end });
-
-    // Now pop all sections still on the stack that have a higher or equal
-    // level. Because we pop sections with equal levels, i.e. siblings, we
-    // ensure that there is only one active section per level on the stack.
-    // That simplifies things because we need to assign popped sections to their
-    // parents and we can assume the relevant parent is always the next on the
-    // stack.
-    loop {
-        let Some((last_level, _, _)) = store_stack.last() else {
-            return Err(anyhow!("Unexpectedly reached the end of the store stack"));
-        };
-
-        if *last_level >= level {
-            if store_stack_pop(&mut store_stack)?.is_some() {
-                return Err(anyhow!("Unexpectedly reached the end of the store stack"));
+        // If we get to this point, `child` is not a section comment.
+        // Recurse into child.
+
+        if active_sections.is_empty() {
+            // If no active section, extend current vector of symbols
+            collect_symbols(&child, contents, current_level, symbols)?;
+        } else {
+            // Otherwise create new store of symbols for the current section
+            let mut child_symbols = Vec::new();
+            collect_symbols(&child, contents, current_level, &mut child_symbols)?;
+
+            // Nest them inside last section
+            if !child_symbols.is_empty() {
+                active_sections
+                    .last_mut()
+                    .unwrap()
+                    .children
+                    .extend(child_symbols);
             }
-            continue;
         }
+    }
 
-        break;
+    // Close any remaining active sections
+    while !active_sections.is_empty() {
+        finalize_section(&mut active_sections, symbols)?;
     }
 
-    store_stack.push((level, Some(symbol), vec![]));
-    Ok(store_stack)
+    Ok(())
 }
 
-fn index_assignment(
+fn collect_assignment(
     node: &Node,
-    mut store: Vec<DocumentSymbol>,
     contents: &Rope,
-) -> anyhow::Result<Vec<DocumentSymbol>> {
+    symbols: &mut Vec<DocumentSymbol>,
+) -> anyhow::Result<()> {
     // Check for assignment
     matches!(
         node.node_type(),
@@ -291,7 +262,7 @@ fn index_assignment(
     let function = lhs.is_identifier_or_string() && rhs.is_function_definition();
 
     if function {
-        return index_assignment_with_function(node, store, contents);
+        return collect_assignment_with_function(node, contents, symbols);
     }
 
     // otherwise, just index as generic object
@@ -301,16 +272,16 @@ fn index_assignment(
     let end = convert_point_to_position(contents, lhs.end_position());
 
     let symbol = new_symbol(name, SymbolKind::VARIABLE, Range { start, end });
-    store.push(symbol);
+    symbols.push(symbol);
 
-    Ok(store)
+    Ok(())
 }
 
-fn index_assignment_with_function(
+fn collect_assignment_with_function(
     node: &Node,
-    mut store: Vec<DocumentSymbol>,
     contents: &Rope,
-) -> anyhow::Result<Vec<DocumentSymbol>> {
+    symbols: &mut Vec<DocumentSymbol>,
+) -> anyhow::Result<()> {
     // check for lhs, rhs
     let lhs = node.child_by_field_name("lhs").into_result()?;
     let rhs = node.child_by_field_name("rhs").into_result()?;
@@ -336,15 +307,36 @@ fn index_assignment_with_function(
 
     let body = rhs.child_by_field_name("body").into_result()?;
 
-    // At this point we increase the nesting level. Recurse into the function
-    // node with a new store of children nodes.
-    let children = index_node(&body, vec![], contents)?;
+    // Process the function body to extract child symbols
+    let mut children = Vec::new();
+    collect_symbols(&body, contents, 0, &mut children)?;
 
     let mut symbol = new_symbol_node(name, SymbolKind::FUNCTION, range, children);
     symbol.detail = Some(detail);
-    store.push(symbol);
+    symbols.push(symbol);
+
+    Ok(())
+}
+
+/// Finalize a section by creating a symbol and adding it to the parent section or output
+fn finalize_section(
+    active_sections: &mut Vec<Section>,
+    symbols: &mut Vec<DocumentSymbol>,
+) -> anyhow::Result<()> {
+    if let Some(section) = active_sections.pop() {
+        let symbol = new_symbol(section.title, SymbolKind::STRING, section.range);
+
+        let mut final_symbol = symbol;
+        final_symbol.children = Some(section.children);
+
+        if let Some(parent) = active_sections.last_mut() {
+            parent.children.push(final_symbol);
+        } else {
+            symbols.push(final_symbol);
+        }
+    }
 
-    Ok(store)
+    Ok(())
 }
 
 // Function to parse a comment and return the section level and title
@@ -374,7 +366,9 @@ mod tests {
         let doc = Document::new(code, None);
         let node = doc.ast.root_node();
 
-        index_node(&node, vec![], &doc.contents).unwrap()
+        let mut symbols = Vec::new();
+        collect_symbols(&node, &doc.contents, 0, &mut symbols).unwrap();
+        symbols
     }
 
     #[test]