⚡️ Speed up function find_helper_functions by 7,426% in PR #1199 (omni-java)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1199

If you approve this dependent PR, these changes will be merged into the original PR branch omni-java.

This PR will be automatically closed if the original PR is merged.

---

📄 7,426% (74.26x) speedup for find_helper_functions in codeflash/languages/java/context.py

⏱️ Runtime : 77.5 milliseconds → 1.03 milliseconds (best of 25 runs)

📝 Explanation and details

The optimized code achieves a 7426% speedup (77.5ms → 1.03ms) by eliminating expensive exception handling for non-existent files.

Key Optimization:

The line profiler reveals that 96.5% of the original runtime (193ms out of 200ms) was spent in logger.warning() calls within exception handlers. The code was attempting to read 165 non-existent helper files, catching FileNotFoundError exceptions, and then logging each failure.

The optimization adds an early file_path.exists() check before attempting to read files:

# New guard clause
if not file_path.exists():
    continue

This prevents:

1. Exception handling overhead: No try-except block execution for missing files
2. Expensive logging: The logger.warning() call consumed 193ms across 165 failures
3. File I/O attempts: No need to even attempt opening non-existent files

The same defensive check is added to _find_same_class_helpers to prevent attempts to read from non-existent function file paths.

Why This Matters:

Based on the function references, find_helper_functions is called during:

• Test discovery and code context extraction (test_integration.py)
• Helper function analysis workflows (test_context.py)

Since the function processes helper files in a loop (181 iterations in the test), avoiding 165 expensive exception-handling cycles per invocation makes this optimization particularly impactful. The test results show this works best when dealing with:

• Many non-existent helper file paths (common in real projects where imports resolve to external dependencies)
• Deep dependency chains with missing files
• Scalability scenarios with 50-100+ helper files where some don't exist

The optimization maintains correctness—all test cases pass with identical output—while dramatically improving performance for the common case of encountering non-existent dependency files during Java code analysis.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 14 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	69.6%

🌀 Click to see Generated Regression Tests

from pathlib import Path
from unittest.mock import MagicMock, patch

# imports
import pytest
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.languages.base import HelperFunction
from codeflash.languages.java.context import find_helper_functions
from codeflash.languages.java.parser import (JavaAnalyzer, JavaMethodNode,
                                             get_java_analyzer)
from codeflash.models.function_types import FunctionParent
from tree_sitter import Node

# Helper function to create a FunctionToOptimize instance for testing
def create_test_function(
    name: str = "testMethod",
    file_path: Path | None = None,
    class_name: str | None = "TestClass",
    start_line: int = 10,
    end_line: int = 20,
) -> FunctionToOptimize:
    """Create a test FunctionToOptimize instance."""
    if file_path is None:
        file_path = Path("/test/TestClass.java")
    
    parents = [FunctionParent(name=class_name, type="ClassDef")] if class_name else []
    
    return FunctionToOptimize(
        function_name=name,
        file_path=file_path,
        starting_line=start_line,
        ending_line=end_line,
        starting_col=0,
        ending_col=10,
        parents=parents,
        is_async=False,
        is_method=class_name is not None,
        language="java",
        doc_start_line=None,
    )

