⚡️ Speed up function _collect_protected_ranges_for_rename by 38% in PR #1693 (fix/java/object-cast-type)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1693

If you approve this dependent PR, these changes will be merged into the original PR branch fix/java/object-cast-type.

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

---

📄 38% (0.38x) speedup for _collect_protected_ranges_for_rename in codeflash/languages/java/instrumentation.py

⏱️ Runtime : 4.44 milliseconds → 3.22 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 38% runtime improvement (from 4.44ms to 3.22ms) by eliminating recursion overhead and optimizing type lookups through two key changes:

Primary Optimization: Iterative Stack-Based Traversal

The original code used recursive calls (_collect_protected_ranges_for_rename(child, out)) which incurred Python function call overhead on every node visit. The line profiler shows this recursion line consumed 44.2% of total time (11.26ms out of 25.45ms).

The optimized version replaces recursion with an explicit stack-based iteration. This eliminates:

• Function call/return overhead for each of the 13,393 recursive calls
• Stack frame allocation and deallocation
• Parameter passing overhead

The performance gains are particularly dramatic for deep/wide trees:

• Deep nesting (100 levels): 89.4% faster (67.7μs → 35.7μs)
• Wide tree (100×100 nodes): 53.6% faster (3.03ms → 1.97ms)
• 1000 sequential nodes: 19.5% faster (527μs → 441μs)

Secondary Optimization: Module-Level Frozenset

Moving protected node types from inline tuple comparisons to a module-level _PROTECTED_RENAME_TYPES frozenset provides O(1) hash-based lookups instead of O(n) linear tuple scans. While the performance difference for 5 items is modest, it eliminates repeated tuple construction and improves cache locality.

Trade-offs

The optimization shows slight regressions (5-38% slower) on small, shallow trees due to the overhead of stack setup and iteration machinery. However, these cases are fast enough (typically <3μs) that the absolute time difference is negligible. The optimization excels where it matters most: real-world AST traversal scenarios with hundreds to thousands of nodes, where the 38% overall speedup translates to meaningful performance gains during Java code instrumentation workflows.

The iterative approach also has better memory characteristics for deep trees, avoiding potential stack overflow issues with Python's recursion limit.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import pytest  # used for our unit tests
from codeflash.languages.java.instrumentation import \
    _collect_protected_ranges_for_rename

class TestNode:
    """A minimal concrete node object with the attributes used by the function.

    We use a real class with an explicit constructor to ensure deterministic,
    real-instance behavior when passed into the function under test.
    """

    def __init__(self, type: str, start_byte: int = 0, end_byte: int = 0, children=None):
        # store the node type (string that the function checks)
        self.type = type
        # start and end bytes that the function will append into 'out'
        self.start_byte = start_byte
        self.end_byte = end_byte
        # children should be an iterable of other TestNode instances
        self.children = list(children) if children is not None else []

    def __repr__(self):
        # helpful repr for debugging failing tests
        return f"TestNode(type={self.type!r}, start={self.start_byte}, end={self.end_byte}, children={len(self.children)})"

def test_collects_range_for_string_literal_node():
    # create a single node that is a string literal; children empty
    node = TestNode("string_literal", start_byte=10, end_byte=20, children=[])
    out = []  # collector list passed to the function
    # call the function under test
    _collect_protected_ranges_for_rename(node, out) # 1.05μs -> 1.55μs (32.2% slower)

def test_collects_range_for_import_declaration_node():
    # import declarations are protected and cause early return
    node = TestNode("import_declaration", start_byte=0, end_byte=50)
    out = []
    _collect_protected_ranges_for_rename(node, out) # 951ns -> 1.04μs (8.65% slower)

def test_non_protected_node_with_no_children_leaves_out_ranges():
    # a node with a type not in the protected list and no children should produce no entries
    node = TestNode("class_declaration", start_byte=5, end_byte=15, children=[])
    out = []
    _collect_protected_ranges_for_rename(node, out) # 871ns -> 1.19μs (26.9% slower)

