⚡️ Speed up function add_global_assignments by 18% in PR #683 (fix/duplicate-global-assignments-when-reverting-helpers)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #683

If you approve this dependent PR, these changes will be merged into the original PR branch fix/duplicate-global-assignments-when-reverting-helpers.

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

---

📄 18% (0.18x) speedup for add_global_assignments in codeflash/code_utils/code_extractor.py

⏱️ Runtime : 1.23 seconds → 1.05 seconds (best of 9 runs)

📝 Explanation and details

The optimized code achieves a 17% speedup by eliminating redundant CST parsing operations, which are the most expensive parts of the function according to the line profiler.

Key optimizations:

1. Eliminate duplicate parsing: The original code parsed src_module_code and dst_module_code multiple times. The optimized version introduces _extract_global_statements_once() that parses each module only once and reuses the parsed CST objects throughout the function.

2. Reuse parsed modules: Instead of re-parsing dst_module_code after modifications, the optimized version conditionally reuses the already-parsed dst_module when no global statements need insertion, avoiding unnecessary cst.parse_module() calls.

3. Early termination: Added an early return when new_collector.assignments is empty, avoiding the expensive GlobalAssignmentTransformer creation and visitation when there's nothing to transform.

4. Minor optimization in uniqueness check: Added a fast-path identity check (stmt is existing_stmt) before the expensive deep_equals() comparison, though this has minimal impact.

Performance impact by test case type:

• Empty/minimal cases: Show the highest gains (59-88% faster) due to early termination optimizations
• Standard cases: Achieve consistent 20-30% improvements from reduced parsing
• Large-scale tests: Benefit significantly (18-23% faster) as parsing overhead scales with code size

The optimization is most effective for workloads with moderate to large code files where CST parsing dominates the runtime, as evidenced by the original profiler showing 70%+ of time spent in cst.parse_module() and module.visit() operations.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_code_context_extractor.py::test_circular_deps	161ms	163ms	-1.03%⚠️

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import libcst as cst
# imports
import pytest  # used for our unit tests
from codeflash.code_utils.code_extractor import add_global_assignments


class GlobalAssignmentCollector(cst.CSTVisitor):
    def __init__(self):
        self.assignments = {}
        self.assignment_order = []

    def visit_Assign(self, node: cst.Assign):
        # Only collect assignments at module level
        if self._in_global_scope(node):
            for target in node.targets:
                if isinstance(target.target, cst.Name):
                    name = target.target.value
                    self.assignments[name] = node
                    self.assignment_order.append(name)

    def _in_global_scope(self, node):
        parent = getattr(node, 'parent', None)
        while parent is not None:
            if isinstance(parent, (cst.FunctionDef, cst.ClassDef)):
                return False
            parent = getattr(parent, 'parent', None)
        return True

class ImportInserter(cst.CSTTransformer):
    def __init__(self, stmts, last_import_line):
        self.stmts = stmts
        self.last_import_line = last_import_line
        self.inserted = False

    def leave_Module(self, original_node, updated_node):
        # Insert after last import line, or at top if none
        new_body = []
        for i, stmt in enumerate(updated_node.body):
            new_body.append(stmt)
            if (
                not self.inserted
                and stmt.body
                and hasattr(stmt.body[0], "start")
                and stmt.body[0].start.line == self.last_import_line
            ):
                for s in self.stmts:
                    new_body.append(s)
                self.inserted = True
        if not self.inserted:
            # No imports found, insert at top
            new_body = list(self.stmts) + list(updated_node.body)
        return updated_node.with_changes(body=new_body)

class GlobalAssignmentTransformer(cst.CSTTransformer):
    def __init__(self, assignments, assignment_order):
        self.assignments = assignments
        self.assignment_order = assignment_order
        self.updated = set()

    def leave_Module(self, original_node, updated_node):
        # Remove old assignments of these names, insert new ones in order
        new_body = []
        seen = set()
        for stmt in updated_node.body:
            # Remove any assignment to a name we're updating
            if (
                isinstance(stmt, cst.SimpleStatementLine)
                and stmt.body
                and isinstance(stmt.body[0], cst.Assign)
                and isinstance(stmt.body[0].targets[0].target, cst.Name)
            ):
                name = stmt.body[0].targets[0].target.value
                if name in self.assignments and name not in seen:
                    # Insert new assignment only once, in order
                    new_body.append(self.assignments[name])
                    seen.add(name)
                elif name not in self.assignments:
                    new_body.append(stmt)
            else:
                new_body.append(stmt)
        # Add any new assignments not already present
        for name in self.assignment_order:
            if name not in seen:
                new_body.insert(0, self.assignments[name])
        return updated_node.with_changes(body=new_body)

