⚡️ Speed up function _extract_function_from_code by 11% in PR #1256 (refactor/use-function-to-optimize-in-js)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1256

If you approve this dependent PR, these changes will be merged into the original PR branch refactor/use-function-to-optimize-in-js.

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

---

📄 11% (0.11x) speedup for _extract_function_from_code in codeflash/code_utils/code_replacer.py

⏱️ Runtime : 1.91 milliseconds → 1.73 milliseconds (best of 79 runs)

📝 Explanation and details

The optimized code achieves a 10% runtime improvement through two key changes that eliminate unnecessary work in common code paths:

1. Deferred splitlines() call
The original code called source_code.splitlines(keepends=True) for every function candidate that matched the target name, even when that candidate had invalid line numbers (missing ending_line or invalid starting_line). The optimization moves this expensive string operation until after validating that both effective_start and func.ending_line exist via an early continue statement. This is particularly effective because:

• String splitting is computationally expensive, especially for large source files
• The validation check is very cheap (just boolean/None checks)
• Test results show significant gains in edge cases: test_missing_ending_line_returns_none runs 36.3% faster and test_extract_with_starting_line_none runs 7.14% faster

2. Guarded debug logging
The original code unconditionally formatted the debug log message string (via f-string evaluation) in exception handlers, even when debug logging was disabled. The optimization wraps this in if logger.isEnabledFor(logging.DEBUG):, preventing unnecessary string formatting in production environments where debug logging is typically off. This shows dramatic improvement in exception cases: test_extract_with_exception_in_discover_functions runs 34.8% faster and test_discover_functions_exception_handling runs 10.9% faster.

Performance characteristics by workload:

• Functions with invalid metadata (None values): 7-36% faster due to avoided splitlines
• Exception handling paths: 10-35% faster due to conditional logging
• Large files with many functions: 5-19% faster as deferred splitlines reduces overhead when iterating through non-matching functions
• Standard extraction cases: 1-6% faster from accumulated micro-optimizations

The optimizations are most beneficial when the function being extracted is not the first candidate checked or when processing large source files with many functions, as they reduce cumulative overhead from repeated unnecessary operations.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from __future__ import annotations

from pathlib import Path  # used by the function signature
from types import SimpleNamespace  # lightweight container for attributes

# imports
import pytest  # used for our unit tests
from codeflash.cli_cmds.console import \
    logger  # logger used by the function (imported by original function)
from codeflash.code_utils.code_replacer import _extract_function_from_code
from codeflash.languages.base import \
    LanguageSupport  # real Protocol used by the function

# ------------------------------
# Tests for _extract_function_from_code
# ------------------------------

# Helper: create a SimpleNamespace representing function metadata
def _make_func_info(name: str, starting_line: int | None, ending_line: int | None, doc_start_line: int | None = None):
    # Use attribute names expected by the function under test
    return SimpleNamespace(function_name=name, starting_line=starting_line, ending_line=ending_line, doc_start_line=doc_start_line)

def _make_lang_support_with_functions(functions_list):
    """
    Create a minimal lang_support-like object with discover_functions_from_source method.
    We use SimpleNamespace to avoid defining classes while still providing the callable attribute.
    """
    def discover_functions_from_source(source_code, file_path):
        # ignore inputs; return the provided function info objects
        return functions_list
    return SimpleNamespace(discover_functions_from_source=discover_functions_from_source)

def _make_lang_support_that_raises(exc):
    """Create a lang_support-like object whose discover function raises an exception."""
    def discover_functions_from_source(source_code, file_path):
        raise exc
    return SimpleNamespace(discover_functions_from_source=discover_functions_from_source)

