⚡️ Speed up function sort_chat_inputs_first by 16%

[codeflash-ai]

📄 16% (0.16x) speedup for sort_chat_inputs_first in src/dsa/nodes.py

⏱️ Runtime : 350 microseconds → 301 microseconds (best of 712 runs)

📝 Explanation and details

The optimized code achieves a 16% speedup through several key algorithmic and structural improvements:

Key Optimizations Applied:

1. Eliminated Redundant List Operations: The original code used layer.remove(vertex_id) which is O(n) for each removal, requiring list shifting. The optimized version builds new layers using list comprehensions [vid for vid in layer if "ChatInput" not in vid], avoiding expensive in-place mutations.

2. Reduced Dependency Checking: The original code checked dependencies with "ChatInput" in vertex_id and self.get_predecessors(...) in a single condition, causing short-circuit evaluation issues. The optimized version separates the string check from dependency checking, only calling expensive graph operations when necessary.

3. Streamlined Data Flow: Instead of first collecting ChatInputs in chat_inputs_first, then extending it, and finally removing from original layers, the optimized version processes everything in a single pass - collecting ChatInputs while immediately checking dependencies, then rebuilding layers without ChatInputs.

4. Eliminated Intermediate Collections: The original code created layer_chat_inputs_first for each layer and used extend() operations. The optimized version directly appends to chatinputs_ids and builds the final result structure more efficiently.

Why These Changes Improve Performance:

• List.remove() elimination: Each remove() call is O(n) and requires shifting elements. With multiple ChatInputs per layer, this becomes expensive. List comprehensions are more cache-efficient and avoid memory moves.
• Better short-circuiting: Early return on first dependency found prevents unnecessary processing of remaining ChatInputs.
• Reduced function call overhead: Fewer intermediate list operations and method calls reduce the per-operation overhead.

Test Case Performance Patterns:

The optimization performs best on:

• Large datasets with no ChatInputs (117% faster): Avoids expensive string checking and graph operations entirely
• Scenarios with many ChatInputs but no dependencies (18-26% faster): Benefits from elimination of list.remove() operations
• Empty or sparse layers (18-20% faster): Reduced overhead in layer processing

The optimization performs worse on small test cases with dependencies because the additional upfront setup (creating collections) has overhead that isn't amortized over enough work, but the algorithmic improvements shine on larger inputs where the O(n) operations in the original become bottlenecks.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import List

# imports
import pytest  # used for our unit tests
from src.dsa.nodes import sort_chat_inputs_first


# Helper class to simulate the required methods and graph structure
class MockGraph:
    def __init__(self, vertices_layers, predecessors_map):
        """
        vertices_layers: List[List[str]] - layers of vertex ids
        predecessors_map: Dict[str, List[str]] - mapping from vertex_id to its predecessors
        """
        self.vertices_layers = vertices_layers
        self.predecessors_map = predecessors_map

    def get_vertex(self, vertex_id):
        # In this mock, vertex_id is used directly
        return vertex_id

    def get_predecessors(self, vertex_id):
        # Return the list of predecessors for the given vertex_id
        return self.predecessors_map.get(vertex_id, [])


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