def test_find_helper_functions_returns_list():
    """Test that find_helper_functions returns a list."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    # Mock the analyzer and helper file discovery
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            mock_find_files.return_value = {}
            mock_get_analyzer.return_value = MagicMock()
            
            codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_empty_helper_files():
    """Test find_helper_functions when there are no helper files."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    # Mock to return empty helper files
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                mock_find_files.return_value = {}
                mock_get_analyzer.return_value = MagicMock()
                mock_same_class.return_value = []
                
                codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_same_class_helpers_included():
    """Test that helpers from same class are included."""
    # Create a test function
    func = create_test_function(class_name="TestClass")
    project_root = Path("/test")
    
    # Create a helper function in the same class
    helper = HelperFunction(
        name="helperMethod",
        qualified_name="TestClass.helperMethod",
        file_path=func.file_path,
        source_code="public void helperMethod() {}",
        start_line=25,
        end_line=26,
    )
    
    # Mock helper discovery
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                mock_find_files.return_value = {}
                mock_get_analyzer.return_value = MagicMock()
                mock_same_class.return_value = [helper]
                
                codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_function_without_class():
    """Test find_helper_functions for a top-level function without a class."""
    # Create a function without a class
    func = create_test_function(class_name=None)
    project_root = Path("/test")
    
    # Mock helper discovery
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                mock_find_files.return_value = {}
                mock_get_analyzer.return_value = MagicMock()
                mock_same_class.return_value = []
                
                codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_nonexistent_file():
    """Test find_helper_functions when a helper file doesn't exist."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    # Mock helper files that don't exist
    nonexistent_path = Path("/nonexistent/Helper.java")
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                mock_find_files.return_value = {nonexistent_path: ["HelperClass"]}
                mock_get_analyzer.return_value = MagicMock()
                mock_same_class.return_value = []
                
                # Should handle the error gracefully
                with patch.object(Path, "read_text") as mock_read:
                    mock_read.side_effect = FileNotFoundError("File not found")
                    codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_duplicate_helpers():
    """Test that duplicate helpers are not added twice."""
    # Create a test function
    func = create_test_function(class_name="TestClass")
    project_root = Path("/test")
    
    # Create a helper function
    helper_path = Path("/test/Helper.java")
    helper = HelperFunction(
        name="helperMethod",
        qualified_name="HelperClass.helperMethod",
        file_path=helper_path,
        source_code="public void helperMethod() {}",
        start_line=10,
        end_line=11,
    )
    
    # Mock helper discovery to return same helper from both sources
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.discover_functions_from_source") as mock_discover:
            with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
                with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                    with patch("codeflash.languages.java.context.extract_function_source") as mock_extract:
                        # Mock helper files with one entry
                        mock_find_files.return_value = {helper_path: ["HelperClass"]}
                        
                        # Mock the discovered functions
                        discovered_func = create_test_function(
                            name="helperMethod",
                            file_path=helper_path,
                            class_name="HelperClass"
                        )
                        mock_discover.return_value = [discovered_func]
                        
                        # Mock analyzer
                        mock_analyzer = MagicMock(spec=JavaAnalyzer)
                        mock_get_analyzer.return_value = mock_analyzer
                        
                        # Mock same class helpers empty
                        mock_same_class.return_value = []
                        
                        # Mock extract_function_source
                        mock_extract.return_value = "public void helperMethod() {}"
                        
                        codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_max_depth_zero():
    """Test find_helper_functions with max_depth of 0."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                mock_find_files.return_value = {}
                mock_get_analyzer.return_value = MagicMock()
                mock_same_class.return_value = []
                
                codeflash_output = find_helper_functions(func, project_root, max_depth=0); result = codeflash_output
                
                # Should call find_helper_files with max_depth=0
                mock_find_files.assert_called_once()