def test_basic_extract_single_function():
    # Construct a small source containing two functions; we will extract the second one.
    source_lines = [
        "def a():\n",
        "    return 'a'\n",
        "\n",
        "def target():\n",
        "    '''doc'''\n",
        "    return 42\n",
        "\n",
    ]
    source = "".join(source_lines)
    # The target function starts at line 4 and ends at line 6 (inclusive)
    func_info = _make_func_info("target", starting_line=4, ending_line=6, doc_start_line=None)
    lang_support = _make_lang_support_with_functions([func_info])

    # Call the function under test
    codeflash_output = _extract_function_from_code(lang_support, source, "target", None); extracted = codeflash_output # 3.50μs -> 3.44μs (1.78% faster)

    # Expect the concatenation of lines 4..6 (1-based index)
    expected = "".join(source_lines[3:6])  # indexes 3..5 in 0-based

def test_include_doc_comment_when_doc_start_line_provided():
    # Source with a comment block preceding the function (e.g., JSDoc or header comment)
    source_lines = [
        "// This is a header comment\n",
        "// More description\n",
        "def target():\n",
        "    return 'ok'\n",
        "\n",
    ]
    source = "".join(source_lines)
    # Provide doc_start_line=1 to include the two comment lines plus the function (1-based)
    func_info = _make_func_info("target", starting_line=3, ending_line=4, doc_start_line=1)
    lang_support = _make_lang_support_with_functions([func_info])

    codeflash_output = _extract_function_from_code(lang_support, source, "target", None); extracted = codeflash_output # 2.83μs -> 3.02μs (6.28% slower)

    # Should include lines 1..4 inclusive
    expected = "".join(source_lines[0:4])

def test_return_none_when_function_not_found():
    source = "def a():\n    pass\n"
    # Provide a function with a different name
    func_info = _make_func_info("other", starting_line=1, ending_line=2)
    lang_support = _make_lang_support_with_functions([func_info])

    codeflash_output = _extract_function_from_code(lang_support, source, "missing", None); extracted = codeflash_output # 921ns -> 971ns (5.15% slower)

def test_discover_functions_exception_handling(monkeypatch, caplog):
    # Create a lang_support that raises an error when discovered
    lang_support = _make_lang_support_that_raises(RuntimeError("parse error"))

    # Ensure the logger's level captures debug for assertion
    caplog.set_level("DEBUG")

    codeflash_output = _extract_function_from_code(lang_support, "some source", "any", None); extracted = codeflash_output # 716μs -> 646μs (10.9% faster)

    # Optionally verify that the debug message was emitted
    # We don't assert exact formatting but ensure the error message text appears
    log_messages = "\n".join(record.getMessage() for record in caplog.records)

def test_starting_line_out_of_range_returns_none():
    source_lines = [
        "line1\n",
        "line2\n",
    ]
    source = "".join(source_lines)
    # starting_line is 100 which is beyond the available number of lines
    func_info = _make_func_info("far", starting_line=100, ending_line=101)
    lang_support = _make_lang_support_with_functions([func_info])

    codeflash_output = _extract_function_from_code(lang_support, source, "far", None); extracted = codeflash_output # 2.40μs -> 2.51μs (4.73% slower)

def test_missing_ending_line_returns_none():
    source = "def target():\n    pass\n"
    # ending_line is None which is falsy and should cause early exit
    func_info = _make_func_info("target", starting_line=1, ending_line=None)
    lang_support = _make_lang_support_with_functions([func_info])

    codeflash_output = _extract_function_from_code(lang_support, source, "target", None); extracted = codeflash_output # 1.88μs -> 1.38μs (36.3% faster)

def test_doc_start_line_zero_falls_back_to_starting_line():
    source_lines = [
        "# stray header that should NOT be included\n",
        "def target():\n",
        "    return 1\n",
    ]
    source = "".join(source_lines)
    # If doc_start_line is 0, the 'or' in effective_start should fallback to starting_line (2)
    func_info = _make_func_info("target", starting_line=2, ending_line=3, doc_start_line=0)
    lang_support = _make_lang_support_with_functions([func_info])

    codeflash_output = _extract_function_from_code(lang_support, source, "target", None); extracted = codeflash_output # 2.65μs -> 2.98μs (11.4% slower)
    expected = "".join(source_lines[1:3])  # lines 2..3

