⚡️ Speed up function _prune_literal_if_trivial by 22%

[codeflash-ai]

📄 22% (0.22x) speedup for _prune_literal_if_trivial in src/uberjob/_transformations/pruning.py

⏱️ Runtime : 5.62 milliseconds → 4.59 milliseconds (best of 34 runs)

📝 Explanation and details

The optimization achieves a 22% speedup by implementing three key improvements:

1. Early-exit iteration over out-edges: Instead of using all() which must evaluate every edge even when the first non-Dependency is found, the optimized version uses a simple for loop that returns immediately upon finding the first non-Dependency edge. The line profiler shows this reduces time spent on edge checking from 75.6% to 80.7% of total time, but with much faster execution overall.

2. Early return for trivial cases: Added explicit checks for m == 0 or n == 0 before the cost comparison, avoiding unnecessary computation when there are no predecessors or successors. This optimization particularly benefits test cases with isolated nodes or nodes with only incoming/outgoing edges.

3. Direct nested loops instead of itertools.product(): Replaced itertools.product(predecessors, successors) with simple nested for loops. For the typical small-to-moderate graph sizes seen in the tests, direct iteration is faster than the generator overhead of itertools.product().

Impact on workloads: Since _prune_literal_if_trivial is called within a loop over all literals in prune_plan(), this 22% improvement compounds significantly during graph optimization phases. The test results show the optimization is particularly effective for:

• Isolated literals (104% faster)
• Literals with only incoming edges (194% faster)
• Large-scale scenarios with many isolated literals (36.4% faster)
• Chain topologies (69.7% faster)

The optimization maintains identical behavior while reducing algorithmic overhead, making graph pruning operations substantially faster across diverse graph structures.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import itertools

# imports
import pytest
from uberjob._transformations.pruning import _prune_literal_if_trivial

# minimal stubs for Plan, Literal, Dependency, and graph to allow isolated testing


class Dependency:
    """Stub for Dependency edge type."""

    pass


class Literal:
    """Stub for Literal node type."""

    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return f"Literal({self.value!r})"

    def __eq__(self, other):
        return isinstance(other, Literal) and self.value == other.value

    def __hash__(self):
        return hash(("Literal", self.value))


class Node:
    """Stub for a generic node (non-literal)."""

    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return f"Node({self.name!r})"

    def __eq__(self, other):
        return isinstance(other, Node) and self.name == other.name

    def __hash__(self):
        return hash(("Node", self.name))


class DummyGraph:
    """A minimal MultiDiGraph-like class with required methods for testing."""

    def __init__(self):
        # Each edge is (from, to, key) -> data
        self._edges = {}  # (u, v, k) -> edge_data
        self._adj = {}  # node -> {neighbor: set(keys)}
        self._pred = {}  # node -> {neighbor: set(keys)}
        self._nodes = set()

    def add_node(self, node):
        self._nodes.add(node)
        self._adj.setdefault(node, {})
        self._pred.setdefault(node, {})

    def add_edge(self, u, v, data):
        self.add_node(u)
        self.add_node(v)
        key = id(data)
        self._edges[(u, v, key)] = data
        self._adj[u].setdefault(v, set()).add(key)
        self._pred[v].setdefault(u, set()).add(key)

    def out_edges(self, node, keys=False):
        # Returns (u, v, key) for all outgoing edges from node
        for v, keyset in self._adj.get(node, {}).items():
            for k in keyset:
                if keys:
                    yield (node, v, self._edges[(node, v, k)])
                else:
                    yield (node, v)

    def predecessors(self, node):
        return list(self._pred.get(node, {}).keys())

    def successors(self, node):
        return list(self._adj.get(node, {}).keys())

    def remove_node(self, node):
        # Remove all edges to and from node
        for v in list(self._adj.get(node, {})):
            for k in self._adj[node][v]:
                del self._edges[(node, v, k)]
                self._pred[v][node].remove(k)
                if not self._pred[v][node]:
                    del self._pred[v][node]
            del self._adj[node][v]
        for u in list(self._pred.get(node, {})):
            for k in self._pred[node][u]:
                del self._edges[(u, node, k)]
                self._adj[u][node].remove(k)
                if not self._adj[u][node]:
                    del self._adj[u][node]
            del self._pred[node][u]
        self._adj.pop(node, None)
        self._pred.pop(node, None)
        self._nodes.discard(node)

    def has_node(self, node):
        return node in self._nodes

    def nodes(self):
        return set(self._nodes)

    def edges(self):
        return set((u, v) for (u, v, _) in self._edges)

    def edge_types(self, node):
        # Return set of types of out edge data for node
        return set(type(data) for _, _, data in self.out_edges(node, keys=True))