# --- UNIT TESTS ---

# 1. BASIC TEST CASES

def test_add_simple_assignment():
    # Adding a new assignment to an empty target
    src = "a = 1"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 382μs -> 310μs (23.2% faster)

def test_add_multiple_assignments():
    # Adding multiple assignments, none present in dst
    src = "a = 1\nb = 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 542μs -> 433μs (25.2% faster)

def test_preserve_existing_assignments():
    # Existing assignment should not be duplicated
    src = "a = 1\nb = 2"
    dst = "a = 1"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 740μs -> 574μs (29.0% faster)

def test_update_existing_assignment():
    # Assignment in dst should be updated to src value
    src = "a = 42"
    dst = "a = 1"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 530μs -> 414μs (28.1% faster)

def test_add_after_imports():
    # New assignments should come after imports
    src = "a = 1"
    dst = "import os\nimport sys"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 693μs -> 548μs (26.4% faster)
    lines = result.splitlines()

def test_preserve_imports_and_assignments():
    # Imports and pre-existing assignments are preserved
    src = "a = 1\nb = 2"
    dst = "import os\na = 1"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 885μs -> 705μs (25.4% faster)

def test_assignment_order():
    # Assignments should appear in the order from src
    src = "x = 1\ny = 2\nz = 3"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 706μs -> 557μs (26.7% faster)
    lines = [line.strip() for line in result.splitlines() if line.strip()]

# 2. EDGE TEST CASES

def test_empty_src_and_dst():
    # Both src and dst are empty
    src = ""
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 124μs -> 78.1μs (59.9% faster)

def test_empty_src_nonempty_dst():
    # src is empty, dst has assignments
    src = ""
    dst = "a = 1"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 362μs -> 192μs (88.3% faster)

def test_nonempty_src_empty_dst():
    # src has assignments, dst is empty
    src = "a = 1"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 368μs -> 296μs (24.3% faster)

def test_assignment_in_function_not_global():
    # Assignments inside functions should not be considered global
    src = "def foo():\n    a = 1\nb = 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 691μs -> 542μs (27.6% faster)

def test_assignment_in_class_not_global():
    # Assignments inside classes should not be considered global
    src = "class Foo:\n    a = 1\nb = 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 635μs -> 491μs (29.2% faster)

def test_dst_has_function_should_remain():
    # Functions in dst should be preserved
    src = "a = 1"
    dst = "def foo():\n    pass"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 642μs -> 509μs (26.1% faster)

def test_duplicate_assignments_in_src():
    # If src has duplicate assignments, only last one should be used
    src = "a = 1\na = 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 522μs -> 413μs (26.4% faster)

def test_duplicate_assignments_in_dst():
    # If dst has duplicate assignments, only one should be replaced
    src = "a = 3"
    dst = "a = 1\na = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 689μs -> 548μs (25.6% faster)

def test_assignment_with_complex_expression():
    # Assignment with a complex expression is handled
    src = "a = [x for x in range(5)]"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 837μs -> 655μs (27.9% faster)

def test_assignment_with_comment_preserved():
    # Comments should be preserved if possible
    src = "# comment\na = 1"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 414μs -> 326μs (27.1% faster)

def test_import_from_statements():
    # ImportFrom statements are handled and preserved
    src = "from math import sqrt\na = 1"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 589μs -> 474μs (24.4% faster)

def test_imports_in_middle_of_dst():
    # Imports in the middle of file, new assignments should still go after all imports
    src = "a = 1"
    dst = "x = 0\nimport os\ny = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 911μs -> 710μs (28.2% faster)

def test_order_with_existing_imports_and_assignments():
    # Order is preserved with existing imports and assignments
    src = "a = 1\nb = 2"
    dst = "import os\nc = 3"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 876μs -> 711μs (23.2% faster)

def test_assignment_with_augassign_not_treated_as_assign():
    # Augmented assignments are not treated as new assignments
    src = "a = 1"
    dst = "a += 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 607μs -> 473μs (28.3% faster)

# 3. LARGE SCALE TEST CASES

def test_large_number_of_assignments():
    # Test with a large number of assignments (up to 1000)
    src = "\n".join([f"var{i} = {i}" for i in range(1000)])
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 153ms -> 128ms (19.0% faster)
    for i in range(1000):
        pass

