Simplify argument extraction

Author: lionel-

crates/ark/src/lsp/diagnostics.rs

@@ -19,13 +19,15 @@ use tower_lsp::lsp_types::DiagnosticSeverity;
 use tree_sitter::Node;
 use tree_sitter::Range;
 
+use crate::lsp;
 use crate::lsp::declarations::top_level_declare;
 use crate::lsp::diagnostics_syntax::syntax_diagnostics;
 use crate::lsp::documents::Document;
 use crate::lsp::encoding::convert_tree_sitter_range_to_lsp_range;
 use crate::lsp::indexer;
 use crate::lsp::inputs::library::Library;
 use crate::lsp::state::WorldState;
+use crate::lsp::traits::node::NodeExt;
 use crate::lsp::traits::rope::RopeExt;
 use crate::treesitter::node_has_error_or_missing;
 use crate::treesitter::BinaryOperatorType;
@@ -799,59 +801,34 @@ fn recurse_call(
     match fun {
         "library" => {
             // Track symbols exported by `library()` calls
-            handle_library_call(node, context, diagnostics)?;
+            handle_library_call(node, context)?;
         },
-        // default case: recurse into each argument
-        _ => recurse_call_like_arguments_default(node, context, diagnostics)?,
+        _ => {},
     };
 
+    // Continue with default recursion to handle any other arguments
+    recurse_call_like_arguments_default(node, context, diagnostics)?;
+
     ().ok()
 }
 
-fn handle_library_call(
-    node: Node,
-    context: &mut DiagnosticContext,
-    diagnostics: &mut Vec<Diagnostic>,
-) -> anyhow::Result<()> {
-    // Get the arguments node
-    let Some(arguments) = node.child_by_field_name("arguments") else {
-        return Ok(());
-    };
-
-    // Find the first argument (package name)
-    let mut cursor = arguments.walk();
-    let mut children = arguments.children_by_field_name("argument", &mut cursor);
-
-    let Some(first_arg) = children.next() else {
-        return Ok(());
-    };
-
-    // Get the package name from the argument value
-    let Some(value) = first_arg.child_by_field_name("value") else {
-        return Ok(());
+fn handle_library_call(node: Node, context: &mut DiagnosticContext) -> anyhow::Result<()> {
+    // Find the first argument (package name). Positionally for now.
+    let Some(value) = node.arguments_values().nth(0) else {
+        return Err(anyhow::anyhow!("Can't unpack `library()` argument"));
     };
 
-    let package_name = if value.is_identifier() {
-        context.contents.node_slice(&value)?.to_string()
-    } else if value.is_string() {
-        // Remove quotes from string literal
-        let raw = context.contents.node_slice(&value)?.to_string();
-        raw.trim_matches('"').trim_matches('\'').to_string()
-    } else {
-        return Ok(());
-    };
+    let package_name = value.get_identifier_or_string_text(context.contents)?;
 
-    // Try to get the package from the library
+    // Insert exports globablly for now
     if let Some(package) = context.library.get(&package_name) {
-        // Add all exported symbols to library_symbols
         for symbol in &package.namespace.exports {
             context.library_symbols.insert(symbol.clone());
         }
+    } else {
+        lsp::log_warn!("Can't get exports from package {package_name} because it is not installed.")
     }
 
-    // Continue with default recursion to handle any other arguments
-    recurse_call_like_arguments_default(node, context, diagnostics)?;
-
     Ok(())
 }
 

crates/ark/src/lsp/traits/node.rs

@@ -90,6 +90,9 @@ pub trait NodeExt: Sized {
     fn bwd_leaf_iter(&self) -> BwdLeafIterator<'_>;
 
     fn ancestors(&self) -> impl Iterator<Item = Self>;
+    fn children_of(node: Self) -> impl Iterator<Item = Self>;
+    fn next_siblings(&self) -> impl Iterator<Item = Self>;
+    fn arguments_values(&self) -> impl Iterator<Item = Self>;
 }
 
 impl<'tree> NodeExt for Node<'tree> {
@@ -199,13 +202,68 @@ impl<'tree> NodeExt for Node<'tree> {
         BwdLeafIterator { node: *self }
     }
 
-    // From rowan. Note that until we switch to rowan, each `parent()` call
-    // causes a linear traversal of the whole tree to find the parent node.
-    // We could do better in the future:
-    // https://github.com/tree-sitter/tree-sitter/pull/3214
     fn ancestors(&self) -> impl Iterator<Item = Node<'tree>> {
+        // We'd ideally use the cursor API here too but
+        // `ts_tree_cursor_goto_parent()` doesn't behave like
+        // `ts_node_parent()`: the latter traverses `ERROR` nodes but not the
+        // former. So for now we accept the performance hit of tree traversal at
+        // each `parent()` call.
         std::iter::successors(Some(*self), |p| p.parent())
     }
+
+    fn next_siblings(&self) -> impl Iterator<Item = Node<'tree>> {
+        let mut cursor = self.walk();
+
+        let first = if cursor.goto_next_sibling() {
+            Some(cursor.node())
+        } else {
+            None
+        };
+
+        std::iter::successors(first, move |_| {
+            if cursor.goto_next_sibling() {
+                Some(cursor.node())
+            } else {
+                None
+            }
+        })
+    }
+
+    fn children_of(node: Node<'tree>) -> impl Iterator<Item = Node<'tree>> {
+        let mut cursor = node.walk();
+        let mut first = true;
+
+        std::iter::from_fn(move || {
+            let advanced = if first {
+                first = false;
+                cursor.goto_first_child()
+            } else {
+                cursor.goto_next_sibling()
+            };
+
+            if advanced {
+                Some(cursor.node())
+            } else {
+                None
+            }
+        })
+    }
+
+    /// Takes a call node and iterates over the values of its arguments
+    fn arguments_values(&self) -> impl Iterator<Item = Node<'tree>> {
+        self.child_by_field_name("arguments")
+            // Create iterator that unpacks Option with `flat_map()`
+            .into_iter()
+            .flat_map(Self::children_of)
+            .filter_map(|node| {
+                // This takes care of non-argument nodes like `(` and `)`
+                if node.kind() == "argument" {
+                    node.child_by_field_name("value")
+                } else {
+                    None
+                }
+            })
+    }
 }
 
 /// First, recurse through children to find the smallest

crates/ark/src/treesitter.rs

@@ -283,6 +283,7 @@ pub trait NodeTypeExt: Sized {
     fn is_identifier(&self) -> bool;
     fn is_string(&self) -> bool;
     fn is_identifier_or_string(&self) -> bool;
+    fn get_identifier_or_string_text(&self, contents: &ropey::Rope) -> anyhow::Result<String>;
     fn is_keyword(&self) -> bool;
     fn is_call(&self) -> bool;
     fn is_subset(&self) -> bool;
@@ -325,6 +326,20 @@ impl NodeTypeExt for Node<'_> {
         matches!(self.node_type(), NodeType::Identifier | NodeType::String)
     }
 
+    fn get_identifier_or_string_text(&self, contents: &ropey::Rope) -> anyhow::Result<String> {
+        match self.node_type() {
+            NodeType::Identifier => return Ok(contents.node_slice(self)?.to_string()),
+            NodeType::String => {
+                // Remove quotes from string literal
+                let string = contents.node_slice(self)?.to_string();
+                Ok(string.trim_matches('"').trim_matches('\'').to_string())
+            },
+            _ => {
+                return Err(anyhow::anyhow!("Not an identifier or string"));
+            },
+        }
+    }
+
     fn is_keyword(&self) -> bool {
         matches!(
             self.node_type(),