Add retrieving types in codequery (#260)

* Add retrieving types in codequery

* Add fuzzy search option (#271)

Author: evandowning

program-model/src/buttercup/program_model/api/tree_sitter.py

@@ -4,10 +4,13 @@
 from dataclasses import dataclass
 from pathlib import Path
 from functools import lru_cache
-from enum import Enum
-
 from buttercup.common.challenge_task import ChallengeTask
-from buttercup.program_model.utils.common import Function, FunctionBody
+from buttercup.program_model.utils.common import (
+    Function,
+    FunctionBody,
+    TypeDefinition,
+    TypeDefinitionType,
+)
 from tree_sitter_language_pack import get_language, get_parser
 from buttercup.common.project_yaml import ProjectYaml
 
@@ -81,26 +84,6 @@
 """
 
 
-@dataclass
-class TypeDefinitionType(str, Enum):
-    """Enum to store type definition type."""
-
-    STRUCT = "struct"
-    UNION = "union"
-    ENUM = "enum"
-    TYPEDEF = "typedef"
-    PREPROC_TYPE = "preproc_type"
-
-
-@dataclass
-class TypeDefinition:
-    """Class to store type definition information."""
-
-    name: str
-    type: TypeDefinitionType
-    definition: str
-
-
 @dataclass
 class CodeTS:
     """Class to extract information about functions in a challenge project using TreeSitter."""
@@ -199,7 +182,9 @@ def get_function(self, function_name: str, file_path: Path) -> Function | None:
         functions = self.get_functions(file_path)
         return functions.get(function_name)
 
-    def parse_types_in_code(self, file_path: Path) -> dict[str, TypeDefinition]:
+    def parse_types_in_code(
+        self, file_path: Path, typename: str | None = None, fuzzy: bool | None = False
+    ) -> dict[str, TypeDefinition]:
         """Parse the definition of a type in a piece of code."""
         logger.debug("Parsing types in code")
         code = self.challenge_task.task_dir.joinpath(file_path).read_bytes()
@@ -233,6 +218,10 @@ def parse_types_in_code(self, file_path: Path) -> dict[str, TypeDefinition]:
 
             type_definition = code[start_byte : definition_node.end_byte].decode()
             name = name_node.text.decode()
+            if typename and not fuzzy and name != typename:
+                continue
+            if typename and fuzzy and typename not in name:
+                continue
             logger.debug("Type name: %s", name)
             logger.debug("Type definition: %s", type_definition)
 
@@ -255,6 +244,7 @@ def parse_types_in_code(self, file_path: Path) -> dict[str, TypeDefinition]:
                 name=name,
                 type=type_def_type,
                 definition=type_definition,
+                definition_line=definition_node.start_point[0],
             )
 
         return res

program-model/src/buttercup/program_model/codequery.py

@@ -8,11 +8,15 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 from itertools import groupby
-from typing import ClassVar
+from typing import ClassVar, Union
+
 
 from buttercup.common.challenge_task import ChallengeTask
 from buttercup.program_model.api.tree_sitter import CodeTS
-from buttercup.program_model.utils.common import Function
+from buttercup.program_model.utils.common import (
+    Function,
+    TypeDefinition,
+)
 from buttercup.common.project_yaml import ProjectYaml
 
 logger = logging.getLogger(__name__)
@@ -215,7 +219,10 @@ def _run_cqsearch(self, *args: str) -> list[CQSearchResult]:
         return [result for result in results if result is not None]
 
     def get_functions(
-        self, function_name: str, file_path: Path | None = None
+        self,
+        function_name: str,
+        file_path: Path | None = None,
+        fuzzy: bool | None = False,
     ) -> list[Function]:
         """Get the definition(s) of a function in the codebase or in a specific file."""
         cqsearch_args = [
@@ -225,31 +232,117 @@ def get_functions(
             "2",
             "-t",
             function_name,
-            "-e",
+            "-f" if fuzzy else "-e",
             "-u",
         ]
         if file_path:
             cqsearch_args += ["-b", file_path.as_posix()]
 
         results = self._run_cqsearch(*cqsearch_args)
+        logger.debug("cqsearch output: %s", results)
 
-        res = []
+        res: list[Function] = []
         results_by_file = groupby(results, key=lambda x: x.file)
         for file, results in results_by_file:
-            if not all(result.value == function_name for result in results):
+            functions_found = [result.value for result in results]
+
+            if not fuzzy and not all(function_name == f for f in functions_found):
                 logger.warning(
-                    "Function name mismatch, this should not happen: %s != %s",
-                    results[0].value,
+                    "Function name mismatch, this should not happen: %s",
                     function_name,
                 )
                 continue
+            if fuzzy and not all(function_name in f for f in functions_found):
+                logger.warning(
+                    "Function name mismatch, this should not happen: %s",
+                    function_name,
+                )
+                continue
+
+            for function in functions_found:
+                f = self.ts.get_function(function, file)
+                if f is None:
+                    logger.warning("Function not found in tree-sitter: %s", function)
+                    continue
+                res.append(f)
 
-            function = self.ts.get_function(function_name, file)
-            if function is None:
-                logger.warning("Function not found in tree-sitter: %s", function_name)
+        return res
+
+    def get_types(
+        self,
+        type_name: Union[bytes, str],
+        file_path: Path | None = None,
+        function_name: str | None = None,
+        fuzzy: bool | None = False,
+    ) -> list[TypeDefinition]:
+        """Finds and return the definition of type named `typename`."""
+        # Build the cqsearch command to find occurences of the typename in the code
+        cqsearch_args = [
+            "-s",
+            self.CODEQUERY_DB,  # Specify the database file path
+            "-p",
+            "1",  # '1' for symbol
+            "-t",
+            type_name,  # The name of the type
+            "-f" if fuzzy else "-e",
+            "-u",  # use full paths
+        ]
+        if file_path:
+            cqsearch_args += ["-b", file_path.as_posix()]
+
+        results = self._run_cqsearch(*cqsearch_args)
+        logger.debug("cqsearch output: %s", results)
+
+        res: list[TypeDefinition] = []
+        results_by_file = groupby(results, key=lambda x: x.file)
+        for file, results in results_by_file:
+            types_found = [result.value for result in results]
+
+            if not fuzzy and not all(type_name == t for t in types_found):
+                logger.warning(
+                    "Type name mismatch, this should not happen: %s",
+                    type_name,
+                )
+                continue
+            if fuzzy and not all(type_name in t for t in types_found):
+                logger.warning(
+                    "Type name mismatch, this should not happen: %s",
+                    type_name,
+                )
                 continue
 
-            res.append(function)
+            typedefs: dict[str, TypeDefinition] = {}
+
+            for typename in types_found:
+                t = self.ts.parse_types_in_code(file, typename, fuzzy)
+                if not t:
+                    logger.warning(
+                        "Type definition not found in tree-sitter: %s", typename
+                    )
+                    continue
+                typedefs.update(t)
+
+            if function_name:
+                # Get the function definition to find its scope
+                function = self.ts.get_function(function_name, file)
+                if function:
+                    # Filter type definitions to only include those within the function's scope
+                    filtered_typedefs = {}
+                    for name, typedef in typedefs.items():
+                        # Check if the type definition is within the function's scope
+                        for body in function.bodies:
+                            if (
+                                body.start_line
+                                <= typedef.definition_line
+                                <= body.end_line
+                            ):
+                                filtered_typedefs[name] = typedef
+                                break
+                    typedefs = filtered_typedefs
+                else:
+                    typedefs = {}
+
+            res.extend(typedefs.values())
 
         return res
 

program-model/src/buttercup/program_model/utils/common.py

@@ -1,5 +1,6 @@
 import logging
 from dataclasses import dataclass, field
+from enum import Enum
 from pathlib import Path
 
 
@@ -35,3 +36,31 @@ class Function:
 
     bodies: list[FunctionBody] = field(default_factory=list)
     """List of function bodies."""
+
+
+@dataclass
+class TypeDefinitionType(str, Enum):
+    """Enum to store type definition type."""
+
+    STRUCT = "struct"
+    UNION = "union"
+    ENUM = "enum"
+    TYPEDEF = "typedef"
+    PREPROC_TYPE = "preproc_type"
+
+
+@dataclass
+class TypeDefinition:
+    """Class to store type definition information."""
+
+    name: str
+    """Name of the type."""
+
+    type: TypeDefinitionType
+    """Type of the type."""
+
+    definition: str
+    """Definition of the type."""
+
+    definition_line: int
+    """Line number of the definition of the type."""

program-model/tests/test_codequery.py

@@ -7,6 +7,7 @@
 from buttercup.common.challenge_task import ChallengeTask
 from buttercup.program_model.codequery import CodeQuery, CodeQueryPersistent
 from buttercup.common.task_meta import TaskMeta
+from buttercup.program_model.utils.common import TypeDefinitionType
 
 
 def setup_dirs(tmp_path: Path) -> Path:
@@ -49,6 +50,13 @@ def setup_dirs(tmp_path: Path) -> Path:
 int function4(char *s) {
     return strlen(s);
 }
+""")
+    (source / "test4.c").write_text("""typedef int myInt;
+myInt function5(myInt a, myInt b) {
+    typedef int myOtherInt;
+    myOtherInt c = a + b;
+    return a + b + c;
+}
 """)
 
     # Create task metadata
@@ -123,6 +131,17 @@ def test_get_functions_multiple(mock_challenge_task: ChallengeTask):
     )
 
 