def test_nested_children_collection_deep_and_shallow():
    # build a tree:
    # root (non-protected)
    #  ├─ child1 (non-protected)
    #  │    └─ grandchild1 (string_literal)
    #  └─ child2 (line_comment)
    grandchild1 = TestNode("string_literal", start_byte=30, end_byte=40)
    child1 = TestNode("method_declaration", start_byte=20, end_byte=45, children=[grandchild1])
    child2 = TestNode("line_comment", start_byte=46, end_byte=60)
    root = TestNode("class_declaration", start_byte=0, end_byte=100, children=[child1, child2])

    out = []
    _collect_protected_ranges_for_rename(root, out) # 2.02μs -> 3.01μs (32.9% slower)

def test_protected_parent_prevents_traversal_of_children():
    # If a parent node is protected, the function should append the parent's range and return,
    # i.e., it should NOT traverse into the parent's children.
    # parent is a protected string_literal but also has a child that is another protected node;
    # only the parent range should be recorded.
    child = TestNode("line_comment", start_byte=100, end_byte=110)
    parent = TestNode("string_literal", start_byte=90, end_byte=120, children=[child])
    out = []
    _collect_protected_ranges_for_rename(parent, out) # 711ns -> 1.11μs (36.0% slower)

def test_various_literal_and_comment_types_collected():
    # Test collection for all protected types mentioned in the function docstring
    protected_nodes = [
        TestNode("string_literal", start_byte=0, end_byte=1),
        TestNode("character_literal", start_byte=2, end_byte=3),
        TestNode("line_comment", start_byte=4, end_byte=5),
        TestNode("block_comment", start_byte=6, end_byte=7),
    ]
    root = TestNode("compilation_unit", start_byte=0, end_byte=10, children=protected_nodes)
    out = []
    _collect_protected_ranges_for_rename(root, out) # 2.04μs -> 2.67μs (23.6% slower)
    # Expect one entry for each protected child in the same order they appear
    expected = [(0, 1), (2, 3), (4, 5), (6, 7)]

def test_missing_attributes_raises_attribute_error():
    # Passing an object that lacks the expected attributes should raise an AttributeError.
    # This verifies that the function relies on the presence of these attributes.
    class Bare:
        # intentionally minimal; no 'type' or 'children' attr
        pass

    with pytest.raises(AttributeError):
        _collect_protected_ranges_for_rename(Bare(), []) # 3.19μs -> 3.55μs (10.2% slower)

def test_large_scale_many_leaves_collection_performance_and_correctness():
    # Construct a large root with 1000 protected leaves to ensure the function scales
    # and collects all ranges. We use 1000 to test performance without being excessive.
    leaf_count = 1000
    # Create leaves with consecutive byte ranges to make assertions straightforward
    leaves = [TestNode("string_literal", start_byte=i * 10, end_byte=i * 10 + 5) for i in range(leaf_count)]
    root = TestNode("compilation_unit", start_byte=0, end_byte=leaf_count * 10 + 10, children=leaves)

    out = []
    _collect_protected_ranges_for_rename(root, out) # 137μs -> 124μs (10.5% faster)

def test_collection_with_duplicate_and_overlapping_ranges():
    # Test that overlapping and duplicate ranges are preserved (function just appends tuples)
    overlap_child1 = TestNode("string_literal", start_byte=0, end_byte=10)
    overlap_child2 = TestNode("line_comment", start_byte=5, end_byte=15)
    duplicate_child = TestNode("block_comment", start_byte=0, end_byte=10)
    root = TestNode("compilation_unit", start_byte=0, end_byte=20, children=[overlap_child1, overlap_child2, duplicate_child])

    out = []
    _collect_protected_ranges_for_rename(root, out) # 1.85μs -> 2.65μs (30.2% slower)
    # Expect 3 entries preserved in that order
    expected = [(0, 10), (5, 15), (0, 10)]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from unittest.mock import Mock

# imports
import pytest
from codeflash.languages.java.instrumentation import \
    _collect_protected_ranges_for_rename

