Merge pull request #1226 from codeflash-ai/feat/find-references-javascript

feat: Add find references functionality for JavaScript/TypeScript

Author: misrasaurabh1

codeflash/code_utils/code_extractor.py

@@ -1563,23 +1563,228 @@ def is_numerical_code(code_string: str, function_name: str | None = None) -> boo
 def get_opt_review_metrics(
     source_code: str, file_path: Path, qualified_name: str, project_root: Path, tests_root: Path, language: Language
 ) -> str:
-    if language != Language.PYTHON:
-        # TODO: {Claude} handle function refrences for other languages
-        return ""
+    """Get function reference metrics for optimization review.
+
+    Uses the LanguageSupport abstraction to find references, supporting both Python and JavaScript/TypeScript.
+
+    Args:
+        source_code: Source code of the file containing the function.
+        file_path: Path to the file.
+        qualified_name: Qualified name of the function (e.g., "module.ClassName.method").
+        project_root: Root of the project.
+        tests_root: Root of the tests directory.
+        language: The programming language.
+
+    Returns:
+        Markdown-formatted string with code blocks showing calling functions.
+    """
+    from codeflash.languages.base import FunctionInfo, ParentInfo, ReferenceInfo
+    from codeflash.languages.registry import get_language_support
+
     start_time = time.perf_counter()
+
     try:
+        # Get the language support
+        lang_support = get_language_support(language)
+        if lang_support is None:
+            return ""
+
+        # Parse qualified name to get function name and class name
         qualified_name_split = qualified_name.rsplit(".", maxsplit=1)
         if len(qualified_name_split) == 1:
-            target_function, target_class = qualified_name_split[0], None
+            function_name, class_name = qualified_name_split[0], None
         else:
-            target_function, target_class = qualified_name_split[1], qualified_name_split[0]
-        matches = get_fn_references_jedi(
-            source_code, file_path, project_root, target_function, target_class
-        )  # jedi is not perfect, it doesn't capture aliased references
-        calling_fns_details = find_occurances(qualified_name, str(file_path), matches, project_root, tests_root)
+            function_name, class_name = qualified_name_split[1], qualified_name_split[0]
+
+        # Create a FunctionInfo for the function
+        # We don't have full line info here, so we'll use defaults
+        parents = ()
+        if class_name:
+            parents = (ParentInfo(name=class_name, type="ClassDef"),)
+
+        func_info = FunctionInfo(
+            name=function_name,
+            file_path=file_path,
+            start_line=1,
+            end_line=1,
+            parents=parents,
+            language=language,
+        )
+
+        # Find references using language support
+        references = lang_support.find_references(func_info, project_root, tests_root, max_files=500)
+
+        if not references:
+            return ""
+
+        # Format references as markdown code blocks
+        calling_fns_details = _format_references_as_markdown(
+            references, file_path, project_root, language
+        )
+
     except Exception as e:
+        logger.debug(f"Error getting function references: {e}")
         calling_fns_details = ""
-        logger.debug(f"Investigate {e}")
+
     end_time = time.perf_counter()
     logger.debug(f"Got function references in {end_time - start_time:.2f} seconds")
     return calling_fns_details
