Rename AdjacencySet to AdjacencySetProxy, add bulk ops and safety guards

Rename class and file to AdjacencySetProxy to clarify write-through
semantics. Add bulk-efficient clear(), __isub__(), __ior__() overrides
and remove_edges() on the Cython core. Guard GraphNode.discard() against
double-destroy via membership check. Filter duplicates in update(). Add
error-path tests for wrong types, cross-graph edges, and self-edges.

Made-with: Cursor

Author: Andy-Jost

…ore/_graph/_graph_def/_adjacency_set.pyx …raph/_graph_def/_adjacency_set_proxy.pyxcuda_core/cuda/core/_graph/_graph_def/_adjacency_set.pyx renamed to cuda_core/cuda/core/_graph/_graph_def/_adjacency_set_proxy.pyx cuda_core/cuda/core/_graph/_graph_def/_adjacency_set.pyx renamed to cuda_core/cuda/core/_graph/_graph_def/_adjacency_set_proxy.pyx

@@ -20,14 +20,16 @@ from collections.abc import MutableSet
 
 # ---- Python MutableSet wrapper ----------------------------------------------
 
-class AdjacencySet(MutableSet):
-    """Mutable set-like view of a node's predecessors or successors."""
+class AdjacencySetProxy(MutableSet):
+    """Mutable set proxy for a node's predecessors or successors.  Mutations
+    write through to the underlying CUDA graph."""
 
     __slots__ = ("_core",)
 
     def __init__(self, node, bint is_fwd):
         self._core = _AdjacencySetCore(node, is_fwd)
 
+    # Used by operators such as &|^ to create non-proxy views when needed.
     @classmethod
     def _from_iterable(cls, it):
         return set(it)
@@ -62,6 +64,22 @@ class AdjacencySet(MutableSet):
 
     # --- override for bulk efficiency ---
 
+    def clear(self):
+        """Remove all edges in a single driver call."""
+        members = (<_AdjacencySetCore>self._core).query()
+        if members:
+            (<_AdjacencySetCore>self._core).remove_edges(members)
+
+    def __isub__(self, it):
+        """Remove edges to all nodes in *it* in a single driver call."""
+        if it is self:
+            self.clear()
+        else:
+            to_remove = [v for v in it if isinstance(v, GraphNode) and v in self]
+            if to_remove:
+                (<_AdjacencySetCore>self._core).remove_edges(to_remove)
+        return self
+
     def update(self, *others):
         """Add edges to multiple nodes at once."""
         nodes = []
@@ -76,21 +94,28 @@ class AdjacencySet(MutableSet):
             if not isinstance(n, GraphNode):
                 raise TypeError(
                     f"expected GraphNode, got {type(n).__name__}")
-        (<_AdjacencySetCore>self._core).add_edges(nodes)
+        new = [n for n in nodes if n not in self]
+        if new:
+            (<_AdjacencySetCore>self._core).add_edges(new)
+
+    def __ior__(self, it):
+        """Add edges to all nodes in *it* in a single driver call."""
+        self.update(it)
+        return self
 
     def __repr__(self):
         return "{" + ", ".join(repr(n) for n in self) + "}"
 
 
-# ---- cdef core holding function pointer ------------------------------------
+# ---- cdef core holding a function pointer ------------------------------------
 
-# Signature shared by _get_preds and _get_succs.
+# Signature shared by driver_get_preds and driver_get_succs.
 ctypedef cydriver.CUresult (*_adj_fn_t)(
     cydriver.CUgraphNode, cydriver.CUgraphNode*, size_t*) noexcept nogil
 
 
 cdef class _AdjacencySetCore:
-    """Cythonized core implementing AdjacencySet"""
+    """Cythonized core implementing AdjacencySetProxy"""
     cdef:
         GraphNodeHandle _h_node
         GraphHandle _h_graph
@@ -101,7 +126,7 @@ cdef class _AdjacencySetCore:
         self._h_node = node._h_node
         self._h_graph = graph_node_get_graph(node._h_node)
         self._is_fwd = is_fwd