# Create mock AST node objects for testing
def create_mock_node(node_type, start_byte, end_byte, children=None):
    """Helper to create a mock AST node with the required attributes."""
    node = Mock()
    node.type = node_type
    node.start_byte = start_byte
    node.end_byte = end_byte
    node.children = children if children is not None else []
    return node

def test_string_literal_protection():
    """Test that string literals are collected as protected ranges."""
    # Create a string literal node at bytes 10-20
    string_node = create_mock_node("string_literal", 10, 20)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(string_node, out) # 902ns -> 1.37μs (34.3% slower)

def test_character_literal_protection():
    """Test that character literals are collected as protected ranges."""
    # Create a character literal node at bytes 5-8
    char_node = create_mock_node("character_literal", 5, 8)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(char_node, out) # 911ns -> 1.19μs (23.6% slower)

def test_line_comment_protection():
    """Test that line comments are collected as protected ranges."""
    # Create a line comment node at bytes 100-150
    comment_node = create_mock_node("line_comment", 100, 150)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(comment_node, out) # 861ns -> 1.16μs (25.9% slower)

def test_block_comment_protection():
    """Test that block comments are collected as protected ranges."""
    # Create a block comment node at bytes 200-300
    block_comment = create_mock_node("block_comment", 200, 300)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(block_comment, out) # 851ns -> 1.10μs (22.8% slower)

def test_import_declaration_protection():
    """Test that import declarations are collected as protected ranges."""
    # Create an import declaration node at bytes 0-40
    import_node = create_mock_node("import_declaration", 0, 40)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(import_node, out) # 882ns -> 1.10μs (20.0% slower)

def test_regular_identifier_not_protected():
    """Test that regular identifiers (class names) are NOT collected."""
    # Create a regular identifier node (not a protected type)
    identifier_node = create_mock_node("identifier", 50, 60)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(identifier_node, out) # 842ns -> 1.07μs (21.4% slower)

def test_method_declaration_not_protected():
    """Test that method declarations are NOT immediately protected."""
    # Create a method declaration node
    method_node = create_mock_node("method_declaration", 100, 200)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(method_node, out) # 772ns -> 961ns (19.7% slower)

def test_single_child_traversal():
    """Test that the function traverses child nodes correctly."""
    # Create a protected child node
    protected_child = create_mock_node("string_literal", 20, 30)
    
    # Create a parent node with the child
    parent = create_mock_node("parent_type", 0, 50, children=[protected_child])
    out = []
    
    # Call the function on parent
    _collect_protected_ranges_for_rename(parent, out) # 1.30μs -> 2.12μs (38.7% slower)

def test_multiple_children_traversal():
    """Test that the function correctly processes multiple child nodes."""
    # Create several child nodes, some protected, some not
    string_child = create_mock_node("string_literal", 10, 20)
    identifier_child = create_mock_node("identifier", 25, 35)
    comment_child = create_mock_node("line_comment", 40, 50)
    
    # Create a parent with these children
    parent = create_mock_node("parent", 0, 60, children=[string_child, identifier_child, comment_child])
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(parent, out) # 2.09μs -> 2.88μs (27.5% slower)

def test_nested_protection_with_depth_two():
    """Test protection collection with two levels of nesting."""
    # Create a protected grandchild
    grandchild = create_mock_node("string_literal", 100, 110)
    
    # Create an intermediate parent
    child = create_mock_node("intermediate", 90, 120, children=[grandchild])
    
    # Create a root node
    root = create_mock_node("root", 0, 200, children=[child])
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 1.68μs -> 2.45μs (31.4% slower)

def test_empty_children_list():
    """Test handling of a node with an empty children list."""
    # Create a node with empty children
    node = create_mock_node("some_type", 0, 50, children=[])
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 761ns -> 1.02μs (25.5% slower)

def test_zero_byte_range():
    """Test handling of protected ranges with zero-width byte range."""
    # Create a string literal with start_byte == end_byte
    node = create_mock_node("string_literal", 100, 100)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 851ns -> 1.16μs (26.8% slower)

def test_very_large_byte_range():
    """Test handling of very large byte ranges."""
    # Create a string literal with a very large range
    node = create_mock_node("string_literal", 0, 1000000)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 721ns -> 921ns (21.7% slower)