def test_large_scale_many_functions_extraction():
    # Build a source consisting of many small functions. Keep the total under 1000 lines.
    n = 200  # number of small functions - well under 1000
    source_lines = []
    functions_meta = []
    current_line = 1
    # Create n functions "def fn_{i}():\n    return {i}\n\n"
    for i in range(n):
        name = f"fn_{i}"
        fn_lines = [f"def {name}():\n", f"    return {i}\n", "\n"]
        source_lines.extend(fn_lines)
        starting = current_line
        ending = current_line + len(fn_lines) - 1  # inclusive
        # no doc comment provided
        functions_meta.append(_make_func_info(name, starting_line=starting, ending_line=ending, doc_start_line=None))
        current_line += len(fn_lines)

    source = "".join(source_lines)

    # Choose a target roughly in the middle to ensure correct indexing
    target_index = 150
    target_name = f"fn_{target_index}"
    lang_support = _make_lang_support_with_functions(functions_meta)

    codeflash_output = _extract_function_from_code(lang_support, source, target_name, None); extracted = codeflash_output # 23.2μs -> 23.3μs (0.086% slower)

    # Expected lines correspond to the function's chunk
    # Calculate slice boundaries for the target
    target_meta = functions_meta[target_index]
    expected = "".join(source.splitlines(keepends=True)[target_meta.starting_line - 1 : target_meta.ending_line])

def test_multiple_functions_same_name_returns_first_match():
    source_lines = [
        "def target():\n",
        "    return 'first'\n",
        "\n",
        "def target():\n",
        "    return 'second'\n",
        "\n",
    ]
    source = "".join(source_lines)
    # Two function entries with same name; first comes earlier in the file.
    func1 = _make_func_info("target", starting_line=1, ending_line=2)
    func2 = _make_func_info("target", starting_line=4, ending_line=5)
    lang_support = _make_lang_support_with_functions([func1, func2])

    codeflash_output = _extract_function_from_code(lang_support, source, "target", None); extracted = codeflash_output # 2.75μs -> 3.03μs (9.29% slower)
    expected = "".join(source_lines[0:2])  # first target block lines

def test_none_file_path_is_allowed_and_ignored():
    source = "def a():\n    pass\n"
    func_info = _make_func_info("a", starting_line=1, ending_line=2)
    lang_support = _make_lang_support_with_functions([func_info])

    # file_path is optional; pass None explicitly and expect successful extraction
    codeflash_output = _extract_function_from_code(lang_support, source, "a", None); extracted = codeflash_output # 2.44μs -> 2.63μs (7.21% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

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

# imports
import pytest
from codeflash.code_utils.code_replacer import _extract_function_from_code

def test_extract_simple_function_without_docstring():
    """Test extracting a simple function without any documentation."""
    source_code = "def simple_func():\n    return 42\n\ndef other_func():\n    pass\n"
    
    # Mock the language support object
    mock_lang_support = Mock()
    
    # Create mock FunctionInfo objects
    func_info = Mock()
    func_info.function_name = "simple_func"
    func_info.doc_start_line = None  # No docstring
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    other_func_info = Mock()
    other_func_info.function_name = "other_func"
    other_func_info.doc_start_line = None
    other_func_info.starting_line = 4
    other_func_info.ending_line = 5
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info, other_func_info]
    
    # Extract the function
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "simple_func", None); result = codeflash_output # 21.1μs -> 20.7μs (1.74% faster)

def test_extract_function_with_docstring():
    """Test extracting a function that has a docstring."""
    source_code = '"""Module docstring"""\n\ndef documented_func():\n    """Function docstring."""\n    return "result"\n'
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "documented_func"
    func_info.doc_start_line = 3  # Start from docstring line
    func_info.starting_line = 4
    func_info.ending_line = 5
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "documented_func", None); result = codeflash_output # 19.8μs -> 19.1μs (3.30% faster)

