refactor(go-segmenter): replace custom GoSegmenter with native Tree-Sitter implementation

Signed-off-by: Vladimir Belousov <[email protected]>

Author: vbelouso

src/vuln_analysis/tools/tests/test_go_segmenter.py

@@ -0,0 +1,72 @@
+from vuln_analysis.utils.go_segmenter_extended import GoSegmenterExtended
+
+
+def _extract(code: str):
+    seg = GoSegmenterExtended(code)
+    return [s.strip() for s in seg.extract_functions_classes()]
+
+
+def test_generic_method_basic():
+    code = """
+    type Box[T any] struct { value T }
+    func (b *Box[T]) Set(v T) { b.value = v }
+    """
+    chunks = _extract(code)
+    assert any("Set" in c for c in chunks), "generic method not extracted"
+
+
+def test_generic_multiple_type_params():
+    code = """
+    func MapKeys[K comparable, V any](m map[K]V) []K {
+        keys := make([]K, 0, len(m))
+        for k := range m {
+            keys = append(keys, k)
+        }
+        return keys
+    }
+    """
+    chunks = _extract(code)
+    assert any("MapKeys" in c for c in chunks), "multiple generics not parsed"
+
+
+def test_function_returning_func():
+    code = """
+    func makeAdder(x int) func(int) int {
+        return func(y int) int { return x + y }
+    }
+    """
+    chunks = _extract(code)
+    assert any("makeAdder" in c for c in chunks), "failed to parse func returning func"
+
+
+def test_inline_anonymous_func():
+    code = """
+    func Worker() {
+        defer func() { cleanup() }()
+        go func() { runTask() }()
+    }
+    """
+    chunks = _extract(code)
+    assert any("Worker" in c for c in chunks), "missed inline anonymous func"
+
+
+def test_double_pointer_receiver():
+    code = """
+    type Conn struct{}
+    func (c **Conn) Reset() {}
+    """
+    chunks = _extract(code)
+    assert any("Reset" in c for c in chunks), "failed to detect pointer receiver"
+
+
+def test_multiline_generic_method():
+    code = """
+    func (r *Repo[
+        T any,
+        E error,
+    ]) Save(v T) (E, error) {
+        return nil, nil
+    }
+    """
+    chunks = _extract(code)
+    assert any("Save" in c for c in chunks), "multiline generic method not parsed"

src/vuln_analysis/tools/tests/test_transitive_code_search.py

@@ -86,7 +86,7 @@ async def test_transitive_search_golang_5():
     (path_found, list_path) = result
 
     assert path_found is False
-    assert len(list_path) is 1
+    assert len(list_path) == 1
 
 # Test fix of https://issues.redhat.com/browse/APPENG-3435
 @pytest.mark.asyncio
@@ -100,7 +100,7 @@ async def test_transitive_search_golang_6():
     (path_found, list_path) = result
     print(result)
     assert path_found is False
