⚡️ Speed up function sort_chat_inputs_first by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for sort_chat_inputs_first in src/dsa/nodes.py

⏱️ Runtime : 160 microseconds → 151 microseconds (best of 448 runs)

📝 Explanation and details

Here’s an optimized rewrite.

• Reduce iteration: Collect ChatInputs locations and references in O(1) with sets and comprehensions.
• Replace repeated removals: Instead of removing while iterating, separate new layers efficiently.
• Avoid redundant list creations and extensions.
• Minimize attribute and method calls.

Here’s the improved code.

Key changes:

• Only traverse input list twice: once for checks, once for creating output.
• Avoids repeated in and remove on each layer.
• Does not mutate input lists in-place, preserving original structure.
• Skips any resulting empty layer for clarity and efficiency.

Return values and function signature are the same; comments are edited for clarity and accuracy.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import pytest
from src.dsa.nodes import sort_chat_inputs_first

# --- Test helpers ---

class DummyVertex:
    """Simple dummy vertex class with an id."""
    def __init__(self, vertex_id):
        self.vertex_id = vertex_id

class DummyGraph:
    """Dummy class to mock the required methods for sort_chat_inputs_first."""
    def __init__(self, edges=None):
        # edges: dict mapping vertex_id -> list of predecessor vertex_ids
        self.edges = edges or {}
        self.vertices = {vertex_id: DummyVertex(vertex_id) for vertex_id in self.edges}

    def get_vertex(self, vertex_id):
        # Return a dummy vertex object
        return self.vertices[vertex_id]

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

# --- Unit Tests ---

# 1. Basic Test Cases

