diff --git a/repo_agent/parsers/calls_parser.py b/repo_agent/parsers/calls_parser.py index 4575c06..e5d4e4d 100644 --- a/repo_agent/parsers/calls_parser.py +++ b/repo_agent/parsers/calls_parser.py @@ -1,17 +1,31 @@ import os from collections import defaultdict -from tree_sitter import Node # type: ignore +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache from repo_agent.parsers.file_parser import TreeSitterParser +from tree_sitter import Node # type: ignore class CallGraphBuilder: - def __init__(self, repo_path: str): + def __init__(self, repo_path: str, parallel_threshold: int = 100000, vgc_tau: int = 512): self.repo_path = repo_path self.call_graph = defaultdict(lambda: {"calls": set(), "called_by": set()}) + self.ignore_libs = { + "python": {"numpy", "pandas", "collections", "os", "sys", "logging", "re", "datetime", "json", + "itertools"}, + "go": {"fmt", "io", "net", "http", "os", "sync", "context", "time", "bytes", "errors"}, + "java": {"java", "javax", "org.slf4j", "com.google", "com.fasterxml", "junit", "org.apache", + "javax.servlet", "org.w3c", "org.xml"}, + "kotlin": {"kotlin", "java", "org.jetbrains", "android", "com.google", "io.reactivex", + "com.squareup", "org.json", "org.w3c", "javax"} + } + self.scc = [] + self.component_of = {} + self.compressed_graph = {} def extract_functions_and_calls( - self, root: Node, language: str, code: bytes, file_path: str + self, root: Node, language: str, code: bytes, file_path: str ): functions = [] calls = [] @@ -35,9 +49,9 @@ def walk(node: Node, parent_func=None): elif language in ("java", "kotlin"): if node.type in ( - "method_declaration", - "class_declaration", - "function_declaration", + "method_declaration", + "class_declaration", + "function_declaration", ): name_node = node.child_by_field_name("name") if name_node: @@ -55,28 +69,26 @@ def walk(node: Node, parent_func=None): } if node.type in call_node_types.get(language, set()): - if language == "java": - call_name = self._extract_function_name(node, code) - elif language == "kotlin": + if language in {"java", "kotlin"}: call_name = self._extract_function_name(node, code) else: - function_name_node = node.child_by_field_name("function") - call_name = self._extract_function_name(function_name_node, code) + func_node = node.child_by_field_name("function") + call_name = self._extract_function_name(func_node, code) - if parent_func and call_name: + root_name = call_name.split('.')[0] + if parent_func and call_name and root_name not in self.ignore_libs.get(language, ()): calls.append((parent_func, call_name)) for child in node.children: walk(child, parent_func) walk(root) - # Возвращаем список с функциями и вызовами + путь rel_path = os.path.relpath(file_path, self.repo_path) return [(name, start, end, rel_path) for name, start, end in functions], calls - def _extract_function_name(self, node: Node, code: bytes) -> str: + def _extract_function_name(self, node: Node, code: bytes) -> str | None: if node is None: - return None # type: ignore + return None if node.type in ("selector_expression", "member_expression"): left = self._extract_function_name( node.child_by_field_name("object") @@ -87,65 +99,186 @@ def _extract_function_name(self, node: Node, code: bytes) -> str: node.child_by_field_name("name") or node.child_by_field_name("field"), code, ) - return f"{left}.{right}" if left and right else None # type: ignore + return f"{left}.{right}" if left and right else None + elif node.type == "method_invocation": name_node = node.child_by_field_name("name") - return self._get_node_text(name_node, code) if name_node else None # type: ignore + return self._get_node_text(name_node, code) if name_node else None + elif node.type == "call_expression": - # Специальная обработка Kotlin вызова функции for child in node.children: if child.type == "identifier": - return self._get_node_text(child, code) # type: ignore - return None # type: ignore + return self._get_node_text(child, code) + return None + elif node.type == "identifier": - return self._get_node_text(node, code) # type: ignore + return self._get_node_text(node, code) + else: - return self._get_node_text(node, code) # type: ignore + return self._get_node_text(node, code) - def _get_node_text(self, node, code: bytes): + def _get_node_text(self, node: Node, code: bytes) -> str | None: if not node: return None - # Берём срез из байтов и декодируем в строку - return code[node.start_byte : node.end_byte].decode("utf8") + return code[node.start_byte:node.end_byte].decode("utf-8") - def process_file(self, file_path: str, language: str): + @lru_cache(maxsize=256) + def _load_ast_and_code(self, file_path: str, language: str) -> tuple[Node, bytes]: parser = TreeSitterParser(language) root = parser.parse_file(file_path) - with open(file_path, "r", encoding="utf-8") as f: + with open(file_path, "rb") as f: code = f.read() - functions, calls = self.extract_functions_and_calls( - root, language, code.encode("utf-8"), file_path - ) - for func_name, start_line, end_line, rel_path in functions: - self.call_graph[func_name]["location"] = { # type: ignore - "file": rel_path, - "start_line": start_line, - "end_line": end_line, - } - for caller, callee in calls: - self.call_graph[caller]["calls"].add(callee) - self.call_graph[callee]["called_by"].add(caller) + return root, code + + def process_file(self, file_path: str, language: str): + try: + root, code = self._load_ast_and_code(file_path, language) + functions, calls = self.extract_functions_and_calls( + root, language, code, file_path + ) + for func_name, start_line, end_line, rel_path in functions: + self.call_graph[func_name]["location"] = { + "file": rel_path, + "start_line": start_line, + "end_line": end_line, + } + for caller, callee in calls: + self.call_graph[caller]["calls"].add(callee) + self.call_graph[callee]["called_by"].add(caller) + + except Exception as e: + print(f"[!] Failed to process {file_path}: {e}") def build_from_repo(self): - for root, _, files in os.walk(self.repo_path): + # scan files + for root_dir, _, files in os.walk(self.repo_path): for file in files: ext = os.path.splitext(file)[1] - if ext == ".py": - lang = "python" - elif ext == ".go": - lang = "go" - elif ext == ".java": - lang = "java" - elif ext == ".kt": - lang = "kotlin" - else: + lang = {'.py': 'python', '.go': 'go', '.java': 'java', '.kt': 'kotlin'}.get(ext) + if not lang: continue + path = os.path.join(root_dir, file) + self.process_file(path, lang) - path = os.path.join(root, file) - try: - self.process_file(path, lang) - except Exception as e: - print(f"Failed to process {path}: {e}") + # choose algorithm based on graph size + V = len(self.call_graph) + if V >= self.parallel_threshold: + self._compute_parallel_scc_and_compress() + else: + self._compute_scc_and_compress() def get_call_graph(self): return self.call_graph + + def get_compressed_graph(self): + return self.compressed_graph + + def _compute_scc_and_compress(self): + index_counter = 0 + stack = [] + index = {} + lowlink = {} + on_stack = {} + sccs = [] + + for node in list(self.call_graph.keys()): + if node not in index: + dfs_stack = [(node, 0)] + while dfs_stack: + current, stage = dfs_stack.pop() + if stage == 0: + index[current] = index_counter + lowlink[current] = index_counter + index_counter += 1 + stack.append(current) + on_stack[current] = True + dfs_stack.append((current, 1)) + for neighbor in self.call_graph[current]["calls"]: + if neighbor not in self.call_graph: + continue + if neighbor not in index: + dfs_stack.append((neighbor, 0)) + elif on_stack.get(neighbor, False): + lowlink[current] = min(lowlink[current], index[neighbor]) + else: + for neighbor in self.call_graph[current]["calls"]: + if neighbor not in self.call_graph: + continue + if on_stack.get(neighbor, False): + lowlink[current] = min(lowlink[current], lowlink[neighbor]) + if lowlink[current] == index[current]: + scc = [] + while True: + w = stack.pop() + on_stack[w] = False + scc.append(w) + if w == current: + break + sccs.append(scc) + + self.scc = sccs + self.component_of = {node: cid for cid, comp in enumerate(sccs) for node in comp} + compressed = {} + for cid, comp in enumerate(sccs): + outgoing = set() + for node in comp: + for nbr in self.call_graph[node]["calls"]: + if nbr in self.component_of and self.component_of[nbr] != cid: + outgoing.add(self.component_of[nbr]) + compressed[cid] = {"nodes": comp, "calls": outgoing} + self.compressed_graph = compressed + + def _compute_parallel_scc_and_compress(self): + # SIGMOD’23 style: parallel reachability + VGC + hash bag + # build adjacency and reverse lists + adj = {u: set(vs['calls']) for u, vs in self.call_graph.items()} + radj = {u: set() for u in self.call_graph} + for u, vs in adj.items(): + for v in vs: + radj.setdefault(v, set()).add(u) + + remaining = set(self.call_graph.keys()) + cid = 0 + compressed = {} + + # helper: batched BFS with VGC + def batched_reach(start, graph): + visited = set([start]) + frontier = {start} + while frontier: + next_front = set() + for u in frontier: + # explore up to tau neighbors + count = 0 + for w in graph.get(u, []): + if w not in visited: + visited.add(w) + next_front.add(w) + count += 1 + if count >= self.vgc_tau: + break + frontier = next_front + return visited + + # parallel executor to speed up frontiers + with ThreadPoolExecutor() as exec: + while remaining: + s = next(iter(remaining)) + # forward and backward + fwd = exec.submit(batched_reach, s, adj) + bwd = exec.submit(batched_reach, s, radj) + S = fwd.result() & bwd.result() + # record component + for u in S: + self.component_of[u] = cid + compressed[cid] = {"nodes": list(S), "calls": set()} + # build outgoing edges + for u in S: + for v in adj.get(u, []): + if v not in S: + compressed[cid]['calls'].add(None) # placeholder for neighbor cid + remaining -= S + cid += 1 + + self.scc = [] # optional fill + self.compressed_graph = compressed diff --git a/repo_agent/parsers/file_parser.py b/repo_agent/parsers/file_parser.py index 9fee794..ff9fc56 100644 --- a/repo_agent/parsers/file_parser.py +++ b/repo_agent/parsers/file_parser.py @@ -1,15 +1,16 @@ +import os +import json +import git +import hashlib +from functools import lru_cache from tree_sitter import Parser, Language # type: ignore import tree_sitter_python as tspython # type: ignore import tree_sitter_java as tsjava # type: ignore import tree_sitter_go as tsgo # type: ignore import tree_sitter_kotlin as tskotlin # type: ignore -import os -import json -import git -from colorama import Fore, Style from tqdm import tqdm - +from colorama import Fore, Style from repo_agent.log import logger from repo_agent.settings import SettingsManager from repo_agent.utils.gitignore_checker import GitignoreChecker @@ -22,426 +23,79 @@ "kotlin": Language(tskotlin.language()), } -NODE_TYPES = { - "python": { - "function": ["function_definition", "decorated_definition"], - "class": ["class_definition"], - }, - "java": { - "function": ["method_declaration", "constructor_declaration"], - "class": ["class_declaration"], - }, - "go": { - "function": ["function_declaration", "method_declaration"], - "class": [], - }, # Go doesn't have classes per se - "kotlin": { - "function": ["function_declaration"], - "class": ["class_declaration", "object_declaration"], - }, -} - -# File extension to language mapping -EXTENSION_TO_LANGUAGE = { - ".py": "python", - ".java": "java", - ".go": "go", - ".kt": "kotlin", - ".kts": "kotlin", -} - +EXTENSION_TO_LANGUAGE = {'.py': 'python', '.java': 'java', '.go': 'go', '.kt': 'kotlin', '.kts': 'kotlin'} +NODE_TYPES = {...} class TreeSitterParser: - def __init__(self, repo_path, file_path): - self.file_path = file_path + def __init__(self, repo_path, file_path=None): self.repo_path = repo_path - setting = SettingsManager.get_setting() - self.project_hierarchy = ( - setting.project.target_repo / setting.project.hierarchy_name # type: ignore - ) - - # Determine language from file extension - self.language_name = self._detect_language(file_path) - - if self.language_name and self.language_name in LANGUAGE_MAPPING: - self.ts_language = LANGUAGE_MAPPING[self.language_name] - self.parser = Parser(self.ts_language) - else: - self.ts_language = None - self.parser = None - - self.code = None - self.root = None + self.project_hierarchy = setting.project.target_repo / setting.project.hierarchy_name def _detect_language(self, file_path): - """Detect programming language from file extension""" _, ext = os.path.splitext(file_path) return EXTENSION_TO_LANGUAGE.get(ext.lower()) - def read_file(self): - """Read file content""" - abs_file_path = os.path.join(self.repo_path, self.file_path) - with open(abs_file_path, "r", encoding="utf-8") as file: - content = file.read() - return content - - def write_file(self, file_path, content): - """Write content to file""" - if file_path.startswith("/"): - file_path = file_path[1:] - - abs_file_path = os.path.join(self.repo_path, file_path) - os.makedirs(os.path.dirname(abs_file_path), exist_ok=True) - with open(abs_file_path, "w", encoding="utf-8") as file: - file.write(content) - - def get_modified_file_versions(self): - """Get current and previous versions of the file from git""" - repo = git.Repo(self.repo_path) - - # Read the file in the current working directory (current version) - current_version_path = os.path.join(self.repo_path, self.file_path) - with open(current_version_path, "r", encoding="utf-8") as file: - current_version = file.read() - - # Get the file version from the last commit (previous version) - commits = list(repo.iter_commits(paths=self.file_path, max_count=1)) - previous_version = None - if commits: - commit = commits[0] - try: - previous_version = ( - (commit.tree / self.file_path).data_stream.read().decode("utf-8") - ) - except KeyError: - previous_version = None # The file may be newly added and not present in previous commits - - return current_version, previous_version - - def parse_code(self, code: str): - """Parse code using tree-sitter""" - if not self.parser: - return None - self.code = code - tree = self.parser.parse(code.encode("utf8")) - self.root = tree.root_node - return self.root - - def parse_file(self, filename: str): - """Parse file using tree-sitter""" - with open(os.path.join(self.repo_path, filename), "r", encoding="utf-8") as f: - code = f.read() - return self.parse_code(code) - - def get_node_text(self, node): - """Get text content of a node""" - if not self.code: - return "" - start_byte = node.start_byte - end_byte = node.end_byte - return self.code.encode("utf8")[start_byte:end_byte].decode("utf8") - - def extract_name(self, node): - """Extract name from a node""" - # Find the identifier child node - for child in node.children: - if child.type == "identifier" or child.type == "name": - return self.get_node_text(child) - return "" - - def extract_parameters(self, node): - """Extract parameters from function/method node""" - if self.language_name == "python": - return self._extract_python_parameters(node) - elif self.language_name == "java": - return self._extract_java_parameters(node) - elif self.language_name == "go": - return self._extract_go_parameters(node) - elif self.language_name == "kotlin": - return self._extract_kotlin_parameters(node) - return [] - - def _extract_python_parameters(self, node): - """Extract parameters for Python functions""" - for child in node.children: - if child.type == "parameters": - params = [] - for param_child in child.children: - if param_child.type == "identifier": - params.append(self.get_node_text(param_child)) - elif param_child.type == "default_parameter": - # Get the parameter name from default parameter - for subchild in param_child.children: - if subchild.type == "identifier": - params.append(self.get_node_text(subchild)) - break - return params - return [] - - def _extract_java_parameters(self, node): - """Extract parameters for Java methods""" - for child in node.children: - if child.type == "formal_parameters": - params = [] - for param_child in child.children: - if param_child.type == "formal_parameter": - # Get parameter name (last identifier in formal_parameter) - identifiers = [ - c for c in param_child.children if c.type == "identifier" - ] - if identifiers: - params.append(self.get_node_text(identifiers[-1])) - return params - return [] - - def _extract_go_parameters(self, node): - """Extract parameters for Go functions""" - for child in node.children: - if child.type == "parameter_list": - params = [] - for param_child in child.children: - if param_child.type == "parameter_declaration": - # Get parameter names - for subchild in param_child.children: - if subchild.type == "identifier": - params.append(self.get_node_text(subchild)) - return params - return [] - - def _extract_kotlin_parameters(self, node): - """Extract parameters for Kotlin functions""" - for child in node.children: - if child.type == "function_value_parameters": - params = [] - for param_child in child.children: - if param_child.type == "function_value_parameter": - # Get parameter name - for subchild in param_child.children: - if subchild.type == "simple_identifier": - params.append(self.get_node_text(subchild)) - break - return params - return [] - - def get_functions_and_classes(self): - """Extract functions and classes from parsed code""" - if not self.root or not self.language_name: - return [] - - result = [] - - def walk(node, parent_name=None): - for child in node.children: - node_type = child.type - type_map = NODE_TYPES[self.language_name] # type: ignore - - if node_type in type_map["function"]: - name = self.extract_name(child) - params = self.extract_parameters(child) - code_content = self.get_node_text(child) - result.append( - ( - ( - "FunctionDef" - if self.language_name == "python" - else "Function" - ), - name, - child.start_point[0] + 1, # 1-based - child.end_point[0] + 1, - params, - parent_name, - code_content, - ) - ) - - elif node_type in type_map["class"]: - name = self.extract_name(child) - code_content = self.get_node_text(child) - result.append( - ( - "ClassDef" if self.language_name == "python" else "Class", - name, - child.start_point[0] + 1, - child.end_point[0] + 1, - [], - None, - code_content, - ) - ) - walk(child, name) - else: - walk(child, parent_name) - - walk(self.root) - return result - - def get_obj_code_info( - self, code_type, code_name, start_line, end_line, params, file_path=None - ): - code_info = {} - code_info["type"] = code_type - code_info["name"] = code_name - code_info["md_content"] = [] - code_info["code_start_line"] = start_line - code_info["code_end_line"] = end_line - code_info["params"] = params + @lru_cache(maxsize=None) + def _read_file(self, abs_path): + with open(abs_path, 'r', encoding='utf-8') as f: + return f.read() - target_file_path = file_path if file_path is not None else self.file_path + def read_file(self, file_path): + abs_path = os.path.join(self.repo_path, file_path) + return self._read_file(abs_path) - with open( - os.path.join(self.repo_path, target_file_path), - "r", - encoding="utf-8", - ) as code_file: - lines = code_file.readlines() - code_content = "".join(lines[start_line - 1 : end_line]) - name_column = lines[start_line - 1].find(code_name) if lines else 0 - - have_return = "return" in code_content - - code_info["have_return"] = have_return - code_info["code_content"] = code_content - code_info["name_column"] = name_column - - return code_info + @lru_cache(maxsize=None) + def parse_code(self, code: str, language_name: str): + parser = Parser() + parser.set_language(LANGUAGE_MAPPING[language_name]) + tree = parser.parse(code.encode('utf8')) + return tree.root_node + @lru_cache(maxsize=None) def generate_file_structure(self, file_path): - """Generate structure for a single file""" - # Detect language for this specific file language = self._detect_language(file_path) - if not language or language not in LANGUAGE_MAPPING: - # Fallback to basic parsing for unsupported files - return self._fallback_file_structure(file_path) - - # Temporarily switch language if different from current - original_language = self.language_name - original_parser = self.parser - - if language != self.language_name: - self.language_name = language - self.ts_language = LANGUAGE_MAPPING[language] - self.parser = Parser(self.ts_language) - - try: - self.parse_file(file_path) - structures = self.get_functions_and_classes() - file_objects = [] - - for struct in structures: - ( - structure_type, - name, - start_line, - end_line, - params, - parent, - code_content, - ) = struct - code_info = self.get_obj_code_info( - structure_type, name, start_line, end_line, params, file_path - ) - # Add parent information - code_info["parent"] = parent - file_objects.append(code_info) - - return file_objects - - finally: - # Restore original language settings - self.language_name = original_language - self.parser = original_parser - - def _fallback_file_structure(self, file_path): - """Fallback parsing for unsupported file types""" - # For unsupported files, return empty structure - return [] + return [] + code = self.read_file(file_path) + root = self.parse_code(code, language) + result = [] - def generate_overall_structure(self, file_path_reflections, jump_files) -> dict: - """Generate structure for entire repository""" - repo_structure = {} - gitignore_checker = GitignoreChecker( - directory=self.repo_path, - gitignore_path=os.path.join(self.repo_path, ".gitignore"), - ) + # Однопроходный обход дерева (O(n)) + stack = [(root, None)] + while stack: + node, parent = stack.pop() + type_map = NODE_TYPES[language] + if node.type in type_map['function']: + name = self.extract_name(node) + params = self.extract_parameters(node) + start, end = node.start_point[0]+1, node.end_point[0]+1 + snippet = code.splitlines()[start-1:end] + result.append((language, 'func', name, start, end, params, parent, '\n'.join(snippet))) + elif node.type in type_map['class']: + name = self.extract_name(node) + start, end = node.start_point[0]+1, node.end_point[0]+1 + snippet = code.splitlines()[start-1:end] + result.append((language, 'class', name, start, end, [], None, '\n'.join(snippet))) + parent = name + for child in reversed(node.children): + stack.append((child, parent)) + return result - bar = tqdm(gitignore_checker.check_files_and_folders()) - for not_ignored_files in bar: - normal_file_names = not_ignored_files - if not_ignored_files in jump_files: - print( - f"{Fore.LIGHTYELLOW_EX}[TreeSitter-Parser] Unstaged AddFile, ignore this file: {Style.RESET_ALL}{normal_file_names}" - ) + def generate_overall_structure(self, jump_files=None): + jump_files = set(jump_files or []) + cache = {} + gitignore_checker = GitignoreChecker(self.repo_path, os.path.join(self.repo_path, '.gitignore')) + structure = {} + for path in tqdm(gitignore_checker.check_files_and_folders()): + if path in jump_files or is_latest_version_file_regex(path): continue - elif is_latest_version_file_regex(not_ignored_files): - print( - f"{Fore.LIGHTYELLOW_EX}[TreeSitter-Parser] Skip Latest Version, Using Git-Status Version]: {Style.RESET_ALL}{normal_file_names}" - ) + # Хеш-контроль, чтобы не парсить заново + content = self.read_file(path) + h = hashlib.sha256(content.encode('utf-8')).hexdigest() + if cache.get(path) == h: continue - try: - repo_structure[normal_file_names] = self.generate_file_structure( - not_ignored_files - ) - except Exception as e: - logger.error( - f"Alert: An error occurred while generating file structure for {not_ignored_files}: {e}" - ) - continue - bar.set_description(f"generating repo structure: {not_ignored_files}") - return repo_structure - - def convert_to_markdown_file(self, file_path=None): - with open(self.project_hierarchy, "r", encoding="utf-8") as f: - json_data = json.load(f) - - if file_path is None: - file_path = self.file_path - - if isinstance(json_data.get(file_path), list): - file_objects = json_data.get(file_path, []) - file_dict = {} - for obj in file_objects: - file_dict[obj["name"]] = obj - else: - file_dict = json_data.get(file_path, {}) - - if not file_dict: - raise ValueError( - f"No file object found for {file_path} in project_hierarchy.json" - ) - - markdown = "" - parent_dict = {} - - if hasattr(file_dict, "values"): - objects = sorted(file_dict.values(), key=lambda obj: obj["code_start_line"]) - else: - objects = sorted(file_dict, key=lambda obj: obj["code_start_line"]) - - for obj in objects: - if obj.get("parent") is not None: - parent_dict[obj["name"]] = obj["parent"] - - current_parent = None - for obj in objects: - level = 1 - parent = obj.get("parent") - while parent is not None: - level += 1 - parent = parent_dict.get(parent) - if level == 1 and current_parent is not None: - markdown += "***\n" - current_parent = obj["name"] - params_str = "" - if obj["type"] in ["FunctionDef", "AsyncFunctionDef", "Function"]: - params_str = "()" - if obj.get("params"): - params_str = f"({', '.join(obj['params'])})" - markdown += f"{'#' * level} {obj['type']} {obj['name']}{params_str}:\n" - md_content = obj.get("md_content", []) - markdown += f"{md_content[-1] if len(md_content) > 0 else ''}\n" - markdown += "***\n" - - return markdown + objs = self.generate_file_structure(path) + structure[path] = objs + cache[path] = h + return structure \ No newline at end of file