def test_large_number_of_existing_assignments():
    # Test with large dst, and src has a few updates
    dst = "\n".join([f"var{i} = {i}" for i in range(1000)])
    src = "\n".join([f"var{i} = {i+1}" for i in range(500, 510)])
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 144ms -> 117ms (23.0% faster)
    for i in range(500, 510):
        pass

def test_large_imports_and_assignments():
    # Large number of imports, assignments should be after all imports
    src = "\n".join([f"var{i} = {i}" for i in range(10)])
    dst = "\n".join([f"import mod{i}" for i in range(100)])
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 12.9ms -> 10.8ms (19.0% faster)
    lines = result.splitlines()
    # All imports first
    for i in range(100):
        pass
    # Then assignments
    for i in range(10):
        pass

def test_performance_on_large_src_and_dst():
    # Ensure function completes on large src and dst
    src = "\n".join([f"var{i} = {i}" for i in range(1000)])
    dst = "\n".join([f"var{i} = {i*2}" for i in range(1000)])
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 274ms -> 227ms (20.4% faster)
    # All assignments should be updated to src values
    for i in range(1000):
        pass

def test_large_mixed_content():
    # Large dst with imports, assignments, functions, classes
    imports = "\n".join([f"import mod{i}" for i in range(10)])
    assigns = "\n".join([f"var{i} = {i}" for i in range(10)])
    funcs = "\n".join([f"def func{i}():\n    pass" for i in range(10)])
    classes = "\n".join([f"class C{i}:\n    pass" for i in range(10)])
    dst = "\n".join([imports, assigns, funcs, classes])
    src = "\n".join([f"var{i} = {i*2}" for i in range(5, 15)])
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 8.00ms -> 6.35ms (25.9% faster)
    # Updated assignments
    for i in range(5, 10):
        pass
    # New assignments
    for i in range(10, 15):
        pass
    # Imports, functions, classes still present
    for i in range(10):
        pass
# 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 libcst as cst
# imports
import pytest  # used for our unit tests
from codeflash.code_utils.code_extractor import add_global_assignments


class GlobalAssignmentCollector(cst.CSTVisitor):
    def __init__(self):
        self.assignments = {}
        self.assignment_order = []

    def visit_SimpleStatementLine(self, node):
        # Only collect global assignments (not inside any function/class)
        if not hasattr(node, 'parent') or not isinstance(node.parent, (cst.FunctionDef, cst.ClassDef)):
            for stmt in node.body:
                if isinstance(stmt, (cst.Assign, cst.AnnAssign, cst.AugAssign)):
                    # Support multiple targets for Assign
                    if isinstance(stmt, cst.Assign):
                        for target in stmt.targets:
                            name = target.target.value
                            self.assignments[name] = node
                            self.assignment_order.append(name)
                    elif isinstance(stmt, cst.AnnAssign):
                        name = stmt.target.value
                        self.assignments[name] = node
                        self.assignment_order.append(name)
                    elif isinstance(stmt, cst.AugAssign):
                        name = stmt.target.value
                        self.assignments[name] = node
                        self.assignment_order.append(name)

class GlobalAssignmentTransformer(cst.CSTTransformer):
    def __init__(self, assignments, assignment_order):
        self.assignments = assignments
        self.assignment_order = assignment_order
        self.found = set()

    def leave_Module(self, original_node, updated_node):
        # Remove existing assignments for these names, then add new ones at the end
        new_body = []
        assigned_names = set(self.assignments.keys())
        for stmt in updated_node.body:
            if isinstance(stmt, cst.SimpleStatementLine):
                keep = True
                for s in stmt.body:
                    if isinstance(s, (cst.Assign, cst.AnnAssign, cst.AugAssign)):
                        # Remove if target in assignments
                        if isinstance(s, cst.Assign):
                            for t in s.targets:
                                if t.target.value in assigned_names:
                                    keep = False
                        elif hasattr(s, "target") and hasattr(s.target, "value"):
                            if s.target.value in assigned_names:
                                keep = False
                if keep:
                    new_body.append(stmt)
            else:
                new_body.append(stmt)
        # Add assignments in order
        for name in self.assignment_order:
            new_body.append(self.assignments[name])
        return updated_node.with_changes(body=new_body)