def test_negative_byte_positions():
    """Test handling of negative byte positions (edge case)."""
    # Note: This tests if the function handles unusual input gracefully
    # In practice, byte positions should be non-negative, but we test robustness
    node = create_mock_node("string_literal", -10, -5)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 761ns -> 1.03μs (26.3% slower)

def test_single_protected_type_only():
    """Test that ONLY string_literal is protected (not other string types)."""
    # Create a node with 'string' type (not 'string_literal')
    node = create_mock_node("string", 10, 20)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 761ns -> 961ns (20.8% slower)

def test_case_sensitive_type_matching():
    """Test that type matching is case-sensitive."""
    # Create a node with capitalized type name (wrong case)
    node = create_mock_node("String_Literal", 10, 20)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 851ns -> 901ns (5.55% slower)

def test_deeply_nested_protection():
    """Test protection collection with deep nesting (10+ levels)."""
    # Start with a protected leaf node
    current = create_mock_node("string_literal", 500, 510)
    
    # Build 15 levels of nesting
    for i in range(15):
        current = create_mock_node(f"parent_{i}", 400 - i*10, 520 + i*10, children=[current])
    
    out = []
    
    # Call the function on root
    _collect_protected_ranges_for_rename(current, out) # 4.89μs -> 6.45μs (24.2% slower)

def test_out_list_is_mutated():
    """Test that the function mutates the provided out list in place."""
    # Create a node
    node = create_mock_node("string_literal", 10, 20)
    
    # Create an out list with existing content
    out = [(0, 5)]
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 912ns -> 1.22μs (25.4% slower)

def test_multiple_protected_types_in_tree():
    """Test collection when multiple different protected types are in the tree."""
    # Create different protected node types
    string_node = create_mock_node("string_literal", 10, 20)
    char_node = create_mock_node("character_literal", 30, 40)
    comment_node = create_mock_node("line_comment", 50, 60)
    import_node = create_mock_node("import_declaration", 70, 80)
    
    # Create a root with all these children
    root = create_mock_node("root", 0, 100, children=[string_node, char_node, comment_node, import_node])
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 2.49μs -> 3.13μs (20.5% slower)

def test_mixed_protected_and_unprotected_siblings():
    """Test handling of siblings where some are protected and some are not."""
    # Create mixed children
    protected_1 = create_mock_node("string_literal", 10, 20)
    unprotected_1 = create_mock_node("class_declaration", 25, 35)
    protected_2 = create_mock_node("block_comment", 40, 50)
    unprotected_2 = create_mock_node("method_invocation", 55, 65)
    protected_3 = create_mock_node("import_declaration", 70, 80)
    
    # Create root with these mixed children
    root = create_mock_node("root", 0, 100, children=[
        protected_1, unprotected_1, protected_2, unprotected_2, protected_3
    ])
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 2.81μs -> 3.01μs (6.64% slower)

def test_reversed_byte_range():
    """Test handling when end_byte < start_byte (invalid but test robustness)."""
    # Create a node with reversed byte range
    node = create_mock_node("string_literal", 100, 50)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(node, out) # 752ns -> 1.15μs (34.7% slower)

def test_overlapping_protected_ranges():
    """Test collection when protected ranges overlap."""
    # Create overlapping protected nodes
    node1 = create_mock_node("string_literal", 10, 50)
    node2 = create_mock_node("line_comment", 30, 70)
    
    # Create root with both children
    root = create_mock_node("root", 0, 100, children=[node1, node2])
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 1.96μs -> 2.66μs (26.3% slower)

def test_large_number_of_siblings():
    """Test handling of a parent with 1000 child nodes."""
    # Create 1000 children, half protected and half not
    children = []
    for i in range(500):
        # Add a protected node
        children.append(create_mock_node("string_literal", i*4, i*4+2))
        # Add an unprotected node
        children.append(create_mock_node("identifier", i*4+2, i*4+4))
    
    # Create root with all children
    root = create_mock_node("root", 0, 2000, children=children)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 527μs -> 441μs (19.5% faster)

