Feature/nested-outline

crates/ark/src/lsp/symbols.rs

@@ -104,8 +104,18 @@ pub(crate) fn document_symbols(
         selection_range: Range { start, end },
     };
 
+    // Maintain a stack to track the hierarchy of comment sections
+    let mut section_stack: Vec<(usize, *mut DocumentSymbol)> =
+        vec![(0, &mut root as *mut DocumentSymbol)];
+
     // index from the root
-    index_node(&node, &contents, &mut root, &mut symbols)?;
+    index_node(
+        &node,
+        &contents,
+        &mut root,
+        &mut symbols,
+        &mut section_stack,
+    )?;
 
     // return the children we found
     Ok(root.children.unwrap_or_default())
@@ -141,58 +151,79 @@ fn index_node(
     contents: &Rope,
     parent: &mut DocumentSymbol,
     symbols: &mut Vec<DocumentSymbol>,
+    section_stack: &mut Vec<(usize, *mut DocumentSymbol)>,
 ) -> Result<bool> {
-    // Check if the node is a comment and matches the markdown-style comment patterns
+    // Maintain hierarchy for comment sections
     if node.node_type() == NodeType::Comment {
         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
-        if let Some((_level, title)) = parse_comment_as_section(&comment_text) {
-            // Create a symbol based on the parsed comment
+        if let Some((level, title)) = parse_comment_as_section(&comment_text) {
             let start = convert_point_to_position(contents, node.start_position());
             let end = convert_point_to_position(contents, node.end_position());
 
             let symbol = DocumentSymbol {
-                name: title,                // Use the title without the trailing '####' or '----'
-                kind: SymbolKind::STRING,   // Treat it as a string section
-                detail: None,               // No need to display level details
-                children: Some(Vec::new()), // Prepare for child symbols if any
+                name: title,
+                kind: SymbolKind::STRING,
+                detail: None,
+                children: Some(Vec::new()),
                 deprecated: None,
                 tags: None,
                 range: Range { start, end },
                 selection_range: Range { start, end },
             };
 
-            // Add the symbol to the parent node
-            parent.children.as_mut().unwrap().push(symbol);
+            // Pop until we find a suitable parent for this level
+            while let Some((current_level, _)) = section_stack.last() {
+                if *current_level >= level {
+                    section_stack.pop();
+                } else {
+                    break;
+                }
+            }
+
+            // Push the new symbol to the current parent
+            if let Some((_, current_parent_ptr)) = section_stack.last() {
+                let current_parent = unsafe { &mut **current_parent_ptr };
+                current_parent.children.as_mut().unwrap().push(symbol);
+                let new_parent = current_parent
+                    .children
+                    .as_mut()
+                    .unwrap()
+                    .last_mut()
+                    .unwrap();
+                section_stack.push((level, new_parent as *mut DocumentSymbol));
+            } else {
+                parent.children.as_mut().unwrap().push(symbol);
+                let new_parent = parent.children.as_mut().unwrap().last_mut().unwrap();
+                section_stack.push((level, new_parent as *mut DocumentSymbol));
+            }
 
-            // Return early to avoid further processing
             return Ok(true);
         }
     }
 
+    // Handle assignments (like `a <- 1`)
     if matches!(
         node.node_type(),
         NodeType::BinaryOperator(BinaryOperatorType::LeftAssignment) |
             NodeType::BinaryOperator(BinaryOperatorType::EqualsAssignment)
     ) {
-        match index_assignment(node, contents, parent, symbols) {
-            Ok(handled) => {
-                if handled {
-                    return Ok(true);
-                }
-            },
-            Err(error) => error!("{:?}", error),
+        if let Some((_, current_parent_ptr)) = section_stack.last() {
+            let current_parent = unsafe { &mut **current_parent_ptr };
+            return index_assignment(node, contents, current_parent, symbols, section_stack);
         }
     }
 
     // Recurse into children
     let mut cursor = node.walk();
     for child in node.children(&mut cursor) {
         if is_indexable(&child) {
-            let result = index_node(&child, contents, parent, symbols);
-            if let Err(error) = result {
-                error!("{:?}", error);
+            if let Some((_, current_parent_ptr)) = section_stack.last() {
+                let current_parent = unsafe { &mut **current_parent_ptr };
+                let result = index_node(&child, contents, current_parent, symbols, section_stack);
+                if let Err(error) = result {
+                    error!("{:?}", error);
+                }
             }
         }
     }
@@ -205,6 +236,7 @@ fn index_assignment(
     contents: &Rope,
     parent: &mut DocumentSymbol,
     symbols: &mut Vec<DocumentSymbol>,
+    section_stack: &mut Vec<(usize, *mut DocumentSymbol)>,
 ) -> Result<bool> {
     // check for assignment
     matches!(
@@ -218,22 +250,21 @@ fn index_assignment(
     let lhs = node.child_by_field_name("lhs").into_result()?;
     let rhs = node.child_by_field_name("rhs").into_result()?;
 
-    // check for identifier on lhs, function on rhs
-    let function = lhs.is_identifier_or_string() && rhs.is_function_definition();
-
-    if function {
-        return index_assignment_with_function(node, contents, parent, symbols);
+    // check if lhs is an identifier and rhs is a function definition
+    if lhs.is_identifier_or_string() && rhs.is_function_definition() {
+        // If the RHS is a function definition, we call `index_assignment_with_function`
+        return index_assignment_with_function(node, contents, parent, symbols, section_stack);
     }
 
-    // otherwise, just index as generic object
+    // otherwise, just index as a generic variable/object
     let name = contents.node_slice(&lhs)?.to_string();
 
     let start = convert_point_to_position(contents, lhs.start_position());
     let end = convert_point_to_position(contents, lhs.end_position());
 
     let symbol = DocumentSymbol {
         name,
-        kind: SymbolKind::OBJECT,
+        kind: SymbolKind::VARIABLE, // Use VARIABLE kind to represent an assignment
         detail: None,
         children: Some(Vec::new()),
         deprecated: None,
@@ -242,7 +273,6 @@ fn index_assignment(
         selection_range: Range::new(start, end),
     };
 
-    // add this symbol to the parent node
     parent.children.as_mut().unwrap().push(symbol);
 
     Ok(true)
@@ -253,27 +283,29 @@ fn index_assignment_with_function(
     contents: &Rope,
     parent: &mut DocumentSymbol,
     symbols: &mut Vec<DocumentSymbol>,
+    section_stack: &mut Vec<(usize, *mut DocumentSymbol)>,
 ) -> Result<bool> {
     // check for lhs, rhs
     let lhs = node.child_by_field_name("lhs").into_result()?;
     let rhs = node.child_by_field_name("rhs").into_result()?;
 
     // start extracting the argument names
     let mut arguments: Vec<String> = Vec::new();
-    let parameters = rhs.child_by_field_name("parameters").into_result()?;
-
-    let mut cursor = parameters.walk();
-    for parameter in parameters.children_by_field_name("parameter", &mut cursor) {
-        let name = parameter.child_by_field_name("name").into_result()?;
-        let name = contents.node_slice(&name)?.to_string();
-        arguments.push(name);
+    if let Some(parameters) = rhs.child_by_field_name("parameters") {
+        let mut cursor = parameters.walk();
+        for parameter in parameters.children_by_field_name("parameter", &mut cursor) {
+            if let Some(name) = parameter.child_by_field_name("name") {
+                let name = contents.node_slice(&name)?.to_string();
+                arguments.push(name);
+            }
+        }
     }
 
     let name = contents.node_slice(&lhs)?.to_string();
     let detail = format!("function({})", arguments.join(", "));
 
     // build the document symbol
-    let symbol = DocumentSymbol {
+    let mut symbol = DocumentSymbol {
         name,
         kind: SymbolKind::FUNCTION,
         detail: Some(detail),
@@ -290,12 +322,21 @@ fn index_assignment_with_function(
         },
     };
 
-    // add this symbol to the parent node
+    // add the function symbol to the parent
     parent.children.as_mut().unwrap().push(symbol);
 
-    // recurse into this node
-    let parent = parent.children.as_mut().unwrap().last_mut().unwrap();
-    index_node(&rhs, contents, parent, symbols)?;
+    // Get a mutable reference to the newly added symbol
+    let new_parent = parent.children.as_mut().unwrap().last_mut().unwrap();
+    let new_parent_ptr = new_parent as *mut DocumentSymbol;
+
+    // Add the function to the section stack
+    section_stack.push((0, new_parent_ptr));
+
+    // Recurse into the function body to add its children
+    index_node(&rhs, contents, new_parent, symbols, section_stack)?;
+
+    // Pop the stack after processing the function body
+    section_stack.pop();
 
     Ok(true)
 }
@@ -327,7 +368,20 @@ mod tests {
             selection_range: Range { start, end },
         };
 
-        index_node(&node, &doc.contents, &mut root, &mut symbols).unwrap();
+        // Create a section_stack to pass to index_node
+        let mut section_stack: Vec<(usize, *mut DocumentSymbol)> =
+            vec![(0, &mut root as *mut DocumentSymbol)];
+
+        // Call index_node with the new argument
+        index_node(
+            &node,
+            &doc.contents,
+            &mut root,
+            &mut symbols,
+            &mut section_stack,
+        )
+        .unwrap();
+
         root.children.unwrap_or_default()
     }