def test_extract_correct_function_from_multiple():
    """Test extracting the correct function when multiple functions exist."""
    source_code = "def func_a():\n    pass\n\ndef func_b():\n    pass\n\ndef func_c():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_a = Mock()
    func_a.function_name = "func_a"
    func_a.doc_start_line = None
    func_a.starting_line = 1
    func_a.ending_line = 2
    
    func_b = Mock()
    func_b.function_name = "func_b"
    func_b.doc_start_line = None
    func_b.starting_line = 4
    func_b.ending_line = 5
    
    func_c = Mock()
    func_c.function_name = "func_c"
    func_c.doc_start_line = None
    func_c.starting_line = 7
    func_c.ending_line = 8
    
    mock_lang_support.discover_functions_from_source.return_value = [func_a, func_b, func_c]
    
    # Request func_b specifically
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func_b", None); result = codeflash_output # 19.9μs -> 19.3μs (2.91% faster)

def test_extract_nonexistent_function_returns_none():
    """Test that extracting a nonexistent function returns None."""
    source_code = "def existing_func():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "existing_func"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "nonexistent_func", None); result = codeflash_output # 16.5μs -> 16.7μs (1.32% slower)

def test_extract_from_empty_source_code():
    """Test extracting from empty source code."""
    source_code = ""
    
    mock_lang_support = Mock()
    mock_lang_support.discover_functions_from_source.return_value = []
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "any_func", None); result = codeflash_output # 16.2μs -> 16.6μs (2.12% slower)

def test_extract_with_exception_in_discover_functions():
    """Test that exceptions in discover_functions are caught and None is returned."""
    source_code = "def func():\n    pass\n"
    
    mock_lang_support = Mock()
    # Simulate an exception during function discovery
    mock_lang_support.discover_functions_from_source.side_effect = ValueError("Parse error")
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", None); result = codeflash_output # 28.5μs -> 21.1μs (34.8% faster)

def test_extract_with_empty_function_name():
    """Test extracting with an empty function name."""
    source_code = "def func():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    # Search for empty string
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "", None); result = codeflash_output # 17.1μs -> 16.6μs (3.20% faster)

def test_extract_with_doc_start_line_none():
    """Test extraction when doc_start_line is None but starting_line exists."""
    source_code = "def func():\n    x = 1\n    return x\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 3
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", None); result = codeflash_output # 18.3μs -> 18.5μs (1.29% slower)

def test_extract_with_starting_line_none():
    """Test extraction when starting_line is None."""
    source_code = "def func():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = None
    func_info.starting_line = None  # None value
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", None); result = codeflash_output # 17.6μs -> 16.4μs (7.14% faster)

def test_extract_with_ending_line_none():
    """Test extraction when ending_line is None."""
    source_code = "def func():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = None  # None value
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", None); result = codeflash_output # 17.8μs -> 16.8μs (6.27% faster)

def test_extract_with_line_numbers_exceeding_source():
    """Test extraction when line numbers exceed the source code length."""
    source_code = "def func():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = None
    func_info.starting_line = 100  # Way beyond source length
    func_info.ending_line = 200
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", None); result = codeflash_output # 17.9μs -> 17.4μs (2.59% faster)

def test_extract_with_doc_start_line_exceeding_source():
    """Test extraction when doc_start_line exceeds source code length."""
    source_code = "def func():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = 50  # Beyond source length
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", None); result = codeflash_output # 17.4μs -> 17.4μs (0.288% faster)

def test_extract_with_file_path_none():
    """Test extraction with file_path parameter as None."""
    source_code = "def func():\n    return 1\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", file_path=None); result = codeflash_output # 19.0μs -> 18.5μs (2.82% faster)
    
    # Verify discover_functions_from_source was called with None file_path
    mock_lang_support.discover_functions_from_source.assert_called_once_with(source_code, None)

def test_extract_with_file_path_provided():
    """Test extraction with file_path parameter provided."""
    source_code = "function test() { return 42; }\n"
    file_path = Path("/home/user/project/test.js")
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "test"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 1
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "test", file_path=file_path); result = codeflash_output # 18.6μs -> 18.4μs (1.20% faster)
    
    # Verify discover_functions_from_source was called with the file_path
    mock_lang_support.discover_functions_from_source.assert_called_once_with(source_code, file_path)