def test_find_helper_functions_multiple_helper_files():
    """Test find_helper_functions with multiple helper files."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    # Create multiple helper paths
    helper1_path = Path("/test/Helper1.java")
    helper2_path = Path("/test/Helper2.java")
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.discover_functions_from_source") as mock_discover:
            with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
                with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                    with patch("codeflash.languages.java.context.extract_function_source") as mock_extract:
                        # Mock multiple helper files
                        mock_find_files.return_value = {
                            helper1_path: ["Helper1Class"],
                            helper2_path: ["Helper2Class"],
                        }
                        
                        # Create discovered functions
                        func1 = create_test_function(
                            name="method1",
                            file_path=helper1_path,
                            class_name="Helper1Class"
                        )
                        func2 = create_test_function(
                            name="method2",
                            file_path=helper2_path,
                            class_name="Helper2Class"
                        )
                        
                        # First call returns func1, second returns func2
                        mock_discover.side_effect = [[func1], [func2]]
                        
                        mock_analyzer = MagicMock(spec=JavaAnalyzer)
                        mock_get_analyzer.return_value = mock_analyzer
                        mock_same_class.return_value = []
                        mock_extract.return_value = "method code"
                        
                        codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_many_helper_files():
    """Test find_helper_functions with many helper files (scalability)."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    # Create many helper file paths
    num_helpers = 100
    helper_files = {
        Path(f"/test/Helper{i}.java"): [f"HelperClass{i}"]
        for i in range(num_helpers)
    }
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.discover_functions_from_source") as mock_discover:
            with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
                with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                    with patch("codeflash.languages.java.context.extract_function_source") as mock_extract:
                        mock_find_files.return_value = helper_files
                        
                        # Mock discover to return functions
                        def make_func(path):
                            return [create_test_function(
                                name=f"method_{path.stem}",
                                file_path=path,
                                class_name=f"Class_{path.stem}"
                            )]
                        
                        mock_discover.side_effect = lambda src, path, analyzer=None: make_func(path)
                        
                        mock_analyzer = MagicMock(spec=JavaAnalyzer)
                        mock_get_analyzer.return_value = mock_analyzer
                        mock_same_class.return_value = []
                        mock_extract.return_value = "method code"
                        
                        codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_large_class_helpers():
    """Test find_helper_functions when there are many same-class helpers."""
    # Create a test function
    func = create_test_function(class_name="LargeClass")
    project_root = Path("/test")
    
    # Create many helpers in the same class
    num_helpers = 50
    helpers = [
        HelperFunction(
            name=f"helperMethod{i}",
            qualified_name=f"LargeClass.helperMethod{i}",
            file_path=func.file_path,
            source_code=f"public void helperMethod{i}() {{}}",
            start_line=20 + i,
            end_line=21 + i,
        )
        for i in range(num_helpers)
    ]
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                mock_find_files.return_value = {}
                mock_get_analyzer.return_value = MagicMock()
                mock_same_class.return_value = helpers
                
                codeflash_output = find_helper_functions(func, project_root); result = codeflash_output
                for i, helper in enumerate(result):
                    pass