def test_deep_nesting_with_protected_at_bottom():
    """Test handling of a tree nested 100 levels deep."""
    # Start with a protected leaf node
    current = create_mock_node("string_literal", 5000, 5010)
    
    # Build 100 levels of nesting
    for i in range(100):
        current = create_mock_node(f"parent_{i}", 5000 - i, 5010 + i, children=[current])
    
    out = []
    
    # Call the function on root
    _collect_protected_ranges_for_rename(current, out) # 67.7μs -> 35.7μs (89.4% faster)

def test_wide_tree_with_many_children_at_each_level():
    """Test handling of a tree with 100 children, each having 100 children."""
    # Create a 2-level tree with many branches
    grandchildren = []
    for j in range(100):
        # Each grandchild is either protected or not
        if j % 2 == 0:
            grandchildren.append(create_mock_node("string_literal", j*2, j*2+1))
        else:
            grandchildren.append(create_mock_node("identifier", j*2, j*2+1))
    
    children = []
    for i in range(100):
        # Create intermediate children, each with grandchildren
        children.append(create_mock_node(f"child_{i}", i*200, i*200+100, children=grandchildren[:]))
    
    # Create root
    root = create_mock_node("root", 0, 20000, children=children)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 3.03ms -> 1.97ms (53.6% faster)

def test_many_protected_nodes_in_sequence():
    """Test collection of 1000 sequential protected nodes."""
    # Create 1000 protected nodes with non-overlapping ranges
    children = []
    for i in range(1000):
        children.append(create_mock_node("string_literal", i*10, i*10+5))
    
    # Create root
    root = create_mock_node("root", 0, 10000, children=children)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 564μs -> 530μs (6.52% faster)

def test_all_import_declarations_in_large_file():
    """Test collecting all import declarations from a simulated large file."""
    # Simulate a Java file with 100 import statements at the top
    children = []
    for i in range(100):
        children.append(create_mock_node("import_declaration", i*50, i*50+40))
    
    # Add some other declarations after imports
    for i in range(50):
        children.append(create_mock_node("class_declaration", 5000 + i*100, 5000 + i*100+80))
    
    # Create root
    root = create_mock_node("root", 0, 10000, children=children)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out) # 64.5μs -> 51.0μs (26.4% faster)

def test_performance_with_complex_nested_structure():
    """Test performance with a complex tree having both depth and width."""
    # Build a tree that's both deep and wide: depth 10, width 50 at each level
    def build_tree(depth, width, byte_offset):
        if depth == 0:
            # Leaf nodes - alternate between protected and unprotected
            children = []
            for i in range(width):
                if i % 3 == 0:
                    children.append(create_mock_node("string_literal", byte_offset + i*2, byte_offset + i*2+1))
                else:
                    children.append(create_mock_node("identifier", byte_offset + i*2, byte_offset + i*2+1))
            return children
        else:
            # Internal nodes
            children = []
            for i in range(width):
                subtree = build_tree(depth - 1, width, byte_offset + i*1000)
                children.append(create_mock_node(f"internal_{depth}_{i}", byte_offset + i*1000, 
                                                byte_offset + i*1000 + 500, children=subtree))
            return children
    
    root_children = build_tree(5, 10, 0)
    root = create_mock_node("root", 0, 100000, children=root_children)
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(root, out)

def test_single_path_deep_nesting_with_alternating_protection():
    """Test a single deep path where protection alternates at each level."""
    # Build a single chain of 100 nodes where every other one is protected
    current = create_mock_node("string_literal", 10000, 10010) if 0 % 2 == 0 else create_mock_node("identifier", 10000, 10010)
    
    for i in range(1, 100):
        node_type = "string_literal" if i % 2 == 0 else "identifier"
        current = create_mock_node(node_type, 10000 - i*10, 10010 + i*10, children=[current])
    
    out = []
    
    # Call the function
    _collect_protected_ranges_for_rename(current, out) # 2.79μs -> 3.60μs (22.3% slower)
# 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-pr1693-2026-02-28T01.17.48 and push.