def test_single_layer_no_chatinput():
    # No ChatInput present
    graph = MockGraph([["A", "B", "C"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 458ns -> 458ns (0.000% faster)

def test_single_layer_with_chatinput_no_dependencies():
    # ChatInput present, no dependencies
    graph = MockGraph([["A", "ChatInput1", "B"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 833ns -> 875ns (4.80% slower)

def test_multiple_layers_chatinput_in_first_layer():
    # ChatInput in first layer, no dependencies
    graph = MockGraph([["ChatInput1", "A"], ["B", "C"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 916ns -> 1.04μs (12.0% slower)

def test_multiple_layers_chatinput_in_second_layer():
    # ChatInput in second layer, no dependencies
    graph = MockGraph([["A"], ["ChatInput1", "B"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 875ns -> 958ns (8.66% slower)

def test_multiple_chatinputs_no_dependencies():
    # Multiple ChatInputs, no dependencies
    graph = MockGraph([["A", "ChatInput1"], ["ChatInput2", "B"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 1.04μs -> 1.12μs (7.38% slower)

def test_chatinput_with_dependencies():
    # ChatInput has dependencies, should not move
    predecessors_map = {"ChatInput1": ["A"]}
    graph = MockGraph([["A"], ["ChatInput1", "B"]], predecessors_map)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 417ns -> 708ns (41.1% slower)

def test_chatinput_with_dependencies_among_multiple():
    # One ChatInput with dependencies, one without
    predecessors_map = {"ChatInput2": ["A"]}
    graph = MockGraph([["A", "ChatInput1"], ["ChatInput2", "B"]], predecessors_map)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 583ns -> 833ns (30.0% slower)

# -------------------
# EDGE TEST CASES
# -------------------

def test_empty_layers():
    # Empty input
    graph = MockGraph([], {})
    codeflash_output = sort_chat_inputs_first(graph, []); result = codeflash_output # 167ns -> 292ns (42.8% slower)

def test_layers_with_empty_lists():
    # Layers are empty lists
    graph = MockGraph([[], [], []], {})
    codeflash_output = sort_chat_inputs_first(graph, [[], [], []]); result = codeflash_output # 500ns -> 416ns (20.2% faster)

def test_chatinput_with_self_dependency():
    # ChatInput depends on itself (cycle)
    predecessors_map = {"ChatInput1": ["ChatInput1"]}
    graph = MockGraph([["ChatInput1"]], predecessors_map)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 375ns -> 625ns (40.0% slower)

def test_chatinput_with_nonexistent_dependency():
    # ChatInput depends on a vertex not in the layers
    predecessors_map = {"ChatInput1": ["Z"]}
    graph = MockGraph([["ChatInput1"]], predecessors_map)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 375ns -> 583ns (35.7% slower)

def test_chatinput_in_all_layers():
    # ChatInputs in every layer, no dependencies
    graph = MockGraph([["ChatInput1"], ["ChatInput2"], ["ChatInput3"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 1.17μs -> 1.17μs (0.000% faster)

def test_layers_with_only_chatinputs():
    # All layers only contain ChatInputs, no dependencies
    graph = MockGraph([["ChatInput1"], ["ChatInput2"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 917ns -> 958ns (4.28% slower)

def test_layers_with_duplicate_chatinputs():
    # Duplicate ChatInputs in different layers (should not happen, but test anyway)
    graph = MockGraph([["ChatInput1"], ["ChatInput1"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 916ns -> 916ns (0.000% faster)

def test_chatinput_with_empty_string():
    # Layer contains empty string and ChatInput
    graph = MockGraph([["", "ChatInput1"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 750ns -> 875ns (14.3% slower)

def test_chatinput_with_special_characters():
    # ChatInput with special characters in name
    graph = MockGraph([["A", "ChatInput$#@!"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 792ns -> 917ns (13.6% slower)

def test_chatinput_substring_in_name():
    # Vertex that contains 'ChatInput' as substring but is not a ChatInput
    graph = MockGraph([["A", "NotAChatInputButHasChatInputInName"]], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in graph.vertices_layers]); result = codeflash_output # 791ns -> 875ns (9.60% slower)

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

def test_large_number_of_layers_and_vertices():
    # 100 layers, each with 10 vertices, ChatInput in every 10th layer
    layers = []
    predecessors_map = {}
    chatinputs = []
    for i in range(100):
        layer = [f"V{i}_{j}" for j in range(10)]
        if i % 10 == 0:
            chatinput = f"ChatInput{i}"
            layer.append(chatinput)
            chatinputs.append(chatinput)
        layers.append(layer)
    graph = MockGraph([layer.copy() for layer in layers], predecessors_map)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 24.0μs -> 26.7μs (10.0% slower)
    # All ChatInputs should be in the first layer, order preserved
    expected_first_layer = chatinputs
    # Remove ChatInputs from their original layers for comparison
    expected_layers = []
    for i, layer in enumerate(layers):
        expected_layers.append([v for v in layer if "ChatInput" not in v])

def test_large_number_of_chatinputs_with_dependencies():
    # 100 ChatInputs, all have dependencies
    layers = []
    predecessors_map = {}
    for i in range(10):
        layer = []
        for j in range(10):
            chatinput = f"ChatInput{i}_{j}"
            layer.append(chatinput)
            # Each ChatInput depends on a dummy vertex
            predecessors_map[chatinput] = [f"Dummy{i}_{j}"]
        layers.append(layer)
    graph = MockGraph([layer.copy() for layer in layers], predecessors_map)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 417ns -> 583ns (28.5% slower)

def test_large_layers_no_chatinputs():
    # 100 layers, 10 vertices each, no ChatInputs
    layers = [[f"V{i}_{j}" for j in range(10)] for i in range(100)]
    graph = MockGraph([layer.copy() for layer in layers], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 20.3μs -> 9.88μs (105% faster)

def test_large_layers_all_chatinputs():
    # 10 layers, each with 100 ChatInputs, no dependencies
    layers = []
    chatinputs = []
    for i in range(10):
        layer = [f"ChatInput{i}_{j}" for j in range(100)]
        chatinputs.extend(layer)
        layers.append(layer)
    graph = MockGraph([layer.copy() for layer in layers], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 112μs -> 88.8μs (26.3% faster)
    # All other layers should be empty
    for l in result[1:]:
        pass

def test_performance_with_max_elements():
    # 100 layers, each with 10 vertices, 10 ChatInputs scattered, no dependencies
    layers = []
    chatinputs = []
    for i in range(100):
        layer = [f"V{i}_{j}" for j in range(10)]
        if i < 10:
            chatinput = f"ChatInput{i}"
            layer.append(chatinput)
            chatinputs.append(chatinput)
        layers.append(layer)
    graph = MockGraph([layer.copy() for layer in layers], {})
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 23.6μs -> 25.8μs (8.71% slower)
    # All other layers should not contain ChatInputs
    for l in result[1:]:
        for v in l:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from src.dsa.nodes import sort_chat_inputs_first


# Helper class to mock self with get_vertex and get_predecessors
class DummyGraph:
    def __init__(self, vertex_dependencies):
        """
        vertex_dependencies: dict mapping vertex_id -> list of predecessor vertex_ids
        """
        self.vertex_dependencies = vertex_dependencies

    def get_vertex(self, vertex_id):
        # In a real graph, this would return a vertex object; here, just return the id
        return vertex_id

    def get_predecessors(self, vertex_id):
        # Return the list of predecessors for the given vertex_id
        return self.vertex_dependencies.get(vertex_id, [])


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

def test_no_chat_inputs():
    # No ChatInput in any layer; should return unchanged
    graph = DummyGraph({})
    layers = [["A", "B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 583ns -> 458ns (27.3% faster)

def test_single_chat_input_first_layer():
    # ChatInput in first layer; should stay in first layer
    graph = DummyGraph({})
    layers = [["ChatInput1", "A"], ["B"]]
    expected = [["ChatInput1"], ["A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 875ns -> 958ns (8.66% slower)

def test_single_chat_input_second_layer():
    # ChatInput in second layer; should be moved to first layer
    graph = DummyGraph({})
    layers = [["A"], ["ChatInput1", "B"], ["C"]]
    expected = [["ChatInput1"], ["A"], ["B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 958ns -> 1.08μs (11.5% slower)

def test_multiple_chat_inputs_multiple_layers():
    # Multiple ChatInputs in different layers; all should be moved to first layer
    graph = DummyGraph({})
    layers = [["A"], ["ChatInput1", "B"], ["ChatInput2", "C"]]
    expected = [["ChatInput1", "ChatInput2"], ["A"], ["B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 1.12μs -> 1.21μs (6.95% slower)

def test_chat_input_with_dependency():
    # ChatInput has a dependency; should NOT move anything
    graph = DummyGraph({"ChatInput1": ["A"]})
    layers = [["A"], ["ChatInput1", "B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 416ns -> 667ns (37.6% slower)

def test_non_chat_input_with_dependency():
    # Non-ChatInput has a dependency; ChatInput should still be moved
    graph = DummyGraph({"A": ["B"]})
    layers = [["A"], ["ChatInput1", "B"], ["C"]]
    expected = [["ChatInput1"], ["A"], ["B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 958ns -> 1.04μs (7.97% slower)

# -------------------- EDGE TEST CASES --------------------

def test_empty_layers():
    # No layers at all
    graph = DummyGraph({})
    layers = []
    codeflash_output = sort_chat_inputs_first(graph, []); result = codeflash_output # 167ns -> 291ns (42.6% slower)

def test_layers_with_empty_lists():
    # Layers are present but some are empty
    graph = DummyGraph({})
    layers = [[], ["ChatInput1"], [], ["A", "ChatInput2"], []]
    expected = [["ChatInput1", "ChatInput2"], [], [], ["A"], []]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 1.38μs -> 1.38μs (0.000% faster)

def test_all_chat_inputs():
    # All vertices are ChatInputs; all should be in first layer
    graph = DummyGraph({})
    layers = [["ChatInput1", "ChatInput2"], ["ChatInput3"]]
    expected = [["ChatInput1", "ChatInput2", "ChatInput3"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 1.21μs -> 1.08μs (11.5% faster)

def test_chat_input_duplicate_names():
    # Duplicate ChatInput IDs in different layers; all should be moved to first
    graph = DummyGraph({})
    layers = [["ChatInput1"], ["ChatInput1", "A"], ["ChatInput2"]]
    expected = [["ChatInput1", "ChatInput1", "ChatInput2"], ["A"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 1.21μs -> 1.29μs (6.42% slower)

def test_chat_input_with_self_dependency():
    # ChatInput depends on itself; should NOT move
    graph = DummyGraph({"ChatInput1": ["ChatInput1"]})
    layers = [["ChatInput1", "A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 375ns -> 584ns (35.8% slower)

def test_chat_input_with_circular_dependency():
    # ChatInput depends on B, B depends on ChatInput; should NOT move
    graph = DummyGraph({"ChatInput1": ["B"], "B": ["ChatInput1"]})
    layers = [["ChatInput1"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 375ns -> 583ns (35.7% slower)


def test_chat_input_substring_in_name():
    # Vertex ID contains 'ChatInput' as a substring (e.g., 'MyChatInputNode')
    graph = DummyGraph({})
    layers = [["MyChatInputNode", "A"], ["ChatInput2"]]
    expected = [["MyChatInputNode", "ChatInput2"], ["A"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 1.12μs -> 1.21μs (6.87% slower)

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

def test_large_number_of_layers_and_vertices():
    # 100 layers, each with 10 vertices; every 10th vertex is a ChatInput
    graph = DummyGraph({})
    layers = []
    chat_inputs = []
    for i in range(100):
        layer = []
        for j in range(10):
            idx = i * 10 + j
            if j == 0:
                vid = f"ChatInput{idx}"
                chat_inputs.append(vid)
            else:
                vid = f"V{idx}"
            layer.append(vid)
        layers.append(layer)
    expected_layers = [chat_inputs] + [[f"V{i*10+j}" for j in range(1,10)] for i in range(100)]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 31.5μs -> 34.0μs (7.12% slower)

def test_large_with_some_chat_inputs_with_dependencies():
    # 50 ChatInputs, 10 have dependencies; only those without dependencies are moved
    dependencies = {f"ChatInput{i}": ["V0"] for i in range(10)}  # first 10 ChatInputs have dependencies
    graph = DummyGraph(dependencies)
    layers = []
    chat_inputs = []
    for i in range(50):
        if i < 10:
            layer = [f"ChatInput{i}", f"V{i}"]
        else:
            layer = [f"ChatInput{i}"]
            chat_inputs.append(f"ChatInput{i}")
        layers.append(layer)
    # Since some ChatInputs have dependencies, function should return as-is (since *any* ChatInput has dependency)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 417ns -> 625ns (33.3% slower)

def test_large_all_chat_inputs_no_dependencies():
    # 500 ChatInputs, no dependencies; all should be in first layer
    graph = DummyGraph({})
    layers = [[f"ChatInput{i}"] for i in range(500)]
    expected = [[f"ChatInput{i}" for i in range(500)]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 84.0μs -> 70.8μs (18.6% faster)

def test_large_no_chat_inputs():
    # 1000 vertices, no ChatInputs; should return unchanged
    graph = DummyGraph({})
    layers = [[f"V{i*10+j}" for j in range(10)] for i in range(100)]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 21.2μs -> 9.79μs (117% faster)

def test_large_sparse_layers():
    # 100 layers, each with 0 or 1 ChatInput, rest empty
    graph = DummyGraph({})
    layers = []
    chat_inputs = []
    for i in range(100):
        if i % 10 == 0:
            layers.append([f"ChatInput{i}"])
            chat_inputs.append(f"ChatInput{i}")
        else:
            layers.append([])
    expected = [chat_inputs] + [[] for _ in range(99)]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 7.75μs -> 6.54μs (18.5% faster)

# -------------------- MUTATION DEFENSE TESTS --------------------

def test_mutation_defense_no_side_effects():
    # Ensure function does not mutate input layers (should not mutate original lists)
    graph = DummyGraph({})
    layers = [["ChatInput1", "A"], ["B"]]
    layers_copy = [layer.copy() for layer in layers]
    sort_chat_inputs_first(graph, [layer.copy() for layer in layers]) # 875ns -> 1.00μs (12.5% slower)

def test_mutation_defense_returns_new_list():
    # Ensure function returns a new list, not the same object
    graph = DummyGraph({})
    layers = [["ChatInput1", "A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output # 875ns -> 917ns (4.58% 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-sort_chat_inputs_first-mdpdgqw8 and push.