-        self._query_fn = _get_succs if is_fwd else _get_preds
+        self._query_fn = driver_get_succs if is_fwd else driver_get_preds
 
     cdef inline void _resolve_edge(
             self, GraphNode other,
@@ -144,15 +169,28 @@ cdef class _AdjacencySetCore:
         cdef cydriver.CUgraphNode c_from, c_to
         self._resolve_edge(other, &c_from, &c_to)
         with nogil:
-            HANDLE_RETURN(_add_edge(as_cu(self._h_graph), &c_from, &c_to, 1))
+            HANDLE_RETURN(driver_add_edges(as_cu(self._h_graph), &c_from, &c_to, 1))
+
+    cdef void add_edges(self, list nodes):
+        cdef size_t n = len(nodes)
+        cdef vector[cydriver.CUgraphNode] from_vec
+        cdef vector[cydriver.CUgraphNode] to_vec
+        from_vec.resize(n)
+        to_vec.resize(n)
+        cdef size_t i
+        for i in range(n):
+            self._resolve_edge(<GraphNode>nodes[i], &from_vec[i], &to_vec[i])
+        with nogil:
+            HANDLE_RETURN(driver_add_edges(
+                as_cu(self._h_graph), from_vec.data(), to_vec.data(), n))
 
     cdef void remove_edge(self, GraphNode other):
         cdef cydriver.CUgraphNode c_from, c_to
         self._resolve_edge(other, &c_from, &c_to)
         with nogil:
-            HANDLE_RETURN(_remove_edge(as_cu(self._h_graph), &c_from, &c_to, 1))
+            HANDLE_RETURN(driver_remove_edges(as_cu(self._h_graph), &c_from, &c_to, 1))
 
-    cdef void add_edges(self, list nodes):
+    cdef void remove_edges(self, list nodes):
         cdef size_t n = len(nodes)
         cdef vector[cydriver.CUgraphNode] from_vec
         cdef vector[cydriver.CUgraphNode] to_vec
@@ -162,13 +200,13 @@ cdef class _AdjacencySetCore:
         for i in range(n):
             self._resolve_edge(<GraphNode>nodes[i], &from_vec[i], &to_vec[i])
         with nogil:
-            HANDLE_RETURN(_add_edge(
+            HANDLE_RETURN(driver_remove_edges(
                 as_cu(self._h_graph), from_vec.data(), to_vec.data(), n))
 
 
 # ---- driver wrappers: absorb CUDA version differences ----
 
-cdef cydriver.CUresult _get_preds(
+cdef inline cydriver.CUresult driver_get_preds(
         cydriver.CUgraphNode node, cydriver.CUgraphNode* out,
         size_t* count) noexcept nogil:
     IF CUDA_CORE_BUILD_MAJOR >= 13:
@@ -177,7 +215,7 @@ cdef cydriver.CUresult _get_preds(
         return cydriver.cuGraphNodeGetDependencies(node, out, count)
 
 
-cdef cydriver.CUresult _get_succs(
+cdef inline cydriver.CUresult driver_get_succs(
         cydriver.CUgraphNode node, cydriver.CUgraphNode* out,
         size_t* count) noexcept nogil:
     IF CUDA_CORE_BUILD_MAJOR >= 13:
@@ -186,7 +224,7 @@ cdef cydriver.CUresult _get_succs(
         return cydriver.cuGraphNodeGetDependentNodes(node, out, count)
 
 
-cdef cydriver.CUresult _add_edge(
+cdef inline cydriver.CUresult driver_add_edges(
         cydriver.CUgraph graph, cydriver.CUgraphNode* from_arr,
         cydriver.CUgraphNode* to_arr, size_t count) noexcept nogil:
     IF CUDA_CORE_BUILD_MAJOR >= 13:
@@ -197,7 +235,7 @@ cdef cydriver.CUresult _add_edge(
             graph, from_arr, to_arr, count)
 
 
-cdef cydriver.CUresult _remove_edge(
+cdef inline cydriver.CUresult driver_remove_edges(
         cydriver.CUgraph graph, cydriver.CUgraphNode* from_arr,
         cydriver.CUgraphNode* to_arr, size_t count) noexcept nogil:
     IF CUDA_CORE_BUILD_MAJOR >= 13:

cuda_core/cuda/core/_graph/_graph_def/_graph_node.pyx

@@ -57,7 +57,7 @@ from cuda.core._graph._utils cimport (
 )
 
 from cuda.core import Device
-from cuda.core._graph._graph_def._adjacency_set import AdjacencySet
+from cuda.core._graph._graph_def._adjacency_set_proxy import AdjacencySetProxy
 from cuda.core._utils.cuda_utils import driver, handle_return
 
 
@@ -124,30 +124,35 @@ cdef class GraphNode:
         return as_py(self._h_node)
 
     def discard(self):
-        """Discard this node and remove all its edges from the parent graph."""
+        """Discard this node and remove all its edges from the parent graph.
+
+        Safe to call on an already-discarded node (no-op).
+        """
+        if self not in self.graph.nodes():
+            return
         cdef cydriver.CUgraphNode node = as_cu(self._h_node)
         with nogil:
             HANDLE_RETURN(cydriver.cuGraphDestroyNode(node))
 
     @property
     def pred(self):
         """A mutable set-like view of this node's predecessors."""
-        return AdjacencySet(self, False)
+        return AdjacencySetProxy(self, False)
 
     @pred.setter
     def pred(self, value):
-        p = AdjacencySet(self, False)
+        p = AdjacencySetProxy(self, False)
         p.clear()
         p.update(value)
 
     @property
     def succ(self):
         """A mutable set-like view of this node's successors."""
-        return AdjacencySet(self, True)
+        return AdjacencySetProxy(self, True)
 
     @succ.setter
     def succ(self, value):
-        s = AdjacencySet(self, True)
+        s = AdjacencySetProxy(self, True)
         s.clear()
         s.update(value)
 

cuda_core/tests/graph/test_graphdef_mutation.py

@@ -4,12 +4,15 @@
 """Tests for mutating a graph definition (edge changes, node removal)."""
 
 import numpy as np
+import pytest
 from helpers.collection_interface_testers import assert_mutable_set_interface
 from helpers.graph_kernels import compile_parallel_kernels
 from helpers.marks import requires_module
 
+from cuda.bindings import driver
 from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource
 from cuda.core._graph._graph_def import GraphDef, KernelNode, MemsetNode
+from cuda.core._utils.cuda_utils import CUDAError
 
 
 class YRig:
@@ -29,6 +32,9 @@ class YRig:
     Node r computes result ``combine(R, A, B) = (A << 16) | (B & 0xFFFF)``,
     encoding both arms' results into a single int. j is a joining (empty) node
     preceeding r.
+
+    The affine operation a * m + b is noncommutative, so we can be sure the
+    graph has exactly the topology we expect by checking the final value.
     """
 
     def __init__(self):
@@ -148,8 +154,8 @@ def test_baseline(self, init_cuda):
         rig.close()
 
     def test_discard_a1(self, init_cuda):
-        """Discard a1 (creates a race). Arm b yields the expected value, and the
-        final result is correctly ordered after b."""
+        """Discard a1 (creates a race on arm a). Arm b yields the expected
+        value, and the final step is correctly ordered after b completes."""
         rig = YRig()
         rig.a[1].discard()
         rig.run()
@@ -204,7 +210,7 @@ def test_insert_b(self, init_cuda):
 
 
 def test_adjacency_set_interface(init_cuda):
-    """Exercise every MutableSet method on AdjacencySet."""
+    """Exercise every MutableSet method on AdjacencySetProxy."""
     g = GraphDef()
     hub = g.join()
     items = [g.join() for _ in range(5)]
@@ -260,6 +266,67 @@ def test_adjacency_set_property_setter(init_cuda):
     assert hub.pred == set()
 
 
+def test_discarded_node(init_cuda):
+    """Test uses of discarded nodes."""
+    mr = LegacyPinnedMemoryResource()
+    buf = mr.allocate(4)
+    arr = np.from_dlpack(buf).view(np.int32)
+    arr[:] = 0
+    ptr = arr[0:].ctypes.data
+
+    g = GraphDef()
+    a = g.memset(ptr, 0, 4)
+    b = a.memset(ptr, 42, 4)
+
+    assert b in g.nodes()
+    assert (a, b) in g.edges()
+
+    b.discard()
+
+    # b is removed from the graph but still usable
+    assert b not in g.nodes()
+    assert (a, b) not in g.edges()
+    assert isinstance(b, MemsetNode)
+    assert b.type == driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_KERNEL
+    assert b.pred == set()
+    assert b.succ == set()
+    assert b.handle != 0
+    assert b.dptr == ptr
+    assert b.value == 42
+    assert b.width == 4
+
+    # Repeated discard succeeds quietly.
+    b.discard()
+
+
+def test_add_wrong_type(init_cuda):
+    """Adding a non-GraphNode raises TypeError."""
+    g = GraphDef()
+    node = g.join()
+    with pytest.raises(TypeError, match="expected GraphNode"):
+        node.succ.add("not a node")
+    with pytest.raises(TypeError, match="expected GraphNode"):
+        node.succ.add(42)
+
+
+def test_cross_graph_edge(init_cuda):
+    """Adding an edge to a node from a different graph raises CUDAError."""
+    g1 = GraphDef()
+    g2 = GraphDef()
+    a = g1.join()
+    b = g2.join()
+    with pytest.raises(CUDAError):
+        a.succ.add(b)
+
+
+def test_self_edge(init_cuda):
+    """Adding a self-edge raises CUDAError."""
+    g = GraphDef()
+    node = g.join()
+    with pytest.raises(CUDAError):
+        node.succ.add(node)
+
+
 @requires_module(np, "2.1")
 def test_convert_linear_to_fan_in(init_cuda):
     """Chain four computations sequentially, then rewire so all pairs run in
@@ -293,17 +360,22 @@ def test_convert_linear_to_fan_in(init_cuda):
     g = GraphDef()
     prev = g
     for i, val in enumerate(values):
-        prev = prev.memset(ptrs[i], val, 1).launch(config, affine, ptrs[i], 2, 1)
+        prev = prev.memset(ptrs[i], val, 1)
+        prev = prev.launch(config, affine, ptrs[i], 2, 1)
     reduce_node = g.launch(config, reduce_kern, ptrs[4], ptrs[0], 4)
 
     # Rewire:
     #   - drop preds from memsets
-    #   - connect results to reduction
+    #   - connect kernel launches to the reduction
+    assert len(g.edges()) == 7
+
     for node in g.nodes():
-        if isinstance(node, MemsetNode) and len(node.pred):
-            node.pred = set()
+        if isinstance(node, MemsetNode):
+            node.pred.clear()
         elif isinstance(node, KernelNode) and node != reduce_node:
-            node.succ = {reduce_node}
+            node.succ.add(reduce_node)
+
+    assert len(g.edges()) == 8
 
     stream = Device().create_stream()
     graph = g.instantiate()