class ImportInserter(cst.CSTTransformer):
    def __init__(self, stmts, last_import_line):
        self.stmts = stmts
        self.last_import_line = last_import_line

    def leave_Module(self, original_node, updated_node):
        # Insert stmts after last import
        new_body = []
        inserted = False
        for stmt in updated_node.body:
            if (
                hasattr(stmt, "start")
                and stmt.start.line == self.last_import_line
                and not inserted
            ):
                new_body.append(stmt)
                for s in self.stmts:
                    new_body.append(s)
                inserted = True
            else:
                new_body.append(stmt)
        if not inserted:
            # No imports, insert at top
            new_body = list(self.stmts) + new_body
        return updated_node.with_changes(body=new_body)

# ======================
# UNIT TESTS START HERE
# ======================

# 1. BASIC TEST CASES

def test_adds_new_global_assignment_to_empty_dst():
    # Adding a simple assignment to an empty file
    src = "x = 1"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 405μs -> 330μs (23.0% faster)

def test_adds_multiple_assignments_to_empty_dst():
    # Adding multiple assignments to an empty file
    src = "x = 1\ny = 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 551μs -> 451μs (22.1% faster)

def test_preserves_existing_dst_assignments():
    # Existing assignment in dst is not overwritten if not present in src
    src = "x = 1"
    dst = "y = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 563μs -> 457μs (23.0% faster)

def test_overwrites_existing_assignment_with_src():
    # Assignment in both src and dst: src should overwrite dst
    src = "x = 42"
    dst = "x = 1"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 532μs -> 419μs (26.8% faster)

def test_assignment_order_preserved():
    # The order of assignments in src should be preserved in result
    src = "a = 1\nb = 2\nc = 3"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 715μs -> 572μs (25.0% faster)
    lines = [line.strip() for line in result.strip().splitlines()]

def test_assignment_with_imports():
    # Global assignments should be placed after last import
    src = "x = 1"
    dst = "import os\nimport sys"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 671μs -> 537μs (25.0% faster)

def test_assignment_with_existing_import_and_assignment():
    # New assignment should be added after imports, old assignment preserved
    src = "x = 1"
    dst = "import os\ny = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 728μs -> 573μs (27.1% faster)
    # x = 1 should appear after import os
    lines = [line.strip() for line in result.strip().splitlines()]

def test_assignment_with_annassign():
    # Annotated assignment should be added/overwritten
    src = "x: int = 5"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 642μs -> 454μs (41.4% faster)

def test_augassign_handling():
    # Augmented assignment should be added/overwritten
    src = "x = 1\nx += 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 749μs -> 635μs (17.9% faster)

def test_multiple_targets_assignment():
    # Multiple targets in assignment should all be added
    src = "a = b = c = 10"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 525μs -> 435μs (20.8% faster)

# 2. EDGE TEST CASES

def test_no_global_assignments_in_src():
    # Src has no global assignments, dst should remain unchanged
    src = "def foo():\n    x = 1"
    dst = "y = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 670μs -> 408μs (64.1% faster)

def test_no_global_assignments_in_dst():
    # Dst has no global assignments, src assignments should be added
    src = "x = 1"
    dst = "def foo():\n    y = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 715μs -> 550μs (29.9% faster)

def test_assignment_inside_function_ignored():
    # Assignments inside functions should not be treated as global
    src = "def foo():\n    x = 1"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 444μs -> 308μs (44.2% faster)

def test_assignment_inside_class_ignored():
    # Assignments inside classes should not be treated as global
    src = "class A:\n    x = 1"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 418μs -> 291μs (43.7% faster)

def test_assignment_with_comments_and_blank_lines():
    # Comments and blank lines should be preserved
    src = "# comment\n\nx = 1  # inline comment"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 452μs -> 365μs (23.6% faster)

def test_duplicate_assignments_in_src():
    # Only the last assignment for a variable in src should be present
    src = "x = 1\nx = 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 519μs -> 415μs (25.0% faster)

def test_duplicate_assignments_in_dst():
    # Only the last assignment for a variable in dst should be replaced
    src = "x = 42"
    dst = "x = 1\nx = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 698μs -> 548μs (27.4% faster)

def test_imports_preserved_and_ordered():
    # Imports in dst should be preserved and assignments added after last import
    src = "x = 1"
    dst = "import sys\nimport os"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 658μs -> 527μs (24.8% faster)
    # x = 1 should be after both imports
    lines = [line.strip() for line in result.strip().splitlines()]

def test_assignment_with_unicode_and_nonascii():
    # Unicode variable names and values
    src = "π = 3.14\n变量 = '值'"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 571μs -> 460μs (24.1% faster)