def test_extract_with_case_sensitive_function_name():
    """Test that function name matching is case sensitive."""
    source_code = "def MyFunc():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "MyFunc"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    # Search for lowercase version
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "myfunc", None); result = codeflash_output # 16.5μs -> 16.5μs (0.363% slower)

def test_extract_with_matching_case():
    """Test successful extraction with correct case."""
    source_code = "def MyFunc():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "MyFunc"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "MyFunc", None); result = codeflash_output # 18.2μs -> 18.1μs (0.779% faster)

def test_extract_with_special_characters_in_function_name():
    """Test extracting function with underscores and numbers in name."""
    source_code = "def _private_func_v2():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "_private_func_v2"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 2
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "_private_func_v2", None); result = codeflash_output # 18.2μs -> 18.3μs (0.602% slower)

def test_extract_multiline_docstring_before_function():
    """Test extraction with multiline docstring before function."""
    source_code = '"""\nMultiline\ndocstring\n"""\ndef func():\n    pass\n'
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "func"
    func_info.doc_start_line = 1  # Start from docstring
    func_info.starting_line = 5
    func_info.ending_line = 6
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func", None); result = codeflash_output # 18.4μs -> 17.6μs (4.38% faster)

def test_extract_function_with_decorator():
    """Test extracting a function with decorator when doc_start_line includes it."""
    source_code = "@decorator\ndef decorated_func():\n    pass\n"
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "decorated_func"
    func_info.doc_start_line = 1  # Include decorator
    func_info.starting_line = 2
    func_info.ending_line = 3
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "decorated_func", None); result = codeflash_output # 18.2μs -> 17.2μs (5.95% faster)

def test_extract_from_large_source_with_many_functions():
    """Test extracting a specific function from source with 100+ functions."""
    # Generate source code with 150 functions
    lines = []
    for i in range(150):
        lines.append(f"def func_{i:03d}():\n")
        lines.append("    pass\n")
    source_code = "".join(lines)
    
    mock_lang_support = Mock()
    
    # Create mock FunctionInfo objects for all functions
    functions = []
    line_num = 1
    for i in range(150):
        func_info = Mock()
        func_info.function_name = f"func_{i:03d}"
        func_info.doc_start_line = None
        func_info.starting_line = line_num
        func_info.ending_line = line_num + 1
        functions.append(func_info)
        line_num += 2
    
    mock_lang_support.discover_functions_from_source.return_value = functions
    
    # Extract function from the middle
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func_075", None); result = codeflash_output # 53.1μs -> 44.7μs (19.0% faster)

def test_extract_large_function_body():
    """Test extracting a function with a very large body (500+ lines)."""
    # Create a function with 500+ lines
    lines = ["def large_func():\n"]
    for i in range(500):
        lines.append(f"    x_{i} = {i}\n")
    lines.append("    return x_499\n")
    source_code = "".join(lines)
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "large_func"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = 502  # 1 def + 500 assignments + 1 return
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "large_func", None); result = codeflash_output # 41.2μs -> 39.4μs (4.40% faster)

def test_extract_with_large_multiline_docstring():
    """Test extracting a function with a very large docstring (100+ lines)."""
    # Create a large docstring
    doc_lines = ['"""\n']
    for i in range(100):
        doc_lines.append(f"Documentation line {i}\n")
    doc_lines.append('"""\n')
    doc_lines.append("def documented_func():\n")
    doc_lines.append("    pass\n")
    
    source_code = "".join(doc_lines)
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "documented_func"
    func_info.doc_start_line = 1  # Include docstring
    func_info.starting_line = 103  # After docstring
    func_info.ending_line = 104
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "documented_func", None); result = codeflash_output # 25.0μs -> 24.8μs (1.13% faster)

