⚡️ Speed up function find_codeflash_output_assignments by 61% in PR #358 (fix-test-reporting)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #358

If you approve this dependent PR, these changes will be merged into the original PR branch fix-test-reporting.

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

---

📄 61% (0.61x) speedup for find_codeflash_output_assignments in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 74.3 milliseconds → 46.1 milliseconds (best of 55 runs)

📝 Explanation and details

Here’s how you can optimize your program for runtime.

Analysis

• Bottleneck: The vast majority of time is spent in visitor.visit(tree) (82.9%). This suggests that.
  • CfoVisitor's implementation (not included) should be optimized, but since its code isn’t given here, we'll focus on the lines given.
• Parsing overhead: ast.parse(source_code) is the next biggest cost (16.8%), but must happen.
• Other lines: Negligible.

Direct External Optimizations

There are limited options without refactoring CfoVisitor. But we can.

• Reuse the AST if the same source gets passed repeatedly, via a simple cache (if that's plausible for your app).
• Remove redundant code. ([The instantiation of CfoVisitor is already minimal.])
• Use __slots__ in CfoVisitor if you control its code (not given).
• Make visitor traversal more efficient, or swap for a faster implementation, if possible (but assuming we can't here).

Safe Minimal Acceleration (with your visible code)

To improve Python's AST speed for repeated jobs you can use the builtin compile cache. Python 3.9+ [via ast.parse does not by itself cache, but compile() can]. However, since ast.parse constructs an AST, and we use CfoVisitor (unknown) we can't avoid it.

1. Use ast.NodeVisitor().visit Directly

This is as direct as your code, but no faster.

2. Use "fast mode" for ast if available ([no such param in stdlib])

3. Use LRU Cache for repeated source (if same string is used multiple times)

If your function may receive duplicates, memoize the result.

• This only helps if the same source_code appears repeatedly.

4. If CfoVisitor doesn't use the source_code string itself.

• Pass the AST only.

But it appears your visitor uses both the AST and source code string.

5. Further Acceleration: Avoid class usage for simple visitors

If you have access to the CfoVisitor code, and it's a simple AST visitor, you could rewrite it as a generator function. This change is NOT possible unless we know what that visitor does.

---

Summing up:

Since the main cost is inside CfoVisitor.visit, and you cannot change CfoVisitor, the only safe optimization at this level is to memoize the parse step if repeat calls for identical inputs arise.

Final Code: Faster for repeated inputs

This form will be notably faster only if source_code is not unique every time.

Otherwise.

• The bottleneck is in CfoVisitor. You would need to optimize that class and its visit logic for further speed gains.

---

If you provide the CfoVisitor code, I can directly optimize the expensive function.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import \
    find_codeflash_output_assignments

# unit tests

# ------------------- BASIC TEST CASES -------------------

def test_single_assignment():
    # Test a single assignment
    code = "codeflash_output = 42"
    codeflash_output = find_codeflash_output_assignments(code) # 29.1μs -> 12.9μs (125% faster)

def test_multiple_assignments():
    # Test multiple assignments on different lines
    code = (
        "codeflash_output = 1\n"
        "x = 2\n"
        "codeflash_output = 3\n"
        "y = codeflash_output\n"
        "codeflash_output = 4"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 59.4μs -> 33.0μs (79.8% faster)

def test_assignment_to_other_var():
    # Test assignment to a different variable
    code = "output = 123"
    codeflash_output = find_codeflash_output_assignments(code) # 25.4μs -> 11.7μs (116% faster)

def test_assignment_in_function():
    # Assignment inside a function
    code = (
        "def foo():\n"
        "    codeflash_output = 10\n"
        "    return codeflash_output"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 43.3μs -> 20.3μs (114% faster)

def test_assignment_in_class():
    # Assignment inside a class method
    code = (
        "class C:\n"
        "    def m(self):\n"
        "        codeflash_output = 99"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 44.9μs -> 20.6μs (118% faster)

def test_assignment_with_annotation():
    # Assignment with type annotation
    code = "codeflash_output: int = 7"
    codeflash_output = find_codeflash_output_assignments(code) # 29.2μs -> 13.8μs (112% faster)

def test_augmented_assignment():
    # Test augmented assignment
    code = (
        "codeflash_output = 1\n"
        "codeflash_output += 2"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 37.0μs -> 18.2μs (103% faster)

def test_assignment_in_tuple():
    # Assignment in tuple unpacking
    code = "codeflash_output, x = 1, 2"
    codeflash_output = find_codeflash_output_assignments(code) # 38.3μs -> 20.4μs (88.0% faster)

def test_assignment_in_tuple_reversed():
    # Assignment in tuple unpacking, reversed order
    code = "x, codeflash_output = 1, 2"
    codeflash_output = find_codeflash_output_assignments(code) # 37.1μs -> 20.5μs (80.9% faster)

def test_assignment_in_for_loop():
    # Assignment in a for loop (should not be detected unless assigned directly)
    code = (
        "for codeflash_output in range(3):\n"
        "    pass"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 37.0μs -> 17.5μs (112% faster)

def test_assignment_in_list_comp():
    # Assignment in a list comprehension (walrus operator)
    code = "[x for x in range(5) if (codeflash_output := x) > 2]"
    codeflash_output = find_codeflash_output_assignments(code) # 58.8μs -> 28.7μs (105% faster)

def test_assignment_in_lambda():
    # Assignment inside a lambda (should not be possible)
    code = "f = lambda x: (codeflash_output := x)"
    codeflash_output = find_codeflash_output_assignments(code) # 39.3μs -> 18.0μs (118% faster)

def test_assignment_with_walrus_operator():
    # Assignment with walrus operator at top level
    code = "(codeflash_output := 42)"
    codeflash_output = find_codeflash_output_assignments(code) # 29.0μs -> 12.5μs (133% faster)


def test_assignment_in_multiline():
    # Assignment split over multiple lines
    code = (
        "codeflash_output = (\n"
        "    123\n"
        ")"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 33.0μs -> 14.9μs (122% faster)

def test_assignment_with_comment():
    # Assignment with inline comment
    code = "codeflash_output = 1  # set output"
    codeflash_output = find_codeflash_output_assignments(code) # 27.4μs -> 13.3μs (106% faster)

def test_assignment_to_attribute():
    # Assignment to an attribute (should not match)
    code = "self.codeflash_output = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 30.0μs -> 15.0μs (99.6% faster)

def test_assignment_to_subscript():
    # Assignment to a subscript (should not match)
    code = "codeflash_output[0] = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 33.7μs -> 16.6μs (103% faster)

def test_assignment_in_nested_function():
    # Assignment in a nested function
    code = (
        "def outer():\n"
        "    def inner():\n"
        "        codeflash_output = 9\n"
        "    return inner"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 50.4μs -> 24.4μs (107% faster)

def test_assignment_in_try_except():
    # Assignment in try/except/finally
    code = (
        "try:\n"
        "    codeflash_output = 1\n"
        "except Exception:\n"
        "    codeflash_output = 2\n"
        "finally:\n"
        "    codeflash_output = 3"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 57.4μs -> 29.1μs (97.0% faster)

def test_assignment_in_with_statement():
    # Assignment in with statement
    code = (
        "with open('x') as f:\n"
        "    codeflash_output = f.read()"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 52.2μs -> 25.2μs (107% faster)

def test_assignment_in_if_else():
    # Assignment in if/else branches
    code = (
        "if True:\n"
        "    codeflash_output = 1\n"
        "else:\n"
        "    codeflash_output = 2"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 45.7μs -> 21.8μs (109% faster)

def test_assignment_in_match_case():
    # Assignment in match/case (Python 3.10+)
    code = (
        "match x:\n"
        "    case 1:\n"
        "        codeflash_output = 10\n"
        "    case _:\n"
        "        codeflash_output = 20"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 58.8μs -> 30.7μs (91.3% faster)

def test_assignment_to_multiple_targets():
    # Assignment to multiple targets at once
    code = "codeflash_output = x = y = 5"
    codeflash_output = find_codeflash_output_assignments(code) # 29.2μs -> 14.3μs (105% faster)

def test_assignment_to_tuple_with_multiple_vars():
    # Tuple assignment with multiple variables, only some are codeflash_output
    code = "a, codeflash_output, b = 1, 2, 3"
    codeflash_output = find_codeflash_output_assignments(code) # 44.8μs -> 24.2μs (85.5% faster)

def test_assignment_with_unicode_identifier():
    # Assignment with unicode in variable name (should not match)
    code = "codeflash_output_β = 5"
    codeflash_output = find_codeflash_output_assignments(code) # 29.9μs -> 11.9μs (151% faster)

def test_assignment_with_similar_name():
    # Assignment to a similar but not identical variable name
    code = "codeflash_outputs = 5"
    codeflash_output = find_codeflash_output_assignments(code) # 25.0μs -> 11.5μs (118% faster)

def test_assignment_in_generator_expression():
    # Assignment in generator expression (walrus operator)
    code = "(x for x in range(5) if (codeflash_output := x) > 2)"
    codeflash_output = find_codeflash_output_assignments(code) # 60.5μs -> 29.4μs (106% faster)

def test_assignment_in_nested_listcomp():
    # Assignment in nested list comprehensions (walrus operator)
    code = "[[y for y in range(2) if (codeflash_output := y)] for x in range(2)]"
    codeflash_output = find_codeflash_output_assignments(code) # 67.0μs -> 33.4μs (101% faster)

def test_assignment_in_deeply_nested_structures():
    # Assignment deeply nested in code
    code = (
        "def a():\n"
        "    for i in range(2):\n"
        "        if i == 1:\n"
        "            codeflash_output = i"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 66.4μs -> 32.1μs (106% faster)

def test_assignment_with_no_newline():
    # Assignment with no trailing newline
    code = "codeflash_output = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 25.1μs -> 12.0μs (110% faster)

def test_assignment_with_semicolon():
    # Multiple assignments on one line separated by semicolon
    code = "x = 1; codeflash_output = 2; y = 3"
    codeflash_output = find_codeflash_output_assignments(code) # 42.0μs -> 22.3μs (88.7% faster)

def test_assignment_with_multiple_semicolons():
    # Multiple assignments to codeflash_output on one line with semicolons
    code = "codeflash_output = 1; codeflash_output = 2"
    codeflash_output = find_codeflash_output_assignments(code) # 33.6μs -> 17.1μs (96.1% faster)

def test_assignment_with_parentheses():
    # Assignment with parentheses (should still match)
    code = "(codeflash_output) = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 26.7μs -> 11.5μs (133% faster)

def test_assignment_with_line_continuation():
    # Assignment with explicit line continuation
    code = "codeflash_output = \\\n    5"
    codeflash_output = find_codeflash_output_assignments(code) # 24.2μs -> 11.8μs (104% faster)

# ------------------- LARGE SCALE TEST CASES -------------------

def test_large_number_of_lines():
    # Test with a large number of lines, only some assign codeflash_output
    code_lines = []
    expected = []
    for i in range(1, 1001):
        if i % 100 == 0:
            code_lines.append(f"codeflash_output = {i}")
            expected.append(i)
        else:
            code_lines.append(f"x{i} = {i}")
    code = "\n".join(code_lines)
    codeflash_output = find_codeflash_output_assignments(code)

def test_large_number_of_assignments():
    # Test with many assignments to codeflash_output
    code = "\n".join([f"codeflash_output = {i}" for i in range(1, 1001)])
    expected = list(range(1, 1001))
    codeflash_output = find_codeflash_output_assignments(code) # 6.89ms -> 4.45ms (54.8% faster)

def test_large_file_with_mixed_assignments():
    # Test with large file and mixed assignments
    code_lines = []
    expected = []
    for i in range(1, 1001):
        if i % 2 == 0:
            code_lines.append(f"codeflash_output = {i}")
            expected.append(i)
        else:
            code_lines.append(f"x = {i}")
    code = "\n".join(code_lines)
    codeflash_output = find_codeflash_output_assignments(code)

def test_large_file_with_nested_structures():
    # Large file with nested structures and assignments
    code_lines = ["def foo():", "    for i in range(500):"]
    expected = []
    for i in range(3, 503):
        code_lines.append(f"        codeflash_output = {i}")
        expected.append(i)
    code = "\n".join(code_lines)
    codeflash_output = find_codeflash_output_assignments(code)

def test_large_file_with_various_assignment_types():
    # Large file mixing all assignment types
    code_lines = []
    expected = []
    for i in range(1, 251):
        code_lines.append(f"codeflash_output = {i}")  # simple assign
        expected.append(len(code_lines))
        code_lines.append(f"codeflash_output: int = {i}")  # annotated assign
        expected.append(len(code_lines))
        code_lines.append(f"codeflash_output += {i}")  # augmented assign
        expected.append(len(code_lines))
        code_lines.append(f"codeflash_output, x = {i}, 0")  # tuple unpack
        expected.append(len(code_lines))
        code_lines.append(f"(codeflash_output := {i})")  # walrus
        expected.append(len(code_lines))
    code = "\n".join(code_lines)
    codeflash_output = find_codeflash_output_assignments(code)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import \
    find_codeflash_output_assignments


class CfoVisitor(ast.NodeVisitor):
    def __init__(self, source_code: str):
        self.results = []
        self.source_code = source_code

    def visit_Assign(self, node):
        # Handles: codeflash_output = ...
        for target in node.targets:
            if isinstance(target, ast.Name) and target.id == "codeflash_output":
                self.results.append(node.lineno)
        self.generic_visit(node)

    def visit_AnnAssign(self, node):
        # Handles: codeflash_output: int = ...
        target = node.target
        if isinstance(target, ast.Name) and target.id == "codeflash_output":
            self.results.append(node.lineno)
        self.generic_visit(node)

    def visit_AugAssign(self, node):
        # Handles: codeflash_output += ...
        target = node.target
        if isinstance(target, ast.Name) and target.id == "codeflash_output":
            self.results.append(node.lineno)
        self.generic_visit(node)

    # Do not collect assignments to attributes or subscripts (e.g., codeflash_output[0] = 1, codeflash_output.x = 2)

# unit tests

# ------------------- Basic Test Cases -------------------

def test_single_assignment():
    # Basic assignment, should detect line 1
    code = "codeflash_output = 42"
    codeflash_output = find_codeflash_output_assignments(code) # 28.8μs -> 12.5μs (131% faster)

def test_multiple_assignments():
    # Multiple assignments, should detect both lines
    code = "codeflash_output = 1\nx = 2\ncodeflash_output = 3"
    codeflash_output = find_codeflash_output_assignments(code) # 44.1μs -> 23.4μs (88.0% faster)


def test_assignment_in_function():
    # Assignment inside a function
    code = "def f():\n    codeflash_output = 7"
    codeflash_output = find_codeflash_output_assignments(code) # 40.3μs -> 20.1μs (100% faster)

def test_assignment_in_if_block():
    # Assignment inside if block
    code = "if True:\n    codeflash_output = 8"
    codeflash_output = find_codeflash_output_assignments(code) # 37.1μs -> 18.2μs (104% faster)

def test_augmented_assignment():
    # Augmented assignment (+=)
    code = "codeflash_output = 1\ncodeflash_output += 2"
    codeflash_output = find_codeflash_output_assignments(code) # 38.4μs -> 19.6μs (95.7% faster)

def test_annotated_assignment():
    # Annotated assignment
    code = "codeflash_output: int = 5"
    codeflash_output = find_codeflash_output_assignments(code) # 29.7μs -> 14.6μs (104% faster)

def test_assignment_with_comment():
    # Assignment with a comment
    code = "codeflash_output = 10  # set output"
    codeflash_output = find_codeflash_output_assignments(code) # 25.7μs -> 12.7μs (102% faster)

def test_assignment_within_class_method():
    # Assignment inside a class method
    code = "class X:\n    def m(self):\n        codeflash_output = 3"
    codeflash_output = find_codeflash_output_assignments(code) # 44.7μs -> 20.8μs (114% faster)

# ------------------- Edge Test Cases -------------------

def test_no_assignment():
    # No assignment to codeflash_output
    code = "x = 1\ny = 2"
    codeflash_output = find_codeflash_output_assignments(code) # 34.7μs -> 18.2μs (90.6% faster)

def test_assignment_to_other_variables():
    # Assignment to variables with similar names
    code = "codeflash_outputx = 1\ncodeflash_outputs = 2"
    codeflash_output = find_codeflash_output_assignments(code) # 34.5μs -> 17.4μs (97.6% faster)

def test_assignment_to_attribute():
    # Assignment to attribute should not match
    code = "obj.codeflash_output = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 29.1μs -> 14.3μs (103% faster)

def test_assignment_to_subscript():
    # Assignment to subscript should not match
    code = "codeflash_output[0] = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 31.8μs -> 16.2μs (96.4% faster)

def test_assignment_in_nested_function():
    # Assignment in nested function
    code = "def f():\n    def g():\n        codeflash_output = 2"
    codeflash_output = find_codeflash_output_assignments(code) # 41.9μs -> 20.8μs (101% faster)

def test_assignment_in_class_body():
    # Assignment in class body (class variable)
    code = "class C:\n    codeflash_output = 11"
    codeflash_output = find_codeflash_output_assignments(code) # 31.3μs -> 14.6μs (115% faster)

def test_assignment_in_lambda():
    # Assignment in lambda is invalid Python, should not raise
    code = "f = lambda: (codeflash_output := 1)"
    # This is valid in Python 3.8+ (walrus operator)
    codeflash_output = find_codeflash_output_assignments(code) # 39.0μs -> 19.2μs (104% faster)

def test_assignment_with_multitarget():
    # Multiple targets, only one is codeflash_output
    code = "a = b = codeflash_output = 5"
    codeflash_output = find_codeflash_output_assignments(code) # 30.5μs -> 15.4μs (97.6% faster)

def test_assignment_with_tuple_unpacking():
    # Tuple unpacking, should not match
    code = "codeflash_output, x = 1, 2"
    codeflash_output = find_codeflash_output_assignments(code) # 38.3μs -> 20.8μs (84.4% faster)

def test_assignment_with_list_unpacking():
    # List unpacking, should not match
    code = "[codeflash_output, x] = [1, 2]"
    codeflash_output = find_codeflash_output_assignments(code) # 42.2μs -> 20.7μs (104% faster)

def test_assignment_with_type_comment():
    # Assignment with type comment
    code = "codeflash_output = 1  # type: int"
    codeflash_output = find_codeflash_output_assignments(code) # 25.0μs -> 12.0μs (109% faster)

def test_assignment_with_unicode_name():
    # Unicode in variable name, should not match
    code = "codeflash_output_α = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 29.9μs -> 11.7μs (155% faster)

def test_assignment_with_non_ascii_chars():
    # Assignment with non-ascii value
    code = "codeflash_output = 'π'"
    codeflash_output = find_codeflash_output_assignments(code) # 25.8μs -> 12.0μs (115% faster)

def test_assignment_in_try_except_finally():
    # Assignment in try/except/finally blocks
    code = (
        "try:\n"
        "    codeflash_output = 1\n"
        "except Exception:\n"
        "    codeflash_output = 2\n"
        "finally:\n"
        "    codeflash_output = 3"
    )
    codeflash_output = find_codeflash_output_assignments(code) # 57.7μs -> 29.1μs (98.5% faster)

def test_assignment_in_with_block():
    # Assignment in with block
    code = "with open('f') as f:\n    codeflash_output = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 45.9μs -> 22.9μs (101% faster)

def test_assignment_in_for_loop():
    # Assignment in for loop
    code = "for i in range(2):\n    codeflash_output = i"
    codeflash_output = find_codeflash_output_assignments(code) # 42.7μs -> 20.4μs (109% faster)

def test_assignment_in_while_loop():
    # Assignment in while loop
    code = "while True:\n    codeflash_output = 1\n    break"
    codeflash_output = find_codeflash_output_assignments(code) # 37.3μs -> 17.5μs (114% faster)

def test_assignment_with_semicolon():
    # Assignment with semicolon on same line
    code = "codeflash_output = 1; x = 2"
    codeflash_output = find_codeflash_output_assignments(code) # 33.9μs -> 17.3μs (95.5% faster)

def test_assignment_with_different_line_endings():
    # Windows line endings
    code = "codeflash_output = 1\r\ncodeflash_output = 2\r\n"
    codeflash_output = find_codeflash_output_assignments(code) # 33.6μs -> 17.6μs (91.5% faster)

def test_assignment_with_blank_lines():
    # Blank lines before/after
    code = "\n\ncodeflash_output = 1\n\n"
    codeflash_output = find_codeflash_output_assignments(code) # 24.8μs -> 12.0μs (107% faster)

def test_assignment_in_decorator():
    # Assignment in decorator (invalid, but should not crash)
    code = "@codeflash_output\ndef f(): pass"
    codeflash_output = find_codeflash_output_assignments(code) # 28.4μs -> 11.7μs (142% faster)

def test_assignment_to_global_and_nonlocal():
    # Assignment in global/nonlocal context
    code = "def f():\n    global codeflash_output\n    codeflash_output = 1"
    codeflash_output = find_codeflash_output_assignments(code) # 40.4μs -> 19.0μs (113% faster)

    code2 = "def f():\n    codeflash_output = 2\n    def g():\n        nonlocal codeflash_output\n        codeflash_output = 3"
    # Only the assignments, not the nonlocal/global statements
    codeflash_output = find_codeflash_output_assignments(code2) # 42.6μs -> 21.9μs (94.3% faster)

# ------------------- Large Scale Test Cases -------------------

def test_many_assignments():
    # 1000 assignments, should return all line numbers
    code = "\n".join(f"codeflash_output = {i}" for i in range(1000))
    expected = list(range(1, 1001))
    codeflash_output = find_codeflash_output_assignments(code) # 6.90ms -> 4.44ms (55.5% faster)

def test_large_file_with_noise():
    # 1000 lines, every 10th line assigns to codeflash_output
    lines = []
    expected = []
    for i in range(1000):
        if i % 10 == 0:
            lines.append(f"codeflash_output = {i}")
            expected.append(i + 1)
        else:
            lines.append(f"x{i} = {i}")
    code = "\n".join(lines)
    codeflash_output = find_codeflash_output_assignments(code)

def test_large_file_with_nested_structures():
    # 500 functions, each with an assignment
    code = "\n".join(
        f"def f{i}():\n    codeflash_output = {i}"
        for i in range(500)
    )
    # Each assignment is on the second line of each function
    expected = [2 + 2 * i for i in range(500)]
    codeflash_output = find_codeflash_output_assignments(code) # 6.45ms -> 3.81ms (69.4% faster)

def test_large_file_with_various_assignment_types():
    # Mix of Assign, AnnAssign, AugAssign, and noise
    code_lines = []
    expected = []
    for i in range(1, 1001):
        if i % 3 == 0:
            code_lines.append(f"codeflash_output = {i}")
            expected.append(i)
        elif i % 3 == 1:
            code_lines.append(f"codeflash_output: int = {i}")
            expected.append(i)
        elif i % 3 == 2:
            code_lines.append(f"codeflash_output += {i}")
            expected.append(i)
    code = "\n".join(code_lines)
    codeflash_output = find_codeflash_output_assignments(code)

def test_large_file_with_mixed_line_endings():
    # Mix of \n and \r\n line endings
    code = ""
    expected = []
    for i in range(1, 1001):
        code += f"codeflash_output = {i}"
        expected.append(i)
        if i % 2 == 0:
            code += "\r\n"
        else:
            code += "\n"
    codeflash_output = find_codeflash_output_assignments(code)
# 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-pr358-2025-06-21T00.11.21 and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #358 by aseembits93 was closed.