⚡️ Speed up function find_last_node by 7,511%

[codeflash-ai]

📄 7,511% (75.11x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 30.4 milliseconds → 399 microseconds (best of 250 runs)

📝 Explanation and details

Refined the optimization to improve code quality by:

1. Removed single-use iterator handling: The complex logic for detecting and handling single-use iterators adds 20+ lines of code for minimal benefit. The profiler shows the re-iterable path is what gets executed in practice.

2. Removed isinstance micro-optimization: The isinstance(edges, list) check before converting to list is unnecessary complexity that doesn't materially affect performance.

3. Restored generator expression: Changed the explicit for-loop back to the original's generator expression style with next(), which is more Pythonic and maintains the original code's elegance.

4. Preserved core optimization: Kept the key performance improvement - converting edges to a set of sources for O(1) membership testing instead of O(m) all() checks for each node. This is the optimization that delivers the 85x speedup.

5. Preserved empty edges behavior: Maintained the important behavioral fix where empty edges returns the first node without accessing n["id"].

The refined code is now much more concise (5 lines vs 30+ lines) while delivering the same performance improvement, making it easier to review and accept.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from src.algorithms.graph import find_last_node

# unit tests

# Basic Test Cases


def test_basic_chain_returns_terminal_node():
    # A simple two-node chain A -> B: B should be identified as the last node.
    nodes = [{"id": "A"}, {"id": "B"}]  # node list in order
    edges = [{"source": "A", "target": "B"}]  # single edge from A to B
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.33μs (46.9% faster)


def test_single_node_no_edges_returns_that_node():
    # With one node and no edges, that node is trivially the last node.
    nodes = [{"id": "solo"}]
    edges = []  # no edges
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 500ns (150% faster)


def test_multiple_sinks_returns_first_sink_in_node_order():
    # If multiple nodes have no outgoing edges (multiple sinks),
    # the function should return the first encountered sink in the nodes list.
    nodes = [{"id": "Z"}, {"id": "A"}, {"id": "B"}]
    # Only A -> B, so Z and B are sinks; Z appears first and should be returned.
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 1.25μs (16.6% faster)


# Edge Test Cases


def test_empty_nodes_returns_none():
    # With no nodes at all, there is no last node; expect None.
    nodes = []
    edges = [{"source": "does-not-matter", "target": "none"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 750ns -> 1.17μs (35.7% slower)


def test_empty_edges_returns_first_node():
    # With no edges, every node has no outgoing edges; the function should
    # return the first node in the nodes list (vacuous truth of all()).
    nodes = [{"id": 1}, {"id": 2}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 500ns (150% faster)


def test_all_nodes_are_sources_returns_none():
    # If every node appears as a source in edges, there are no sinks and
    # the function should return None.
    nodes = [{"id": "A"}, {"id": "B"}]
    # Both A and B are listed as sources.
    edges = [{"source": "A", "target": "X"}, {"source": "B", "target": "Y"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.33μs (46.9% faster)


def test_missing_node_id_key_raises_key_error():
    # Nodes must be mappings with an 'id' key. If a node lacks 'id',
    # attempting to access n["id"] should raise KeyError.
    nodes = [{}]  # missing 'id' key
    edges = [{"source": "anything", "target": "nothing"}]
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.79μs -> 1.58μs (13.2% faster)


def test_missing_edge_source_key_raises_key_error():
    # Edges are expected to have a 'source' key. If an edge lacks 'source',
    # accessing e["source"] will raise KeyError.
    nodes = [{"id": "A"}]
    edges = [{"target": "B"}]  # missing 'source'
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.67μs -> 958ns (74.0% faster)


def test_edges_reference_unknown_ids_should_still_find_existing_sinks():
    # If edges reference node ids that are not present in nodes,
    # these edges should not prevent valid sinks from being detected.
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "X", "target": "Y"}]  # 'X' not among node ids
    # Since neither A nor B are sources here, the first node 'A' should be returned.
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 1.25μs (20.0% faster)


def test_duplicate_node_ids_honor_node_order_and_exclude_sources():
    # When node ids are duplicated, any id that appears as a source should exclude
    # all nodes with that id; the first node that does not have its id used as a source
    # must be returned.
    nodes = [{"id": "A"}, {"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    # All nodes with id "A" are sources and must be excluded; only "B" remains as sink.
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.21μs -> 1.33μs (65.6% faster)


# Large Scale Test Cases (kept under 1000 elements and loops under 1000 iterations)


def test_large_chain_of_nodes_returns_terminal_node():
    # Build a long chain of nodes n0 -> n1 -> ... -> n998 (999 nodes).
    # The last node n998 has no outgoing edges and should be returned.
    size = 999  # under the 1000 element limit
    nodes = [
        {"id": f"n{i}"} for i in range(size)
    ]  # list comprehension uses <1000 steps
    # Create edges for a chain from n0 to n998-1 (each is a source)
    edges = [{"source": f"n{i}", "target": f"n{i+1}"} for i in range(size - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 19.6ms -> 114μs (16940% faster)


def test_large_many_sinks_returns_first_sink_quickly():
    # Create 200 nodes (well under 1000). Make most nodes be sinks except a block of sources.
    size = 200
    nodes = [{"id": i} for i in range(size)]
    # Make nodes 50..149 be sources pointing to some other ids, leaving nodes 0..49 and 150..199 as sinks.
    edges = [{"source": i, "target": i + 1} for i in range(50, 150)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 6.42μs -> 5.00μs (28.3% faster)


# Mutation-resistance targeted test:
def test_order_sensitivity_and_exact_matching_of_source_field():
    # This test guards against common faulty mutations such as checking 'target' instead of 'source'
    # or returning the last element unconditionally. We construct nodes and edges where the real sink
    # is neither the last node nor associated with 'target' exclusively.
    nodes = [{"id": "first"}, {"id": "mid"}, {"id": "last"}]
    # Edge uses 'source' == 'mid' so 'first' and 'last' are sinks; should return 'first'.
    edges = [{"source": "mid", "target": "last"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 1.25μs (20.0% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from src.algorithms.graph import find_last_node

# unit tests


class TestFindLastNodeBasic:
    """Test basic functionality under normal conditions"""

    def test_single_node_no_edges(self):
        """Test with a single node and no edges - node should be last"""
        nodes = [{"id": "A", "data": "node_a"}]
        edges = []
        # With no edges, the single node is the last node
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.25μs -> 542ns (131% faster)

    def test_linear_chain_three_nodes(self):
        """Test linear chain A→B→C where C is the last node"""
        nodes = [
            {"id": "A", "name": "start"},
            {"id": "B", "name": "middle"},
            {"id": "C", "name": "end"},
        ]
        edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]
        # C has no outgoing edges, so it's the last node
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.38μs -> 1.42μs (67.6% faster)

    def test_simple_two_node_graph(self):
        """Test simple A→B graph where B is last"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [{"source": "A", "target": "B"}]
        # B is not a source of any edge
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.88μs -> 1.25μs (50.0% faster)

    def test_tree_structure_returns_first_leaf(self):
        """Test tree structure where multiple leaves exist, should return first"""
        nodes = [
            {"id": "root"},
            {"id": "left"},
            {"id": "right"},
            {"id": "leaf1"},
            {"id": "leaf2"},
        ]
        edges = [
            {"source": "root", "target": "left"},
            {"source": "root", "target": "right"},
            {"source": "left", "target": "leaf1"},
            {"source": "right", "target": "leaf2"},
        ]
        # Both leaf1 and leaf2 have no outgoing edges
        # Function returns first match in iteration order
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 3.00μs -> 1.62μs (84.6% faster)


class TestFindLastNodeEmptyCases:
    """Test edge cases with empty inputs"""

    def test_empty_nodes_empty_edges(self):
        """Test with both empty lists"""
        nodes = []
        edges = []
        # No nodes means no last node
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 792ns -> 583ns (35.8% faster)

    def test_empty_nodes_with_edges(self):
        """Test with no nodes but edges defined (inconsistent state)"""
        nodes = []
        edges = [{"source": "A", "target": "B"}]
        # No nodes to iterate over
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 792ns -> 1.12μs (29.6% slower)

    def test_multiple_nodes_empty_edges(self):
        """Test with multiple nodes but no edges - returns first node"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = []
        # All nodes are "last" nodes, return first
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.21μs -> 500ns (142% faster)


class TestFindLastNodeCycles:
    """Test graphs with cycles (no last node)"""

    def test_self_loop_single_node(self):
        """Test node with self-loop - it's a source so not a last node"""
        nodes = [{"id": "A"}]
        edges = [{"source": "A", "target": "A"}]
        # A is a source, so it's not a last node
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.46μs -> 1.29μs (12.8% faster)

    def test_simple_cycle_two_nodes(self):
        """Test A→B→A cycle - no last node"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "A"}]
        # Both nodes are sources
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.92μs -> 1.42μs (35.4% faster)

    def test_cycle_with_tail(self):
        """Test A→B→C→B cycle with entry point - no last node"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "B", "target": "C"},
            {"source": "C", "target": "B"},
        ]
        # All nodes are sources
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.38μs -> 1.50μs (58.3% faster)

    def test_complete_graph(self):
        """Test complete graph where every node connects to every other"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "A", "target": "C"},
            {"source": "B", "target": "A"},
            {"source": "B", "target": "C"},
            {"source": "C", "target": "A"},
            {"source": "C", "target": "B"},
        ]
        # All nodes are sources
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.58μs -> 1.62μs (59.0% faster)


class TestFindLastNodeMultipleSinks:
    """Test graphs with multiple sink nodes"""

    def test_two_sink_nodes_returns_first(self):
        """Test graph with two sink nodes - returns first found"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = [{"source": "A", "target": "B"}, {"source": "A", "target": "C"}]
        # Both B and C are sinks
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.96μs -> 1.42μs (38.2% faster)

    def test_disconnected_components(self):
        """Test disconnected graph with sinks in each component"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
        edges = [{"source": "A", "target": "B"}, {"source": "C", "target": "D"}]
        # B and D are both sinks
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.96μs -> 1.38μs (42.4% faster)

    def test_star_graph_center_to_leaves(self):
        """Test star graph where center points to all leaves"""
        nodes = [{"id": "center"}, {"id": "leaf1"}, {"id": "leaf2"}, {"id": "leaf3"}]
        edges = [
            {"source": "center", "target": "leaf1"},
            {"source": "center", "target": "leaf2"},
            {"source": "center", "target": "leaf3"},
        ]
        # All leaves are sinks
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.04μs -> 1.38μs (48.5% faster)


class TestFindLastNodeDataIntegrity:
    """Test with various data integrity issues"""

    def test_edges_reference_nonexistent_nodes(self):
        """Test edges referencing node IDs not in nodes list"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [
            {"source": "A", "target": "Z"},  # Z doesn't exist
            {"source": "X", "target": "B"},  # X doesn't exist
        ]
        # A is a source so not last, B is not a source
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.96μs -> 1.42μs (38.2% faster)

    def test_node_with_additional_fields(self):
        """Test nodes with extra metadata fields"""
        nodes = [
            {"id": "A", "label": "Start", "x": 100, "y": 200},
            {"id": "B", "label": "End", "x": 300, "y": 200},
        ]
        edges = [{"source": "A", "target": "B", "weight": 5}]
        # B is last, should return entire node dict
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.92μs -> 1.33μs (43.7% faster)

    def test_duplicate_node_ids(self):
        """Test behavior with duplicate node IDs"""
        nodes = [
            {"id": "A", "version": 1},
            {"id": "B"},
            {"id": "A", "version": 2},  # Duplicate ID
        ]
        edges = [{"source": "B", "target": "C"}]
        # First "A" should be returned as it's not a source
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.50μs -> 1.25μs (20.0% faster)

    def test_edges_with_extra_fields(self):
        """Test edges with additional metadata don't affect logic"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [
            {
                "source": "A",
                "target": "B",
                "weight": 10,
                "color": "blue",
                "type": "directed",
            }
        ]
        # B is still the last node
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.79μs -> 1.29μs (38.7% faster)

    def test_numeric_node_ids(self):
        """Test with numeric node IDs instead of strings"""
        nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
        edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}]
        # Node with id=3 is last
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.33μs -> 1.46μs (59.9% faster)

    def test_mixed_type_node_ids(self):
        """Test with mixed type node IDs (strings and numbers)"""
        nodes = [{"id": "A"}, {"id": 2}, {"id": "C"}]
        edges = [{"source": "A", "target": 2}, {"source": 2, "target": "C"}]
        # "C" is last
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.50μs -> 1.42μs (76.4% faster)


class TestFindLastNodeComplexGraphs:
    """Test more complex graph structures"""

    def test_diamond_graph(self):
        """Test diamond shape: A→B,C; B,C→D"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "A", "target": "C"},
            {"source": "B", "target": "D"},
            {"source": "C", "target": "D"},
        ]
        # D is the only sink
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.79μs -> 1.58μs (76.4% faster)

    def test_long_chain_with_branch(self):
        """Test long chain with a branch at the end"""
        nodes = [
            {"id": "A"},
            {"id": "B"},
            {"id": "C"},
            {"id": "D"},
            {"id": "E"},
            {"id": "F"},
        ]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "B", "target": "C"},
            {"source": "C", "target": "D"},
            {"source": "D", "target": "E"},
            {"source": "D", "target": "F"},
        ]
        # E and F are both sinks
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 3.08μs -> 1.62μs (89.7% faster)

    def test_multiple_sources_one_sink(self):
        """Test multiple starting nodes converging to one sink"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
        edges = [
            {"source": "A", "target": "D"},
            {"source": "B", "target": "D"},
            {"source": "C", "target": "D"},
        ]
        # D is the only sink
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.67μs -> 1.50μs (77.7% faster)

    def test_graph_with_isolated_node(self):
        """Test graph with one isolated node (no edges to/from it)"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "isolated"}]
        edges = [{"source": "A", "target": "B"}]
        # Both B and isolated are not sources
        # Function returns first found
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.79μs -> 1.29μs (38.7% faster)


class TestFindLastNodeLargeScale:
    """Test performance and scalability with large datasets"""

    def test_large_linear_chain(self):
        """Test with a long linear chain of 500 nodes"""
        size = 500
        nodes = [{"id": f"node_{i}"} for i in range(size)]
        edges = [
            {"source": f"node_{i}", "target": f"node_{i+1}"} for i in range(size - 1)
        ]
        # Last node in chain should be returned
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 4.77ms -> 63.7μs (7391% faster)

    def test_large_graph_single_sink(self):
        """Test with 300 nodes all pointing to a single sink"""
        size = 300
        nodes = [{"id": f"node_{i}"} for i in range(size)]
        nodes.append({"id": "sink"})
        edges = [{"source": f"node_{i}", "target": "sink"} for i in range(size)]
        # Sink node should be returned
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.74ms -> 38.4μs (4422% faster)

    def test_large_graph_many_sinks(self):
        """Test with 200 source nodes pointing to 200 different sinks"""
        sources = 200
        nodes = [{"id": f"source_{i}"} for i in range(sources)]
        nodes.extend([{"id": f"sink_{i}"} for i in range(sources)])
        edges = [
            {"source": f"source_{i}", "target": f"sink_{i}"} for i in range(sources)
        ]
        # One of the sink nodes should be returned
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 800μs -> 28.7μs (2694% faster)

    def test_large_binary_tree(self):
        """Test with a complete binary tree structure"""
        # Create a binary tree with 7 levels (127 nodes)
        levels = 7
        total_nodes = 2**levels - 1
        nodes = [{"id": i} for i in range(total_nodes)]
        edges = []

        # Build tree edges: node i has children 2i+1 and 2i+2
        for i in range(total_nodes // 2):
            left_child = 2 * i + 1
            right_child = 2 * i + 2
            if left_child < total_nodes:
                edges.append({"source": i, "target": left_child})
            if right_child < total_nodes:
                edges.append({"source": i, "target": right_child})

        # Result should be a leaf node (nodes in second half are leaves)
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 166μs -> 6.71μs (2382% faster)

    def test_dense_graph_structure(self):
        """Test with a dense graph (many edges between nodes)"""
        size = 50
        nodes = [{"id": i} for i in range(size)]
        edges = []

        # Create edges: each node points to next 5 nodes
        for i in range(size - 1):
            for j in range(i + 1, min(i + 6, size)):
                edges.append({"source": i, "target": j})

        # Last node (size-1) should have no outgoing edges
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 230μs -> 9.00μs (2460% faster)

    def test_wide_graph_one_level(self):
        """Test with one source node connecting to 500 nodes"""
        size = 500
        nodes = [{"id": "root"}]
        nodes.extend([{"id": f"child_{i}"} for i in range(size)])
        edges = [{"source": "root", "target": f"child_{i}"} for i in range(size)]

        # One of the children should be returned
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 22.2μs -> 12.7μs (74.4% faster)

    def test_many_disconnected_pairs(self):
        """Test with 250 disconnected pairs (A→B)"""
        pairs = 250
        nodes = []
        edges = []

        for i in range(pairs):
            nodes.append({"id": f"A_{i}"})
            nodes.append({"id": f"B_{i}"})
            edges.append({"source": f"A_{i}", "target": f"B_{i}"})

        # Should return one of the B nodes
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 12.5μs -> 19.0μs (34.0% slower)

    def test_ladder_graph_structure(self):
        """Test ladder graph: two parallel chains with cross connections"""
        size = 200
        nodes = []
        edges = []

        # Create two parallel chains
        for i in range(size):
            nodes.append({"id": f"top_{i}"})
            nodes.append({"id": f"bottom_{i}"})

        # Connect chains horizontally
        for i in range(size - 1):
            edges.append({"source": f"top_{i}", "target": f"top_{i+1}"})
            edges.append({"source": f"bottom_{i}", "target": f"bottom_{i+1}"})

        # Add rungs (vertical connections)
        for i in range(0, size, 10):
            edges.append({"source": f"top_{i}", "target": f"bottom_{i}"})

        # Last nodes on both chains have no outgoing edges
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.99ms -> 57.5μs (5101% faster)


# 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-find_last_node-mk3kdhgh and push.