def test_no_chat_inputs():
    # No ChatInput nodes present; output should be unchanged
    graph = DummyGraph(edges={
        "A": [],
        "B": [],
        "C": []
    })
    layers = [["A"], ["B", "C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_single_chat_input_no_deps():
    # Single ChatInput node, no dependencies; should be moved to first layer
    graph = DummyGraph(edges={
        "ChatInput1": [],
        "A": [],
        "B": []
    })
    layers = [["ChatInput1"], ["A", "B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_multiple_chat_inputs_no_deps():
    # Multiple ChatInput nodes, no dependencies; all should be moved to first layer
    graph = DummyGraph(edges={
        "ChatInput1": [],
        "ChatInput2": [],
        "A": [],
        "B": []
    })
    layers = [["ChatInput1", "A"], ["ChatInput2", "B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_chat_input_with_other_nodes():
    # ChatInput mixed with other nodes in a layer, no dependencies
    graph = DummyGraph(edges={
        "ChatInput1": [],
        "A": [],
        "B": []
    })
    layers = [["A", "ChatInput1"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

# 2. Edge Test Cases

def test_chat_input_with_dependency():
    # ChatInput node with a dependency; should NOT be moved
    graph = DummyGraph(edges={
        "ChatInput1": ["A"],
        "A": [],
        "B": []
    })
    layers = [["ChatInput1", "A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_multiple_chat_inputs_some_with_deps():
    # Some ChatInputs have dependencies, others don't; if any has deps, none are moved
    graph = DummyGraph(edges={
        "ChatInput1": [],
        "ChatInput2": ["A"],
        "A": [],
        "B": []
    })
    layers = [["ChatInput1", "A"], ["ChatInput2", "B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_all_chat_inputs_with_deps():
    # All ChatInputs have dependencies; nothing should be moved
    graph = DummyGraph(edges={
        "ChatInput1": ["B"],
        "ChatInput2": ["A"],
        "A": [],
        "B": []
    })
    layers = [["ChatInput1", "A"], ["ChatInput2", "B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_empty_layers():
    # Empty layers list; should return empty list
    graph = DummyGraph(edges={})
    layers = []
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_empty_inner_layers():
    # Layers contain empty sublists; should handle gracefully
    graph = DummyGraph(edges={})
    layers = [[], [], []]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_non_chatinput_with_chatinput_in_name():
    # Node name contains 'ChatInput' but is not a ChatInput node
    graph = DummyGraph(edges={
        "NotAChatInput": [],
        "A": []
    })
    layers = [["NotAChatInput", "A"]]
    # By the function's logic, it will treat any node with 'ChatInput' in its name as ChatInput
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

# 3. Large Scale Test Cases

def test_large_number_of_layers_and_nodes():
    # Many layers and nodes, some ChatInputs, no dependencies
    num_layers = 50
    nodes_per_layer = 10
    layers = []
    edges = {}
    # Place a ChatInput in every 5th layer
    for i in range(num_layers):
        layer = []
        for j in range(nodes_per_layer):
            node_id = f"N{i}_{j}"
            layer.append(node_id)
            edges[node_id] = []
        if i % 5 == 0:
            chat_id = f"ChatInput_{i}"
            layer.append(chat_id)
            edges[chat_id] = []
        layers.append(layer)
    graph = DummyGraph(edges=edges)
    input_layers = [layer.copy() for layer in layers]
    codeflash_output = sort_chat_inputs_first(graph, input_layers); result = codeflash_output
    # All ChatInputs should be in the first layer, order preserved
    expected_chat_inputs = [f"ChatInput_{i}" for i in range(num_layers) if i % 5 == 0]
    # All other nodes remain in their respective layers, with ChatInputs removed
    expected_layers = []
    for i in range(num_layers):
        layer = [f"N{i}_{j}" for j in range(nodes_per_layer)]
        expected_layers.append(layer)

def test_large_number_of_chatinputs_with_deps():
    # Many ChatInputs, all have dependencies; nothing should be moved
    num_layers = 20
    nodes_per_layer = 5
    layers = []
    edges = {}
    for i in range(num_layers):
        layer = []
        for j in range(nodes_per_layer):
            node_id = f"N{i}_{j}"
            layer.append(node_id)
            edges[node_id] = []
        chat_id = f"ChatInput_{i}"
        layer.append(chat_id)
        # Each ChatInput depends on one node in its layer
        edges[chat_id] = [f"N{i}_{0}"]
        layers.append(layer)
    graph = DummyGraph(edges=edges)
    input_layers = [layer.copy() for layer in layers]
    codeflash_output = sort_chat_inputs_first(graph, input_layers); result = codeflash_output

def test_large_layers_with_some_chatinputs_with_deps():
    # Some ChatInputs have deps, some don't; if any has deps, none are moved
    num_layers = 10
    nodes_per_layer = 10
    layers = []
    edges = {}
    for i in range(num_layers):
        layer = []
        for j in range(nodes_per_layer):
            node_id = f"N{i}_{j}"
            layer.append(node_id)
            edges[node_id] = []
        chat_id = f"ChatInput_{i}"
        layer.append(chat_id)
        if i == 5:
            # Only one ChatInput has a dependency
            edges[chat_id] = ["N5_0"]
        else:
            edges[chat_id] = []
        layers.append(layer)
    graph = DummyGraph(edges=edges)
    input_layers = [layer.copy() for layer in layers]
    codeflash_output = sort_chat_inputs_first(graph, input_layers); result = codeflash_output

def test_large_layers_all_chatinputs_no_deps():
    # All ChatInputs, no dependencies; all should be moved to first layer
    num_layers = 10
    layers = []
    edges = {}
    for i in range(num_layers):
        chat_id = f"ChatInput_{i}"
        edges[chat_id] = []
        layers.append([chat_id])
    graph = DummyGraph(edges=edges)
    input_layers = [layer.copy() for layer in layers]
    codeflash_output = sort_chat_inputs_first(graph, input_layers); result = codeflash_output
    expected_chat_inputs = [f"ChatInput_{i}" for i in range(num_layers)]
    expected_layers = [[] for _ in range(num_layers)]
# 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 simulate the environment for the function
class DummyGraph:
    def __init__(self, predecessors_map):
        """
        predecessors_map: dict mapping vertex_id -> list of predecessor ids
        """
        self.predecessors_map = predecessors_map

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

    def get_predecessors(self, vertex_id):
        # Returns list of predecessor ids for the given vertex
        return self.predecessors_map.get(vertex_id, [])


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

def test_no_chat_inputs():
    # No ChatInput nodes at all
    graph = DummyGraph({})
    layers = [["A", "B"], ["C"], ["D", "E"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_single_chat_input_no_dependencies():
    # One ChatInput node, no dependencies
    graph = DummyGraph({})
    layers = [["ChatInput1", "A"], ["B", "C"]]
    expected = [["ChatInput1"], ["A"], ["B", "C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_multiple_chat_inputs_no_dependencies():
    # Multiple ChatInput nodes, all with no dependencies
    graph = DummyGraph({})
    layers = [["ChatInput1", "A"], ["ChatInput2", "B"], ["C"]]
    expected = [["ChatInput1", "ChatInput2"], ["A"], ["B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_chat_input_with_dependency():
    # ChatInput node with a dependency, should NOT move
    graph = DummyGraph({"ChatInput1": ["X"]})
    layers = [["ChatInput1", "A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_mixed_chat_inputs_some_with_dependencies():
    # Some ChatInputs have dependencies, some don't
    graph = DummyGraph({
        "ChatInput1": [],
        "ChatInput2": ["A"],
    })
    layers = [["ChatInput1", "ChatInput2", "A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_chat_input_in_later_layer():
    # ChatInput appears in a later layer, no dependencies
    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

def test_chat_input_empty_layers():
    # Layers with empty lists, ChatInput present
    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

def test_no_layers():
    # Completely empty input
    graph = DummyGraph({})
    layers = []
    codeflash_output = sort_chat_inputs_first(graph, []); result = codeflash_output

def test_all_chat_inputs():
    # All nodes are ChatInputs, none with dependencies
    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

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

def test_chat_input_with_self_dependency():
    # ChatInput depends on itself (edge case)
    graph = DummyGraph({"ChatInput1": ["ChatInput1"]})
    layers = [["ChatInput1", "A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_chat_input_name_substring():
    # Node with name containing 'ChatInput' as substring but not a ChatInput
    graph = DummyGraph({})
    layers = [["NotAChatInput1", "A"], ["B"]]
    # The function only checks for substring, so will treat as ChatInput
    expected = [["NotAChatInput1"], ["A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_multiple_layers_multiple_chat_inputs():
    # ChatInputs spread across multiple layers
    graph = DummyGraph({})
    layers = [["A", "ChatInput1"], ["B", "ChatInput2"], ["C", "ChatInput3"]]
    expected = [["ChatInput1", "ChatInput2", "ChatInput3"], ["A"], ["B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_duplicate_chat_input_ids():
    # Duplicate ChatInput IDs in different layers
    graph = DummyGraph({})
    layers = [["ChatInput1", "A"], ["ChatInput1", "B"]]
    expected = [["ChatInput1", "ChatInput1"], ["A"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_layers_with_only_chat_inputs():
    # Some layers contain only ChatInputs
    graph = DummyGraph({})
    layers = [["ChatInput1"], ["ChatInput2"], ["A", "B"]]
    expected = [["ChatInput1", "ChatInput2"], ["A", "B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_layers_with_only_non_chat_inputs():
    # All layers are non-ChatInput nodes
    graph = DummyGraph({})
    layers = [["A"], ["B"], ["C"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_empty_and_nonempty_layers():
    # Mix of empty and non-empty layers, some ChatInputs
    graph = DummyGraph({})
    layers = [[], ["ChatInput1"], [], ["A", "ChatInput2"], [], ["B"]]
    expected = [["ChatInput1", "ChatInput2"], [], ["A"], [], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_chat_input_with_multiple_dependencies():
    # ChatInput with multiple dependencies
    graph = DummyGraph({"ChatInput1": ["A", "B"]})
    layers = [["A", "ChatInput1"], ["B"]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output



def test_large_number_of_layers_and_nodes():
    # Large number of layers, each with several nodes, some ChatInputs
    graph = DummyGraph({})
    num_layers = 50
    num_nodes_per_layer = 10
    layers = []
    expected_first_layer = []
    for i in range(num_layers):
        layer = []
        for j in range(num_nodes_per_layer):
            # Every 7th node is a ChatInput
            if j % 7 == 0:
                node_id = f"ChatInput_{i}_{j}"
                expected_first_layer.append(node_id)
            else:
                node_id = f"Node_{i}_{j}"
            layer.append(node_id)
        layers.append(layer)
    # Expected: all ChatInputs in first layer, rest as in original order but with ChatInputs removed
    expected_layers = [expected_first_layer]
    for i in range(num_layers):
        expected_layers.append([n for n in layers[i] if "ChatInput" not in n])
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_large_with_some_chat_inputs_having_dependencies():
    # Large input, but some ChatInputs have dependencies, so no moving
    num_layers = 20
    num_nodes_per_layer = 20
    predecessors_map = {}
    layers = []
    for i in range(num_layers):
        layer = []
        for j in range(num_nodes_per_layer):
            if j % 11 == 0:
                node_id = f"ChatInput_{i}_{j}"
                # Every 3rd ChatInput has a dependency
                if (i + j) % 3 == 0:
                    predecessors_map[node_id] = [f"Node_{i}_{j-1}"]
            else:
                node_id = f"Node_{i}_{j}"
            layer.append(node_id)
        layers.append(layer)
    graph = DummyGraph(predecessors_map)
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_large_no_chat_inputs():
    # Large input, no ChatInputs at all
    graph = DummyGraph({})
    num_layers = 100
    num_nodes_per_layer = 10
    layers = [[f"Node_{i}_{j}" for j in range(num_nodes_per_layer)] for i in range(num_layers)]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_large_all_chat_inputs():
    # All nodes are ChatInputs, none with dependencies
    graph = DummyGraph({})
    num_layers = 10
    num_nodes_per_layer = 50
    layers = [[f"ChatInput_{i}_{j}" for j in range(num_nodes_per_layer)] for i in range(num_layers)]
    expected = [[f"ChatInput_{i}_{j}" for i in range(num_layers) for j in range(num_nodes_per_layer)]]
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output

def test_large_sparse_layers():
    # Large number of sparse layers, some empty, some with ChatInputs
    graph = DummyGraph({})
    num_layers = 100
    layers = []
    expected_first_layer = []
    for i in range(num_layers):
        if i % 10 == 0:
            # Layer with ChatInput only
            node_id = f"ChatInput_{i}"
            layers.append([node_id])
            expected_first_layer.append(node_id)
        elif i % 15 == 0:
            # Empty layer
            layers.append([])
        else:
            layers.append([f"Node_{i}_A", f"Node_{i}_B"])
    expected_layers = [expected_first_layer]
    for i in range(num_layers):
        if i % 10 == 0:
            continue  # Already moved ChatInput
        elif i % 15 == 0:
            expected_layers.append([])
        else:
            expected_layers.append([f"Node_{i}_A", f"Node_{i}_B"])
    codeflash_output = sort_chat_inputs_first(graph, [layer.copy() for layer in layers]); result = codeflash_output
# 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-mb8dhm5m and push.