From ab36f80ad9d3224fef224142511d5e85a6f07fd4 Mon Sep 17 00:00:00 2001
From: David Cole <76058020+dlcole3@users.noreply.github.com>
Date: Thu, 22 Aug 2024 18:59:46 -0500
Subject: [PATCH] Added `JuMP.delete` extension for link constraints (#118)

* Added delete extension for link constraints

* added unregister functions

* added unregister function for edges

* refactored JuMP.delete methods
---
 src/optiedge.jl        |  7 +++++++
 src/optielement.jl     | 12 ++++++++++++
 src/optigraph.jl       |  4 ++++
 src/optinode.jl        | 12 ++----------
 test/test_optigraph.jl | 20 ++++++++++++++++++++
 5 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/src/optiedge.jl b/src/optiedge.jl
index 08281b2..2bed1d9 100644
--- a/src/optiedge.jl
+++ b/src/optiedge.jl
@@ -76,6 +76,13 @@ function JuMP.all_variables(edge::OptiEdge)
     return unique(vars)
 end
 
+function _set_dirty(edge::OptiEdge)
+    for graph in containing_optigraphs(edge)
+        graph.is_model_dirty = true
+    end
+    return nothing
+end
+
 ### Edge Constraints
 
 # NOTE: could use one method for node and edge
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 d99322e..e125e19 100644
--- a/src/optigraph.jl
+++ b/src/optigraph.jl
@@ -1423,3 +1423,7 @@ function _set_objective_coefficient(
     )
     return nothing
 end
+
+function JuMP.unregister(graph::OptiGraph, key::Symbol)
+    delete!(object_dictionary(graph), key)
+end
diff --git a/src/optinode.jl b/src/optinode.jl
index c9640b7..425dacd 100644
--- a/src/optinode.jl
+++ b/src/optinode.jl
@@ -121,16 +121,8 @@ 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
+function JuMP.unregister(node::OptiNode, key::Symbol)
+    delete!(object_dictionary(node), (node, key))
 end
 
 ### Duals
diff --git a/test/test_optigraph.jl b/test/test_optigraph.jl
index aa4eaff..48dc865 100644
--- a/test/test_optigraph.jl
+++ b/test/test_optigraph.jl
@@ -440,6 +440,26 @@ function test_nlp_exceptions()
     @test_throws Exception @NLconstraint(graph, graph[1][:x]^3 >= 0)
 end
 
+function test_delete_extensions()
+    graph = OptiGraph()
+    @optinode(graph, nodes[1:2])
+
+    @variable(nodes[1], x >= 1)
+    @variable(nodes[2], x >= 2)
+    @constraint(nodes[1], n_con, nodes[1][:x]^2 >= 1)
+    @linkconstraint(graph, l_con, nodes[1][:x] + nodes[2][:x] == 4)
+
+    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(edge, l_con)
+    @test !(l_con in all_link_constraints(graph))
+end
+
 function run_tests()
     for name in names(@__MODULE__; all=true)
         if !startswith("$(name)", "test_")