+def test_get_functions_fuzzy(mock_challenge_task: ChallengeTask):
+    """Test that we can get functions (fuzzy search) in codebase"""
+    codequery = CodeQuery(mock_challenge_task)
+    functions = codequery.get_functions("function", fuzzy=True)
+    assert len(functions) == 4
+    functions = codequery.get_functions("function", Path("test3.c"), fuzzy=True)
+    assert len(functions) == 2
+    functions = codequery.get_functions("function3", Path("test3.c"), fuzzy=True)
+    assert len(functions) == 1
+
+
 def test_keep_status(
     mock_challenge_task: ChallengeTask,
     mock_challenge_task_ro: ChallengeTask,
@@ -161,6 +180,46 @@ def test_keep_status(
         assert mock_challenge_task_ro.task_dir.exists()
 
 
+def test_get_types(mock_challenge_task: ChallengeTask):
+    """Test that we can get types in codebase"""
+    codequery = CodeQuery(mock_challenge_task)
+    types = codequery.get_types("myInt", Path("test3.c"))
+    assert len(types) == 0
+    types = codequery.get_types("myInt")
+    assert len(types) == 1
+    types = codequery.get_types("myInt", Path("test4.c"))
+    assert len(types) == 1
+    assert types[0].name == "myInt"
+    assert types[0].type == TypeDefinitionType.TYPEDEF
+    assert types[0].definition == "typedef int myInt;"
+    assert types[0].definition_line == 0
+    types = codequery.get_types("myInt", Path("test4.c"), function_name="function5")
+    assert len(types) == 0
+    types = codequery.get_types(
+        "myOtherInt", Path("test4.c"), function_name="function5"
+    )
+    assert len(types) == 1
+    assert types[0].name == "myOtherInt"
+    assert types[0].type == TypeDefinitionType.TYPEDEF
+    assert types[0].definition == "    typedef int myOtherInt;"
+    assert types[0].definition_line == 2
+
+
+def test_get_types_fuzzy(mock_challenge_task: ChallengeTask):
+    """Test that we can get types (fuzzy search) in codebase"""
+    codequery = CodeQuery(mock_challenge_task)
+    types = codequery.get_types("my", Path("test4.c"), fuzzy=True)
+    assert len(types) == 2
+    types = codequery.get_types("myInt", Path("test4.c"), fuzzy=True)
+    assert len(types) == 1
+    types = codequery.get_types("myOtherInt", Path("test4.c"), fuzzy=True)
+    assert len(types) == 1
+    types = codequery.get_types("my", fuzzy=True)
+    assert len(types) == 2
+    types = codequery.get_types("myOtherInt", Path("test4.c"), "function5", fuzzy=True)
+    assert len(types) == 1
+
+
 @pytest.fixture
 def libjpeg_oss_fuzz_task(tmp_path: Path) -> ChallengeTask:
     """Create a challenge task using a real OSS-Fuzz repository."""