def test_assignment_with_multiline_strings():
    # Multiline string assignment
    src = 'x = """hello\nworld"""'
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 366μs -> 291μs (25.4% faster)

def test_assignment_with_type_annotations_only():
    # Type annotation without value is not an assignment, should not be added
    src = "x: int"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 518μs -> 359μs (44.2% faster)

def test_assignment_with_if_main_guard():
    # Assignments inside if __name__ == "__main__" should not be global
    src = 'if __name__ == "__main__":\n    x = 1'
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 626μs -> 461μs (35.9% faster)

def test_assignment_with_trailing_whitespace_and_blank_lines():
    # Trailing whitespace and blank lines should not cause issues
    src = "x = 1   \n\ny = 2\n"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 578μs -> 464μs (24.5% faster)

def test_assignment_with_semicolon_statements():
    # Multiple assignments on one line with semicolons
    src = "x = 1; y = 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 520μs -> 423μs (22.8% faster)

def test_assignment_with_different_spacing():
    # Assignments with different spacing should be recognized as same
    src = "x=1"
    dst = "x = 2"
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 533μs -> 416μs (28.0% faster)

def test_assignment_with_tuple_unpacking():
    # Tuple unpacking assignment
    src = "a, b = 1, 2"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 554μs -> 418μs (32.5% faster)

def test_assignment_with_list_unpacking():
    # List unpacking assignment
    src = "[a, b] = [1, 2]"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 674μs -> 509μs (32.2% faster)

def test_assignment_with_complex_expression():
    # Assignment with a complex expression
    src = "x = [i for i in range(5)]"
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 811μs -> 655μs (23.7% faster)

# 3. LARGE SCALE TEST CASES

def test_large_number_of_assignments():
    # Test with a large number of assignments
    src_lines = [f"x{i} = {i}" for i in range(500)]
    dst_lines = [f"y{i} = {i}" for i in range(500)]
    src = "\n".join(src_lines)
    dst = "\n".join(dst_lines)
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 143ms -> 118ms (20.7% faster)
    for i in range(500):
        pass

def test_large_src_overwrites_large_dst():
    # Test that src assignments overwrite dst assignments for same names
    src_lines = [f"x{i} = {i}" for i in range(500)]
    dst_lines = [f"x{i} = 0" for i in range(500)]
    src = "\n".join(src_lines)
    dst = "\n".join(dst_lines)
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 133ms -> 108ms (23.0% faster)
    for i in range(500):
        pass

def test_large_with_imports():
    # Test with lots of imports and assignments
    import_lines = [f"import mod{i}" for i in range(10)]
    src_lines = [f"x{i} = {i}" for i in range(100)]
    dst_lines = import_lines + [f"y{i} = {i}" for i in range(100)]
    src = "\n".join(src_lines)
    dst = "\n".join(dst_lines)
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 30.0ms -> 25.2ms (18.7% faster)
    for i in range(10):
        pass
    for i in range(100):
        pass
    # All x assignments should appear after the last import
    result_lines = result.splitlines()
    last_import_idx = max(i for i, line in enumerate(result_lines) if line.startswith("import mod"))
    for i in range(100):
        idx = result_lines.index(f"x{i} = {i}")

def test_large_assignment_with_overlapping_names():
    # Some assignments in dst are overwritten by src, others not
    src_lines = [f"x{i} = {i}" for i in range(250, 500)]
    dst_lines = [f"x{i} = 0" for i in range(500)]
    src = "\n".join(src_lines)
    dst = "\n".join(dst_lines)
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 100ms -> 82.6ms (22.0% faster)
    for i in range(250):
        pass
    for i in range(250, 500):
        pass

def test_large_assignment_with_blank_lines_and_comments():
    # Large file with interleaved blank lines and comments
    src_lines = []
    for i in range(200):
        src_lines.append(f"# comment {i}")
        src_lines.append(f"x{i} = {i}")
        src_lines.append("")
    src = "\n".join(src_lines)
    dst = ""
    codeflash_output = add_global_assignments(src, dst); result = codeflash_output # 40.4ms -> 33.7ms (19.9% faster)
    for i in range(200):
        pass
# 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-pr683-2025-08-25T18.50.33 and push.

[misrasaurabh1]

@mohammedahmed18 can you see if this looks good? The logic seems good at a high level to me. AS you review this keep an eye our for how you would make it easier for our users to merge this code.