feat: add queries for function parameters and identifiers

Author: SunBK201

scubatrace/cpp/language.py

@@ -8,6 +8,10 @@ class C(Language):
     extensions = ["c", "h", "cc", "cpp", "cxx", "hxx", "hpp"]
     tslanguage = TSLanguage(tscpp.language())
 
+    query_identifier = """
+        (identifier)@name
+        (field_identifier)@name
+    """
     query_call = "(call_expression)@name"
     query_import_identifier = """
         (preproc_include
@@ -17,6 +21,37 @@ class C(Language):
             ]
         )
     """
+    query_function_parameter = """
+        (parameter_declaration
+        	declarator: [
+            	(identifier)@name
+                (pointer_declarator
+                	(identifier)@name
+                )
+                (pointer_declarator
+                	(pointer_declarator
+                		(identifier)@name
+                    )
+                )
+                (pointer_declarator
+                	(pointer_declarator
+                        (pointer_declarator
+                            (identifier)@name
+                        )
+                    )
+                )
+                (pointer_declarator
+                    (pointer_declarator
+                        (pointer_declarator
+                            (pointer_declarator
+                                (identifier)@name
+                            )
+                        )
+                    )
+                )
+            ]
+        )
+    """
 
     query_struct = "(struct_specifier)@name"
     query_class = "(class_specifier)@name"

scubatrace/file.py

@@ -454,3 +454,47 @@ def query_oneshot(self, query: str) -> Statement | None:
         if len(matched_statements) == 0:
             return None
         return matched_statements[0]
+
+    def query_identifiers(
+        self, query: str, node: Node | None = None
+    ) -> list[Identifier]:
+        """
+        Executes a tree-sitter query to find identifiers in the file.
+
+        Args:
+            identifier (str): The identifier to search for.
+            node (Node | None): The tree-sitter node to query. If None, uses the root node of the file.
+
+        Returns:
+            list[Identifier]: A list of identifiers that contain the specified identifier.
+        """
+        if node is None:
+            node = self.node
+        matched_nodes = self.parser.query_all(node, query)
+        matched_identifiers = []
+        for identifier in self.identifiers:
+            for node in matched_nodes:
+                if (
+                    identifier.node.start_byte >= node.start_byte
+                    and identifier.node.end_byte <= node.end_byte
+                ):
+                    matched_identifiers.append(identifier)
+                    break
+        return matched_identifiers
+
+    def query_identifier(
+        self, query: str, node: Node | None = None
+    ) -> Identifier | None:
+        """
+        Executes a tree-sitter oneshot query to find an identifier in the file.
+
+        Args:
+            query (str): The tree-sitter oneshot query to execute.
+
+        Returns:
+            Identifier | None: The first identifier that matches the query, or None if no match is found.
+        """
+        matched_identifiers = self.query_identifiers(query, node)
+        if len(matched_identifiers) == 0:
+            return None
+        return matched_identifiers[0]

scubatrace/function.py

@@ -10,6 +10,7 @@
 
 from . import language as lang
 from .call import Call
+from .identifier import Identifier
 from .statement import BlockStatement, Statement
 
 if TYPE_CHECKING:
@@ -163,12 +164,20 @@ def parameter_lines(self) -> list[int]:
         """
         The lines where the parameters of the function are defined.
         """
-        parameters_node = self.node.child_by_field_name("parameters")
-        if parameters_node is None:
+        params = self.query_identifiers(self.language.query_function_parameter)
+        if len(params) == 0:
             return [self.start_line]
-        param_start_line = parameters_node.start_point[0] + 1
-        param_end_line = parameters_node.end_point[0] + 1
-        return list(range(param_start_line, param_end_line + 1))
+        return list(range(params[0].start_line, params[-1].end_line + 1))
+
+    @cached_property
+    def parameters(self) -> list[Identifier]:
+        """
+        The parameter statements of the function.
+        """
+        params = self.query_identifiers(self.language.query_function_parameter)
+        if len(params) == 0:
+            return self.block_variables
+        return params
 
     @cached_property
     def name_node(self) -> Node:

scubatrace/java/language.py

@@ -19,6 +19,11 @@ class JAVA(Language):
             (identifier)@name
         )
     """
+    query_function_parameter = """
+        (formal_parameter
+        	name: (identifier)@name
+        )
+    """
 
     query_package = "(package_declaration)@name"
     query_class = "(class_declaration)@name"

scubatrace/language.py

@@ -41,6 +41,11 @@ class Language:
     For example, in C/C++, this would match the `header.h` in `#include <header.h>`.
     """
 
+    query_function_parameter: str
+    """
+    The tree-sitter query to match function parameters.
+    """
+
     EXIT_STATEMENTS: list[str] = []
     """
     The tree-sitter AST types of exit statements.

scubatrace/python/language.py

@@ -27,6 +27,14 @@ class PYTHON(Language):
         ]
     )
     """
+    query_function_parameter = """
+        (parameters
+            (identifier)@name
+        )
+        (typed_parameter
+            (identifier)@name
+        )
+    """
 
     query_class = "(class_definition)@name"
 

scubatrace/statement.py

@@ -1024,3 +1024,30 @@ def query_oneshot(self, query: str) -> Statement | None:
         if len(matched_statements) == 0:
             return None
         return matched_statements[0]
+
+    def query_identifiers(self, query: str) -> list[Identifier]:
+        """
+        Executes a tree-sitter query to find identifiers in the block.
+
+        Args:
+            query (str): The tree-sitter query to execute.
+
+        Returns:
+            list[Identifier]: A list of identifiers that match the query.
+        """
+        return self.file.query_identifiers(query, self.node)
+
+    def query_identifier(self, query: str) -> Identifier | None:
+        """
+        Executes a tree-sitter oneshot query to find an identifier in the block.
+
+        Args:
+            query (str): The tree-sitter oneshot query to execute.
+
+        Returns:
+            Identifier | None: The identifier that matches the query, or None if not found.
+        """
+        matched_identifiers = self.query_identifiers(query)
+        if len(matched_identifiers) == 0:
+            return None
+        return matched_identifiers[0]