-    assert len(list_path) is 0
+    assert len(list_path) == 0
 
 
 def set_input_for_next_run(git_repository: str, git_ref: str, included_extensions: list[str],

src/vuln_analysis/utils/chain_of_calls_retriever.py

@@ -190,14 +190,15 @@ def __init__(self, documents: List[Document], ecosystem: Ecosystem, manifest_pat
             logger.debug(f"no_macro_documents len : {len(no_macro_documents)}")  
 
         if not self.language_parser.is_script_language():            
-            # filter out types and full code documents, retaining only functions/methods documents in this attribute.
+        # filter out types and full code documents, retaining only functions/methods documents in this attribute.
             self.documents = [doc for doc in no_macro_documents if self.language_parser.is_function(doc)]
             self.documents_of_functions = self.documents
         else:
             self.documents = filtered_documents
             self.documents_of_functions = [doc for doc in self.documents
                           if doc.page_content.startswith(self.language_parser.get_function_reserved_word())]
-        
+        # sort documents to ensure deterministic behavior
+        self.documents.sort(key=lambda doc: doc.metadata.get('source', ''))
         logger.debug(f"self.documents  len : {len(self.documents)}")        
         logger.debug("Chain of Calls Retriever - retaining only types/classes docs "
                     "documents_of_types len %d", len(self.documents_of_types))
@@ -248,6 +249,8 @@ def __find_caller_function_dfs(self, document_function: Document, function_packa
             # Add same package itself to search path.
         # direct_parents.extend([function_package])
         # gets list of documents to search in only from parents of function' package.
+        # Fixes non-deterministic behavior in chain-of-calls resolution where identical inputs produced different call-chain lengths
+        direct_parents.sort()
         function_name_to_search = self.language_parser.get_function_name(document_function)
         if function_name_to_search == self.language_parser.get_constructor_method_name():
             function_name_to_search = self.language_parser.get_class_name_from_class_function(document_function)

src/vuln_analysis/utils/document_embedding.py

@@ -15,35 +15,37 @@
 
 import copy
 import json
-import logging
 import os
 import pickle
 import sys
 import time
 import typing
 from hashlib import sha512
-from pathlib import Path
-from pathlib import PurePath
+from pathlib import Path, PurePath
 
 from langchain.docstore.document import Document
-from langchain.text_splitter import Language
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
 from langchain_community.document_loaders.generic import GenericLoader
-from langchain_community.document_loaders.parsers.language.code_segmenter import CodeSegmenter
-from langchain_community.document_loaders.parsers.language.language_parser import LANGUAGE_EXTENSIONS
-from langchain_community.document_loaders.parsers.language.language_parser import LANGUAGE_SEGMENTERS
-from langchain_community.document_loaders.parsers.language.language_parser import LanguageParser
+from langchain_community.document_loaders.parsers.language.code_segmenter import (
+    CodeSegmenter,
+)
+from langchain_community.document_loaders.parsers.language.language_parser import (
+    LANGUAGE_EXTENSIONS,
+    LANGUAGE_SEGMENTERS,
+    LanguageParser,
+)
 from langchain_community.vectorstores import FAISS
 from langchain_core.document_loaders.blob_loaders import Blob
 
 from vuln_analysis.data_models.input import SourceDocumentsInfo
-from vuln_analysis.utils.go_segmenters_with_methods import GoSegmenterWithMethods
-from vuln_analysis.utils.python_segmenters_with_classes_methods import PythonSegmenterWithClassesMethods
+from vuln_analysis.logging.loggers_factory import LoggingFactory
+from vuln_analysis.utils.git_utils import sanitize_git_url_for_path
+from vuln_analysis.utils.go_segmenter_extended import GoSegmenterExtended
 from vuln_analysis.utils.js_extended_parser import ExtendedJavaScriptSegmenter
+from vuln_analysis.utils.python_segmenters_with_classes_methods import (
+    PythonSegmenterWithClassesMethods,
+)
 from vuln_analysis.utils.source_code_git_loader import SourceCodeGitLoader
-from vuln_analysis.utils.git_utils import sanitize_git_url_for_path
-from vuln_analysis.utils.transitive_code_searcher_tool import TransitiveCodeSearcher
-from vuln_analysis.logging.loggers_factory import LoggingFactory
 
 from vuln_analysis.utils.c_segmenter_custom import CSegmenterExtended
 
@@ -144,7 +146,8 @@ class ExtendedLanguageParser(LanguageParser):
             "javascript": ExtendedJavaScriptSegmenter,
             "js": ExtendedJavaScriptSegmenter,
         }
-    additional_segmenters["go"] = GoSegmenterWithMethods
+    # additional_segmenters["go"] = GoSegmenterWithMethods
+    additional_segmenters["go"] = GoSegmenterExtended
     additional_segmenters["python"] = PythonSegmenterWithClassesMethods
     additional_segmenters["c"] = CSegmenterExtended