def test_find_helper_functions_deep_dependency_chain():
    """Test find_helper_functions with deep dependency chains."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    # Create a chain of helper files
    helper_paths = [Path(f"/test/Helper{i}.java") for i in range(10)]
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.discover_functions_from_source") as mock_discover:
            with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
                with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                    with patch("codeflash.languages.java.context.extract_function_source") as mock_extract:
                        # Mock helper files chain
                        helper_dict = {
                            path: [f"HelperClass{i}"]
                            for i, path in enumerate(helper_paths)
                        }
                        mock_find_files.return_value = helper_dict
                        
                        # Mock discover
                        def make_func(path):
                            idx = helper_paths.index(path)
                            return [create_test_function(
                                name=f"method{idx}",
                                file_path=path,
                                class_name=f"Class{idx}"
                            )]
                        
                        mock_discover.side_effect = lambda src, path, analyzer=None: make_func(path)
                        
                        mock_analyzer = MagicMock(spec=JavaAnalyzer)
                        mock_get_analyzer.return_value = mock_analyzer
                        mock_same_class.return_value = []
                        mock_extract.return_value = "method code"
                        
                        codeflash_output = find_helper_functions(func, project_root, max_depth=5); result = codeflash_output

def test_find_helper_functions_complex_qualified_names():
    """Test find_helper_functions with complex nested qualified names."""
    # Create a test function with complex nesting
    func = create_test_function(
        name="method",
        class_name="OuterClass$InnerClass"
    )
    project_root = Path("/test")
    
    # Create helpers with complex names
    helper = HelperFunction(
        name="nestedHelper",
        qualified_name="OuterClass$InnerClass$NestedClass.nestedHelper",
        file_path=func.file_path,
        source_code="public void nestedHelper() {}",
        start_line=30,
        end_line=31,
    )
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
            with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                mock_find_files.return_value = {}
                mock_get_analyzer.return_value = MagicMock()
                mock_same_class.return_value = [helper]
                
                codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_performance_many_duplicates():
    """Test find_helper_functions performance with many duplicate helpers."""
    # Create a test function
    func = create_test_function()
    project_root = Path("/test")
    
    # Create helper files with duplicate entries
    num_files = 50
    helper_files = {}
    for i in range(num_files):
        helper_files[Path(f"/test/Helper{i}.java")] = ["DuplicateHelper"]
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.discover_functions_from_source") as mock_discover:
            with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
                with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                    with patch("codeflash.languages.java.context.extract_function_source") as mock_extract:
                        mock_find_files.return_value = helper_files
                        
                        # All files discover the same function name
                        mock_discover.return_value = [create_test_function(
                            name="duplicateMethod",
                            class_name="DuplicateHelper"
                        )]
                        
                        mock_analyzer = MagicMock(spec=JavaAnalyzer)
                        mock_get_analyzer.return_value = mock_analyzer
                        mock_same_class.return_value = []
                        mock_extract.return_value = "method code"
                        
                        codeflash_output = find_helper_functions(func, project_root); result = codeflash_output

def test_find_helper_functions_qualified_name_construction():
    """Test that qualified names are correctly constructed for helpers."""
    # Create a test function
    func = create_test_function(class_name="TestClass")
    project_root = Path("/test")
    
    helper_path = Path("/test/Helper.java")
    
    with patch("codeflash.languages.java.context.find_helper_files") as mock_find_files:
        with patch("codeflash.languages.java.context.discover_functions_from_source") as mock_discover:
            with patch("codeflash.languages.java.context.get_java_analyzer") as mock_get_analyzer:
                with patch("codeflash.languages.java.context._find_same_class_helpers") as mock_same_class:
                    with patch("codeflash.languages.java.context.extract_function_source") as mock_extract:
                        mock_find_files.return_value = {helper_path: ["HelperClass"]}
                        
                        # Create a discovered function
                        discovered_func = create_test_function(
                            name="helperMethod",
                            file_path=helper_path,
                            class_name="HelperClass"
                        )
                        mock_discover.return_value = [discovered_func]
                        
                        mock_analyzer = MagicMock(spec=JavaAnalyzer)
                        mock_get_analyzer.return_value = mock_analyzer
                        mock_same_class.return_value = []
                        mock_extract.return_value = "public void helperMethod() {}"
                        
                        codeflash_output = find_helper_functions(func, project_root); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1199-2026-02-20T22.33.34 and push.

[claude]

PR Review Summary

Prek Checks

Fixed the following lint issues (committed as ea20983):

• codeflash/languages/java/context.py: W293 blank-line-with-whitespace, extra blank line removed
• codeflash/languages/__init__.py: I001 unsorted-imports (blank line between same-group imports)
• codeflash/languages/registry.py: RUF100 unused-noqa comments removed

Pre-existing issue (not introduced by this PR):

• registry.py has F401 unused-import warnings on the side-effect import support as _ lines. These are intentional imports for language registration and the F401/RUF100 conflict exists on the base branch.

Mypy

9 pre-existing type errors in context.py (none introduced by this PR). All relate to missing type annotations and int | None handling in the broader file, not the optimization changes.

Code Review

No critical issues found. The optimization is clean and correct:

1. Line 664: if not file_path.exists(): continue — Skips non-existent helper files before read_text(), avoiding FileNotFoundError exception handling and logger.warning() overhead.
2. Line 721: if not function.file_path.exists(): return helpers — Same guard for _find_same_class_helpers.

Both changes are behavior-preserving: the except block at line 690 already catches FileNotFoundError and logs a warning. The early guard avoids the exception overhead entirely while producing the same result.

Test Coverage

File	Stmts	Miss	Coverage
codeflash/languages/java/context.py	469	54	88%

• Changed lines coverage: Lines 664 and 721 (the file_path.exists() checks) are executed by tests ✅
• Branch coverage note: The continue/return branches (lines 665, 722) are not taken during tests (all test files exist), but the guard conditions are exercised
• This file is new vs main (part of the omni-java feature branch), so no main-branch comparison is applicable
• 88% coverage exceeds the 75% threshold for new files ✅

Test Results

• 3,159 passed, 58 skipped, 21 failed
• All failures are pre-existing and unrelated to this PR:
  • test_comparator.py (1 failure) — test data issue
  • test_tracer.py (20 failures) — Tracer attribute errors

---

Last updated: 2026-02-20T22:45:00Z