def test_extract_with_100_consecutive_functions():
    """Test extracting specific function when 100 functions exist consecutively."""
    # Build source with 100 consecutive functions
    lines = []
    for i in range(100):
        lines.append(f"def func_{i}():\n")
        lines.append(f"    return {i}\n")
    source_code = "".join(lines)
    
    mock_lang_support = Mock()
    
    # Create function info for all
    functions = []
    for i in range(100):
        func_info = Mock()
        func_info.function_name = f"func_{i}"
        func_info.doc_start_line = None
        func_info.starting_line = (i * 2) + 1
        func_info.ending_line = (i * 2) + 2
        functions.append(func_info)
    
    mock_lang_support.discover_functions_from_source.return_value = functions
    
    # Extract from different positions
    for target_idx in [0, 49, 99]:
        codeflash_output = _extract_function_from_code(mock_lang_support, source_code, f"func_{target_idx}", None); result = codeflash_output # 85.7μs -> 80.9μs (5.96% faster)

def test_extract_performance_with_large_source():
    """Test that extraction performs reasonably with large source code (10k+ lines)."""
    # Create large source code
    lines = []
    for i in range(500):
        lines.append(f"def func_{i}():\n")
        for j in range(20):
            lines.append(f"    statement_{j} = {j}\n")
    source_code = "".join(lines)
    
    mock_lang_support = Mock()
    
    # Create a target function in the middle
    target_func = Mock()
    target_func.function_name = "func_250"
    target_func.doc_start_line = None
    target_func.starting_line = (250 * 21) + 1
    target_func.ending_line = target_func.starting_line + 20
    
    other_functions = []
    for i in range(500):
        if i != 250:
            func = Mock()
            func.function_name = f"func_{i}"
            func.doc_start_line = None
            func.starting_line = (i * 21) + 1
            func.ending_line = func.starting_line + 20
            other_functions.append(func)
    
    all_functions = other_functions + [target_func]
    mock_lang_support.discover_functions_from_source.return_value = all_functions
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func_250", None); result = codeflash_output # 539μs -> 458μs (17.7% faster)

def test_extract_handles_unicode_in_large_source():
    """Test extraction works correctly with unicode characters in large source."""
    # Create source with unicode characters
    lines = ["# -*- coding: utf-8 -*-\n"]
    for i in range(50):
        lines.append(f"def func_{i}():\n")
        lines.append(f'    x = "{i} \u2764"  # Unicode emoji\n')
    source_code = "".join(lines)
    
    mock_lang_support = Mock()
    
    target_func = Mock()
    target_func.function_name = "func_25"
    target_func.doc_start_line = None
    target_func.starting_line = (25 * 2) + 1
    target_func.ending_line = target_func.starting_line + 1
    
    functions = []
    for i in range(50):
        func = Mock()
        func.function_name = f"func_{i}"
        func.doc_start_line = None
        func.starting_line = (i * 2) + 1
        func.ending_line = func.starting_line + 1
        functions.append(func)
    
    mock_lang_support.discover_functions_from_source.return_value = functions
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "func_25", None); result = codeflash_output # 33.0μs -> 30.4μs (8.47% faster)

def test_extract_with_complex_nested_structures():
    """Test extracting function with complex nested structures from large file."""
    # Create a function with nested classes, loops, and conditionals
    lines = ["def complex_func():\n"]
    lines.append("    class InnerClass:\n")
    for i in range(10):
        lines.append("        def method_{i}(self):\n")
        lines.append("            for j in range(10):\n")
        lines.append("                if j > 5:\n")
        lines.append("                    x = j\n")
    lines.append("    return InnerClass\n")
    
    source_code = "".join(lines)
    
    mock_lang_support = Mock()
    
    func_info = Mock()
    func_info.function_name = "complex_func"
    func_info.doc_start_line = None
    func_info.starting_line = 1
    func_info.ending_line = len(source_code.splitlines())
    
    mock_lang_support.discover_functions_from_source.return_value = [func_info]
    
    codeflash_output = _extract_function_from_code(mock_lang_support, source_code, "complex_func", None); result = codeflash_output # 22.6μs -> 21.9μs (3.25% faster)
# 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-pr1256-2026-02-03T00.46.30 and push.