class Plan:
    """Stub Plan with a 'graph' attribute."""

    def __init__(self):
        self.graph = DummyGraph()


from uberjob._transformations.pruning import _prune_literal_if_trivial

# unit tests

# --- Basic Test Cases ---


def test_literal_with_no_edges_is_pruned():
    """Literal with no edges should be pruned (removed from graph)."""
    plan = Plan()
    lit = Literal(1)
    plan.graph.add_node(lit)
    _prune_literal_if_trivial(plan, lit)  # 8.73μs -> 4.28μs (104% faster)


def test_literal_with_only_incoming_edges_is_pruned():
    """Literal with only incoming edges should be pruned."""
    plan = Plan()
    lit = Literal(2)
    pred = Node("pred")
    plan.graph.add_edge(pred, lit, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 10.4μs -> 3.54μs (194% faster)


def test_literal_with_only_outgoing_edges_is_pruned():
    """Literal with only outgoing edges should be pruned."""
    plan = Plan()
    lit = Literal(3)
    succ = Node("succ")
    plan.graph.add_edge(lit, succ, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.79μs -> 3.01μs (26.0% faster)


def test_literal_with_one_pred_and_one_succ_is_pruned_and_edge_added():
    """Literal with one predecessor and one successor: should be pruned and edge added."""
    plan = Plan()
    lit = Literal(4)
    pred = Node("pred")
    succ = Node("succ")
    plan.graph.add_edge(pred, lit, Dependency())
    plan.graph.add_edge(lit, succ, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.54μs -> 2.82μs (25.6% faster)


def test_literal_with_multiple_preds_and_succs_prunes_when_mn_leq_m_plus_n():
    """Literal with 2 preds and 2 succs: 2*2=4 <= 2+2=4, should prune."""
    plan = Plan()
    lit = Literal(5)
    preds = [Node("pred1"), Node("pred2")]
    succs = [Node("succ1"), Node("succ2")]
    for p in preds:
        plan.graph.add_edge(p, lit, Dependency())
    for s in succs:
        plan.graph.add_edge(lit, s, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.47μs -> 2.86μs (21.4% faster)
    # all combinations of pred->succ exist
    for p in preds:
        for s in succs:
            pass


def test_literal_with_preds_succs_not_pruned_when_mn_gt_m_plus_n():
    """Literal with 3 preds and 3 succs: 3*3=9 > 3+3=6, should NOT prune."""
    plan = Plan()
    lit = Literal(6)
    preds = [Node("pred1"), Node("pred2"), Node("pred3")]
    succs = [Node("succ1"), Node("succ2"), Node("succ3")]
    for p in preds:
        plan.graph.add_edge(p, lit, Dependency())
    for s in succs:
        plan.graph.add_edge(lit, s, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.44μs -> 2.97μs (15.7% faster)
    # no new pred->succ edges should be present
    for p in preds:
        for s in succs:
            pass


def test_literal_with_non_dependency_out_edge_is_not_pruned():
    """Literal with an outgoing edge that is not a Dependency should not be pruned."""

    class NotDependency:
        pass

    plan = Plan()
    lit = Literal(7)
    succ = Node("succ")
    plan.graph.add_edge(lit, succ, NotDependency())
    _prune_literal_if_trivial(plan, lit)  # 3.32μs -> 2.85μs (16.3% faster)


def test_literal_with_non_dependency_in_edge_is_pruned():
    """Literal with a non-Dependency incoming edge should be pruned (since only out_edges are checked)."""

    class NotDependency:
        pass

    plan = Plan()
    lit = Literal(8)
    pred = Node("pred")
    plan.graph.add_edge(pred, lit, NotDependency())
    _prune_literal_if_trivial(plan, lit)  # 10.5μs -> 3.72μs (181% faster)


# --- Edge Test Cases ---


def test_literal_with_no_predecessors_and_multiple_successors():
    """Literal with 0 preds, >0 succs: should be pruned (m=0, n>0, 0*X <= 0+X)."""
    plan = Plan()
    lit = Literal(9)
    succs = [Node("succ1"), Node("succ2")]
    for s in succs:
        plan.graph.add_edge(lit, s, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.40μs -> 2.90μs (17.1% faster)
    # No new edges should be added (no preds)


def test_literal_with_multiple_predecessors_and_no_successors():
    """Literal with >0 preds, 0 succs: should be pruned (m>0, n=0, 0*X <= X+0)."""
    plan = Plan()
    lit = Literal(10)
    preds = [Node("pred1"), Node("pred2")]
    for p in preds:
        plan.graph.add_edge(p, lit, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 11.7μs -> 3.60μs (224% faster)
    # No new edges should be added (no succs)


def test_literal_with_no_predecessors_and_no_successors():
    """Literal with 0 preds, 0 succs: should be pruned."""
    plan = Plan()
    lit = Literal(11)
    plan.graph.add_node(lit)
    _prune_literal_if_trivial(plan, lit)  # 6.71μs -> 3.31μs (102% faster)


def test_literal_with_self_loop():
    """Literal with self-loop should not be pruned if out_edge is not Dependency."""

    class NotDependency:
        pass

    plan = Plan()
    lit = Literal(12)
    plan.graph.add_edge(lit, lit, NotDependency())
    _prune_literal_if_trivial(plan, lit)  # 3.48μs -> 2.80μs (24.3% faster)


def test_literal_with_self_loop_dependency():
    """Literal with self-loop as Dependency: should be pruned (m=1, n=1, 1<=2)."""
    plan = Plan()
    lit = Literal(13)
    plan.graph.add_edge(lit, lit, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.27μs -> 2.74μs (19.4% faster)


def test_literal_with_duplicate_edges():
    """Literal with duplicate Dependency edges should be pruned."""
    plan = Plan()
    lit = Literal(14)
    pred = Node("pred")
    succ = Node("succ")
    plan.graph.add_edge(pred, lit, Dependency())
    plan.graph.add_edge(pred, lit, Dependency())  # duplicate in-edge
    plan.graph.add_edge(lit, succ, Dependency())
    plan.graph.add_edge(lit, succ, Dependency())  # duplicate out-edge
    _prune_literal_if_trivial(plan, lit)  # 3.29μs -> 2.70μs (21.6% faster)


def test_literal_with_no_edges_is_pruned_and_graph_remains_consistent():
    """After pruning, graph should not contain any reference to the literal."""
    plan = Plan()
    lit = Literal(15)
    plan.graph.add_node(lit)
    _prune_literal_if_trivial(plan, lit)  # 6.91μs -> 3.53μs (95.4% faster)
    for u, v in plan.graph.edges():
        pass


# --- Large Scale Test Cases ---


def test_large_number_of_preds_and_succs_prunes_when_mn_leq_m_plus_n():
    """For m=20, n=20, 400 > 40, so should NOT prune. For m=5, n=5, 25=10, should NOT prune."""
    plan = Plan()
    lit = Literal(16)
    preds = [Node(f"pred{i}") for i in range(5)]
    succs = [Node(f"succ{i}") for i in range(5)]
    for p in preds:
        plan.graph.add_edge(p, lit, Dependency())
    for s in succs:
        plan.graph.add_edge(lit, s, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.21μs -> 2.90μs (10.9% faster)
    # now try with m=2, n=2 (4=4, should prune)
    plan2 = Plan()
    lit2 = Literal(17)
    preds2 = [Node("predA"), Node("predB")]
    succs2 = [Node("succA"), Node("succB")]
    for p in preds2:
        plan2.graph.add_edge(p, lit2, Dependency())
    for s in succs2:
        plan2.graph.add_edge(lit2, s, Dependency())
    _prune_literal_if_trivial(plan2, lit2)  # 1.88μs -> 1.55μs (20.9% faster)
    for p in preds2:
        for s in succs2:
            pass


def test_large_graph_with_many_literals_only_trivial_are_pruned():
    """Test with 100 literals, each with 1 pred and 1 succ: all should be pruned."""
    plan = Plan()
    nodes = [Node(f"n{i}") for i in range(101)]
    literals = [Literal(i) for i in range(100)]
    for i, lit in enumerate(literals):
        plan.graph.add_edge(nodes[i], lit, Dependency())
        plan.graph.add_edge(lit, nodes[i + 1], Dependency())
    for lit in literals:
        _prune_literal_if_trivial(plan, lit)  # 143μs -> 115μs (24.0% faster)
    for lit in literals:
        pass
    # All node-to-node edges should exist
    for i in range(100):
        pass


def test_large_graph_literal_with_many_preds_and_succs_not_pruned_when_not_trivial():
    """Literal with 30 preds, 30 succs: 900>60, should not be pruned."""
    plan = Plan()
    lit = Literal(99)
    preds = [Node(f"pred{i}") for i in range(30)]
    succs = [Node(f"succ{i}") for i in range(30)]
    for p in preds:
        plan.graph.add_edge(p, lit, Dependency())
    for s in succs:
        plan.graph.add_edge(lit, s, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 3.50μs -> 2.73μs (27.8% faster)
    # No new pred->succ edges
    for p in preds:
        for s in succs:
            pass


def test_large_graph_literal_with_many_preds_and_one_succ_is_pruned():
    """Literal with 999 preds, 1 succ: 999*1=999 <= 999+1=1000, should prune."""
    plan = Plan()
    lit = Literal(100)
    preds = [Node(f"pred{i}") for i in range(999)]
    succ = Node("succ")
    for p in preds:
        plan.graph.add_edge(p, lit, Dependency())
    plan.graph.add_edge(lit, succ, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 5.09μs -> 4.32μs (17.8% faster)
    for p in preds:
        pass


def test_large_graph_literal_with_one_pred_and_many_succs_is_pruned():
    """Literal with 1 pred, 999 succs: 1*999=999 <= 1+999=1000, should prune."""
    plan = Plan()
    lit = Literal(101)
    pred = Node("pred")
    succs = [Node(f"succ{i}") for i in range(999)]
    plan.graph.add_edge(pred, lit, Dependency())
    for s in succs:
        plan.graph.add_edge(lit, s, Dependency())
    _prune_literal_if_trivial(plan, lit)  # 5.58μs -> 4.85μs (15.1% faster)
    for s in succs:
        pass


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

import itertools

import networkx as nx

# imports
import pytest
from uberjob._transformations.pruning import _prune_literal_if_trivial


# Mock classes to simulate the required classes from uberjob
class Dependency:
    pass


class Literal:
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return f"Literal({self.value!r})"

    def __eq__(self, other):
        return isinstance(other, Literal) and self.value == other.value

    def __hash__(self):
        return hash((self.value, "Literal"))


class Plan:
    def __init__(self):
        # Use a MultiDiGraph to support keys and multiple edges
        self.graph = nx.MultiDiGraph()


from uberjob._transformations.pruning import _prune_literal_if_trivial

# ----------------- UNIT TESTS -----------------

# Helper functions for test clarity
def get_edges(graph):
    # Return edges as (u, v) ignoring keys and edge data
    return set((u, v) for u, v, _ in graph.edges(keys=True))


# ----------- 1. BASIC TEST CASES ------------


def test_literal_with_no_predecessors_or_successors():
    # The literal is isolated; should be pruned (removed)
    plan = Plan()
    l = Literal(1)
    plan.graph.add_node(l)
    _prune_literal_if_trivial(plan, l)  # 22.1μs -> 18.8μs (17.6% faster)


def test_literal_with_one_predecessor_one_successor():
    # Should be pruned and predecessor connected to successor
    plan = Plan()
    pred = "pred"
    succ = "succ"
    l = Literal(2)
    plan.graph.add_node(pred)
    plan.graph.add_node(succ)
    plan.graph.add_node(l)
    plan.graph.add_edge(pred, l, Dependency())
    plan.graph.add_edge(l, succ, Dependency())
    _prune_literal_if_trivial(plan, l)  # 13.4μs -> 13.0μs (3.06% faster)


def test_literal_with_multiple_predecessors_and_successors():
    # Should be pruned and all predecessors connected to all successors
    plan = Plan()
    preds = ["p1", "p2"]
    succs = ["s1", "s2"]
    l = Literal(3)
    for node in preds + succs + [l]:
        plan.graph.add_node(node)
    for p in preds:
        plan.graph.add_edge(p, l, Dependency())
    for s in succs:
        plan.graph.add_edge(l, s, Dependency())
    _prune_literal_if_trivial(plan, l)  # 12.9μs -> 12.5μs (3.53% faster)
    # All pred->succ edges exist
    for p in preds:
        for s in succs:
            pass


def test_literal_with_no_successors():
    # Should be pruned (no successors)
    plan = Plan()
    pred = "pred"
    l = Literal(4)
    plan.graph.add_node(pred)
    plan.graph.add_node(l)
    plan.graph.add_edge(pred, l, Dependency())
    _prune_literal_if_trivial(plan, l)  # 17.2μs -> 13.5μs (27.2% faster)


def test_literal_with_no_predecessors():
    # Should be pruned (no predecessors)
    plan = Plan()
    succ = "succ"
    l = Literal(5)
    plan.graph.add_node(succ)
    plan.graph.add_node(l)
    plan.graph.add_edge(l, succ, Dependency())
    _prune_literal_if_trivial(plan, l)  # 12.2μs -> 11.8μs (3.10% faster)


def test_literal_with_non_dependency_out_edge():
    # Should NOT prune if any out-edge is not a Dependency
    class NotDependency:
        pass

    plan = Plan()
    pred = "pred"
    succ = "succ"
    l = Literal(6)
    plan.graph.add_node(pred)
    plan.graph.add_node(succ)
    plan.graph.add_node(l)
    plan.graph.add_edge(pred, l, Dependency())
    plan.graph.add_edge(l, succ, NotDependency())
    _prune_literal_if_trivial(plan, l)  # 12.4μs -> 11.4μs (8.92% faster)


def test_literal_with_dependency_and_non_dependency_out_edges():
    # Should NOT prune if any out-edge is not a Dependency
    class NotDependency:
        pass

    plan = Plan()
    pred = "pred"
    succ1 = "succ1"
    succ2 = "succ2"
    l = Literal(7)
    plan.graph.add_node(pred)
    plan.graph.add_node(succ1)
    plan.graph.add_node(succ2)
    plan.graph.add_node(l)
    plan.graph.add_edge(pred, l, Dependency())
    plan.graph.add_edge(l, succ1, Dependency())
    plan.graph.add_edge(l, succ2, NotDependency())
    _prune_literal_if_trivial(plan, l)  # 12.5μs -> 11.7μs (6.72% faster)


def test_literal_with_multiple_edges_between_same_nodes():
    # Should still prune if all are Dependency
    plan = Plan()
    pred = "pred"
    succ = "succ"
    l = Literal(8)
    plan.graph.add_node(pred)
    plan.graph.add_node(succ)
    plan.graph.add_node(l)
    # Add two edges pred->l
    plan.graph.add_edge(pred, l, Dependency())
    plan.graph.add_edge(pred, l, Dependency())
    # Add two edges l->succ
    plan.graph.add_edge(l, succ, Dependency())
    plan.graph.add_edge(l, succ, Dependency())
    _prune_literal_if_trivial(plan, l)  # 11.7μs -> 11.4μs (3.17% faster)


# ----------- 2. EDGE TEST CASES ------------


def test_literal_with_large_number_of_predecessors_and_successors_no_prune():
    # Should NOT prune if m*n > m+n
    plan = Plan()
    preds = [f"p{i}" for i in range(10)]
    succs = [f"s{i}" for i in range(10)]
    l = Literal(9)
    for node in preds + succs + [l]:
        plan.graph.add_node(node)
    for p in preds:
        plan.graph.add_edge(p, l, Dependency())
    for s in succs:
        plan.graph.add_edge(l, s, Dependency())
    _prune_literal_if_trivial(plan, l)  # 12.0μs -> 11.5μs (4.64% faster)


def test_literal_with_large_number_of_predecessors_and_few_successors_prune():
    # Should prune if m*n <= m+n
    plan = Plan()
    preds = [f"p{i}" for i in range(5)]
    succs = [f"s{i}" for i in range(2)]
    l = Literal(10)
    for node in preds + succs + [l]:
        plan.graph.add_node(node)
    for p in preds:
        plan.graph.add_edge(p, l, Dependency())
    for s in succs:
        plan.graph.add_edge(l, s, Dependency())
    _prune_literal_if_trivial(plan, l)  # 12.0μs -> 11.6μs (3.79% faster)
    for p in preds:
        for s in succs:
            pass


def test_literal_with_self_loop():
    # Should prune if all out-edges are Dependency, even with self-loop
    plan = Plan()
    pred = "pred"
    succ = "succ"
    l = Literal(11)
    plan.graph.add_node(pred)
    plan.graph.add_node(succ)
    plan.graph.add_node(l)
    plan.graph.add_edge(pred, l, Dependency())
    plan.graph.add_edge(l, succ, Dependency())
    plan.graph.add_edge(l, l, Dependency())  # self-loop
    _prune_literal_if_trivial(plan, l)  # 11.7μs -> 11.4μs (2.71% faster)


def test_literal_with_no_edges():
    # Should prune isolated node
    plan = Plan()
    l = Literal(12)
    plan.graph.add_node(l)
    _prune_literal_if_trivial(plan, l)  # 17.1μs -> 14.2μs (20.3% faster)


def test_literal_with_only_incoming_edges():
    # Should prune node with only incoming edges
    plan = Plan()
    pred1 = "pred1"
    pred2 = "pred2"
    l = Literal(13)
    plan.graph.add_node(pred1)
    plan.graph.add_node(pred2)
    plan.graph.add_node(l)
    plan.graph.add_edge(pred1, l, Dependency())
    plan.graph.add_edge(pred2, l, Dependency())
    _prune_literal_if_trivial(plan, l)  # 17.1μs -> 12.9μs (32.1% faster)


def test_literal_with_only_outgoing_edges():
    # Should prune node with only outgoing edges
    plan = Plan()
    succ1 = "succ1"
    succ2 = "succ2"
    l = Literal(14)
    plan.graph.add_node(succ1)
    plan.graph.add_node(succ2)
    plan.graph.add_node(l)
    plan.graph.add_edge(l, succ1, Dependency())
    plan.graph.add_edge(l, succ2, Dependency())
    _prune_literal_if_trivial(plan, l)  # 12.2μs -> 11.8μs (3.80% faster)


def test_literal_with_non_literal_node():
    # Should not affect non-Literal nodes
    plan = Plan()
    node = "not_a_literal"
    plan.graph.add_node(node)
    # Should not raise
    _prune_literal_if_trivial(plan, node)  # 14.4μs -> 12.0μs (20.2% faster)


# ----------- 3. LARGE SCALE TEST CASES ------------


def test_large_scale_pruning():
    # Test with 100 predecessors and 2 successors (should prune)
    plan = Plan()
    preds = [f"p{i}" for i in range(100)]
    succs = [f"s{i}" for i in range(2)]
    l = Literal("large")
    for node in preds + succs + [l]:
        plan.graph.add_node(node)
    for p in preds:
        plan.graph.add_edge(p, l, Dependency())
    for s in succs:
        plan.graph.add_edge(l, s, Dependency())
    _prune_literal_if_trivial(plan, l)  # 13.5μs -> 13.0μs (4.24% faster)
    for p in preds:
        for s in succs:
            pass


def test_large_scale_no_prune():
    # Test with 32 predecessors and 32 successors (should NOT prune, 32*32 > 32+32)
    plan = Plan()
    preds = [f"p{i}" for i in range(32)]
    succs = [f"s{i}" for i in range(32)]
    l = Literal("large_no_prune")
    for node in preds + succs + [l]:
        plan.graph.add_node(node)
    for p in preds:
        plan.graph.add_edge(p, l, Dependency())
    for s in succs:
        plan.graph.add_edge(l, s, Dependency())
    _prune_literal_if_trivial(plan, l)  # 13.4μs -> 12.4μs (8.17% faster)


def test_large_scale_all_isolated_literals():
    # 500 isolated literals, all should be pruned
    plan = Plan()
    literals = [Literal(i) for i in range(500)]
    for l in literals:
        plan.graph.add_node(l)
    for l in literals:
        _prune_literal_if_trivial(plan, l)  # 2.24ms -> 1.65ms (36.4% faster)
    for l in literals:
        pass


def test_large_scale_chain_of_literals():
    # Chain: l0->l1->l2->...->ln, all should be pruned in sequence
    plan = Plan()
    n = 50
    literals = [Literal(i) for i in range(n)]
    for l in literals:
        plan.graph.add_node(l)
    for i in range(n - 1):
        plan.graph.add_edge(literals[i], literals[i + 1], Dependency())
    # Prune from last to first
    for l in reversed(literals):
        _prune_literal_if_trivial(plan, l)  # 259μs -> 152μs (69.7% faster)
    for l in literals:
        pass


def test_large_scale_star_topology():
    # One literal with 999 predecessors and 1 successor (should prune)
    plan = Plan()
    preds = [f"p{i}" for i in range(999)]
    succ = "s"
    l = Literal("center")
    for node in preds + [succ, l]:
        plan.graph.add_node(node)
    for p in preds:
        plan.graph.add_edge(p, l, Dependency())
    plan.graph.add_edge(l, succ, Dependency())
    _prune_literal_if_trivial(plan, l)  # 19.5μs -> 19.6μs (0.688% slower)
    for p in preds:
        pass


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

from uberjob._plan import Plan
from uberjob._transformations.pruning import _prune_literal_if_trivial
from uberjob.graph import Literal
import pytest


def test__prune_literal_if_trivial():
    with pytest.raises(
        NetworkXError,
        match="nbunch\\ is\\ not\\ a\\ node\\ or\\ a\\ sequence\\ of\\ nodes\\.",
    ):
        _prune_literal_if_trivial(Plan(), Literal("", scope=""))

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_teststest_plan_py_teststest_render_py_teststest_scheduler_py_teststest_unpack_py__replay_test_0.py::test_uberjob__transformations_pruning__prune_literal_if_trivial	1.65ms	1.49ms	11.1%✅
test_pytest_teststest_registry_py_teststest_version_py_teststest_progress_py_teststest_atoms_py__replay_test_0.py::test_uberjob__transformations_pruning__prune_literal_if_trivial	958μs	888μs	7.89%✅

To edit these changes git checkout codeflash/optimize-_prune_literal_if_trivial-mi6qok6l and push.