From adc32a8e12f835f18d01a50b57c5f0894dd41582 Mon Sep 17 00:00:00 2001
From: dlcole3 <dlcole3@wisc.edu>
Date: Thu, 22 Aug 2024 15:18:37 -0500
Subject: [PATCH] refactored JuMP.delete methods

---
 src/optiedge.jl        | 23 ++++-------------------
 src/optielement.jl     | 12 ++++++++++++
 src/optigraph.jl       |  4 ++--
 src/optinode.jl        | 12 ------------
 test/test_optigraph.jl |  6 ++++--
 5 files changed, 22 insertions(+), 35 deletions(-)

diff --git a/src/optiedge.jl b/src/optiedge.jl
index e6705cc..2bed1d9 100644
--- a/src/optiedge.jl
+++ b/src/optiedge.jl
@@ -76,26 +76,11 @@ function JuMP.all_variables(edge::OptiEdge)
     return unique(vars)
 end
 
-function JuMP.delete(graph::OptiGraph, cref::ConstraintRef)
-    if typeof(JuMP.owner_model(cref)) == OptiNode{OptiGraph}
-        error(
-            "You have passed a node constraint but specified an OptiGraph." *
-            "Use `JuMP.delete(::OptiNode, ::ConstraintRef)` instead",
-        )
-    end
-    if graph !== source_graph(JuMP.owner_model(cref))
-        error(
-            "The constraint reference you are trying to delete does not" *
-            "belong to the specified graph",
-        )
+function _set_dirty(edge::OptiEdge)
+    for graph in containing_optigraphs(edge)
+        graph.is_model_dirty = true
     end
-    _set_dirty(graph)
-    MOI.delete(JuMP.owner_model(cref), cref)
-    #TODO: Probably need to delete the edge altogether if it is the only constraint on the edge
-end
-
-function JuMP.unregister(graph::OptiGraph, key::Symbol)
-    delete!(object_dictionary(graph), key)
+    return nothing
 end
 
 ### Edge Constraints
diff --git a/src/optielement.jl b/src/optielement.jl
index 2a0ec34..3427a74 100644
--- a/src/optielement.jl
+++ b/src/optielement.jl
@@ -144,3 +144,15 @@ function JuMP.all_constraints(
     end
     return constraints
 end
+
+function JuMP.delete(element::OptiElement, cref::ConstraintRef)
+    if element !== JuMP.owner_model(cref)
+        error(
+            "The constraint reference you are trying to delete does not " *
+            "belong to the model.",
+        )
+    end
+    _set_dirty(element)
+    MOI.delete(element, cref)
+    return nothing
+end
diff --git a/src/optigraph.jl b/src/optigraph.jl
index f1ef772..e125e19 100644
--- a/src/optigraph.jl
+++ b/src/optigraph.jl
@@ -1424,6 +1424,6 @@ function _set_objective_coefficient(
     return nothing
 end
 
-function _set_dirty(graph::OptiGraph)
-    graph.is_model_dirty = true
+function JuMP.unregister(graph::OptiGraph, key::Symbol)
+    delete!(object_dictionary(graph), key)
 end
diff --git a/src/optinode.jl b/src/optinode.jl
index 595095a..7680884 100644
--- a/src/optinode.jl
+++ b/src/optinode.jl
@@ -121,18 +121,6 @@ function JuMP.all_variables(node::OptiNode)
     return NodeVariableRef.(Ref(node), var_inds)
 end
 
-function JuMP.delete(node::OptiNode, cref::ConstraintRef)
-    if node !== JuMP.owner_model(cref)
-        error(
-            "The constraint reference you are trying to delete does not " *
-            "belong to the model.",
-        )
-    end
-    _set_dirty(node)
-    MOI.delete(node, cref)
-    return nothing
-end
-
 function JuMP.unregister(node::OptiNode, key::Symbol)
     delete!(object_dictionary(node), (node, key))
 end
diff --git a/test/test_optigraph.jl b/test/test_optigraph.jl
index b48fa4e..48dc865 100644
--- a/test/test_optigraph.jl
+++ b/test/test_optigraph.jl
@@ -449,12 +449,14 @@ function test_delete_extensions()
     @constraint(nodes[1], n_con, nodes[1][:x]^2 >= 1)
     @linkconstraint(graph, l_con, nodes[1][:x] + nodes[2][:x] == 4)
 
-    @test_throws ErrorException JuMP.delete(graph, n_con)
+    edge = JuMP.owner_model(l_con)
+
+    @test_throws ErrorException JuMP.delete(edge, n_con)
     @test_throws ErrorException JuMP.delete(nodes[1], l_con)
 
     JuMP.delete(nodes[1], n_con)
     @test !(n_con in all_constraints(nodes[1]))
-    JuMP.delete(graph, l_con)
+    JuMP.delete(edge, l_con)
     @test !(l_con in all_link_constraints(graph))
 end