+
+
+def _format_references_as_markdown(
+    references: list, file_path: Path, project_root: Path, language: Language
+) -> str:
+    """Format references as markdown code blocks with calling function code.
+
+    Args:
+        references: List of ReferenceInfo objects.
+        file_path: Path to the source file (to exclude).
+        project_root: Root of the project.
+        language: The programming language.
+
+    Returns:
+        Markdown-formatted string.
+    """
+    # Group references by file
+    refs_by_file: dict[Path, list] = {}
+    for ref in references:
+        # Exclude the source file's definition/import references
+        if ref.file_path == file_path and ref.reference_type in ("import", "reexport"):
+            continue
+
+        if ref.file_path not in refs_by_file:
+            refs_by_file[ref.file_path] = []
+        refs_by_file[ref.file_path].append(ref)
+
+    fn_call_context = ""
+    context_len = 0
+
+    for ref_file, file_refs in refs_by_file.items():
+        if context_len > MAX_CONTEXT_LEN_REVIEW:
+            break
+
+        try:
+            path_relative = ref_file.relative_to(project_root)
+        except ValueError:
+            continue
+
+        # Get syntax highlighting language
+        ext = ref_file.suffix.lstrip(".")
+        if language == Language.PYTHON:
+            lang_hint = "python"
+        elif ext in ("ts", "tsx"):
+            lang_hint = "typescript"
+        else:
+            lang_hint = "javascript"
+
+        # Read the file to extract calling function context
+        try:
+            file_content = ref_file.read_text(encoding="utf-8")
+            lines = file_content.splitlines()
+        except Exception:
+            continue
+
+        # Get unique caller functions from this file
+        callers_seen: set[str] = set()
+        caller_contexts: list[str] = []
+
+        for ref in file_refs:
+            caller = ref.caller_function or "<module>"
+            if caller in callers_seen:
+                continue
+            callers_seen.add(caller)
+
+            # Extract context around the reference
+            if ref.caller_function:
+                # Try to extract the full calling function
+                func_code = _extract_calling_function(file_content, ref.caller_function, ref.line, language)
+                if func_code:
+                    caller_contexts.append(func_code)
+                    context_len += len(func_code)
+            else:
+                # Module-level call - show a few lines of context
+                start_line = max(0, ref.line - 3)
+                end_line = min(len(lines), ref.line + 2)
+                context_code = "\n".join(lines[start_line:end_line])
+                caller_contexts.append(context_code)
+                context_len += len(context_code)
+
+        if caller_contexts:
+            fn_call_context += f"```{lang_hint}:{path_relative}\n"
+            fn_call_context += "\n".join(caller_contexts)
+            fn_call_context += "\n```\n"
+
+    return fn_call_context
+
+
+def _extract_calling_function(source_code: str, function_name: str, ref_line: int, language: Language) -> str | None:
+    """Extract the source code of a calling function.
+
+    Args:
+        source_code: Full source code of the file.
+        function_name: Name of the function to extract.
+        ref_line: Line number where the reference is.
+        language: The programming language.
+
+    Returns:
+        Source code of the function, or None if not found.
+    """
+    if language == Language.PYTHON:
+        return _extract_calling_function_python(source_code, function_name, ref_line)
+    else:
+        return _extract_calling_function_js(source_code, function_name, ref_line)
+
+
+def _extract_calling_function_python(source_code: str, function_name: str, ref_line: int) -> str | None:
+    """Extract the source code of a calling function in Python."""
+    try:
+        import ast
+
+        tree = ast.parse(source_code)
+        lines = source_code.splitlines()
+
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                if node.name == function_name:
+                    # Check if the reference line is within this function
+                    start_line = node.lineno
+                    end_line = node.end_lineno or start_line
+                    if start_line <= ref_line <= end_line:
+                        return "\n".join(lines[start_line - 1 : end_line])
+        return None
+    except Exception:
+        return None
+
+
+def _extract_calling_function_js(source_code: str, function_name: str, ref_line: int) -> str | None:
+    """Extract the source code of a calling function in JavaScript/TypeScript.
+
+    Args:
+        source_code: Full source code of the file.
+        function_name: Name of the function to extract.
+        ref_line: Line number where the reference is (helps identify the right function).
+
+    Returns:
+        Source code of the function, or None if not found.
+    """
+    try:
+        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+
+        # Try TypeScript first, fall back to JavaScript
+        for lang in [TreeSitterLanguage.TYPESCRIPT, TreeSitterLanguage.TSX, TreeSitterLanguage.JAVASCRIPT]:
+            try:
+                analyzer = TreeSitterAnalyzer(lang)
+                functions = analyzer.find_functions(source_code, include_methods=True)
+
+                for func in functions:
+                    if func.name == function_name:
+                        # Check if the reference line is within this function
+                        if func.start_line <= ref_line <= func.end_line:
+                            return func.source_text
+                break
+            except Exception:
+                continue
+
+        return None
+    except Exception:
+        return None

codeflash/languages/base.py

@@ -236,6 +236,37 @@ class FunctionFilterCriteria:
     max_lines: int | None = None
 
 
+@dataclass
+class ReferenceInfo:
+    """Information about a reference (call site) to a function.
+
+    This class captures information about where a function is called
+    from, including the file, line number, context, and caller function.
+
+    Attributes:
+        file_path: Path to the file containing the reference.
+        line: Line number (1-indexed).
+        column: Column number (0-indexed).
+        end_line: End line number (1-indexed).
+        end_column: End column number (0-indexed).
+        context: The line of code containing the reference.
+        reference_type: Type of reference ("call", "callback", "memoized", "import", "reexport").
+        import_name: Name used to import the function (may differ from original).
+        caller_function: Name of the function containing this reference (or None for module-level).
+
+    """
+
+    file_path: Path
+    line: int
+    column: int
+    end_line: int
+    end_column: int
+    context: str
+    reference_type: str
+    import_name: str | None
+    caller_function: str | None = None
+
+
 @runtime_checkable
 class LanguageSupport(Protocol):
     """Protocol defining what a language implementation must provide.
@@ -357,6 +388,29 @@ def find_helper_functions(self, function: FunctionInfo, project_root: Path) -> l
         """
         ...
 
+    def find_references(
+        self, function: FunctionInfo, project_root: Path, tests_root: Path | None = None, max_files: int = 500
+    ) -> list[ReferenceInfo]:
+        """Find all references (call sites) to a function across the codebase.
+
+        This method finds all places where a function is called, including:
+        - Direct calls
+        - Callbacks (passed to other functions)
+        - Memoized versions
+        - Re-exports
+
+        Args:
+            function: The function to find references for.
+            project_root: Root of the project to search.
+            tests_root: Root of tests directory (references in tests are excluded).
+            max_files: Maximum number of files to search.
+
+        Returns:
+            List of ReferenceInfo objects describing each reference location.
+
+        """
+        ...
+
     # === Code Transformation ===
 
     def replace_function(self, source: str, function: FunctionInfo, new_source: str) -> str: