From c9703248f886abef27dd1ea93ff08156c6b2719c Mon Sep 17 00:00:00 2001
From: Jordan Jalving <jhjalving@gmail.com>
Date: Wed, 21 Aug 2024 13:02:17 -0700
Subject: [PATCH] fix inefficient variable check; refactor all_variables (#116)

---
 src/backends/moi_backend.jl |  9 +++++++++
 src/optinode.jl             | 13 +++++++++----
 test/test_optigraph.jl      |  5 +++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/backends/moi_backend.jl b/src/backends/moi_backend.jl
index dfbf9a2..76e3ebf 100644
--- a/src/backends/moi_backend.jl
+++ b/src/backends/moi_backend.jl
@@ -235,6 +235,15 @@ function MOI.get(backend::GraphMOIBackend, attr::MOI.NumberOfVariables, node::Op
     return length(backend.node_variables[node])
 end
 
+function MOI.get(backend::GraphMOIBackend, attr::MOI.ListOfVariableIndices, node::OptiNode)
+    graph_indices = backend.node_variables[node]
+    var_indices = MOI.VariableIndex[]
+    for graph_index in graph_indices
+        push!(var_indices, backend.graph_to_element_map[graph_index].index)
+    end
+    return var_indices
+end
+
 function MOI.get(
     backend::GraphMOIBackend, attr::MOI.ListOfConstraintTypesPresent, element::OptiElement
 )
diff --git a/src/optinode.jl b/src/optinode.jl
index 87c92b2..a2486cc 100644
--- a/src/optinode.jl
+++ b/src/optinode.jl
@@ -117,9 +117,8 @@ function JuMP.num_variables(node::OptiNode)
 end
 
 function JuMP.all_variables(node::OptiNode)
-    gb = graph_backend(node)
-    graph_indices = gb.node_variables[node]
-    return getindex.(Ref(gb.graph_to_element_map), graph_indices)
+    var_inds = MOI.get(node, MOI.ListOfVariableIndices(), node)
+    return NodeVariableRef.(Ref(node), var_inds)
 end
 
 function JuMP.delete(node::OptiNode, cref::ConstraintRef)
@@ -211,7 +210,13 @@ function _check_node_variables(
         NodeVariableRef,JuMP.GenericAffExpr,JuMP.GenericQuadExpr,JuMP.GenericNonlinearExpr
     },
 )
-    return isempty(setdiff(_extract_variables(jump_func), JuMP.all_variables(node)))
+    extract_vars = _extract_variables(jump_func)
+    for var in extract_vars
+        if var.node != node
+            error("Variable $var does not belong to node $node")
+        end
+    end
+    return nothing
 end
 
 ### Objective
diff --git a/test/test_optigraph.jl b/test/test_optigraph.jl
index 30c3b8a..aa4eaff 100644
--- a/test/test_optigraph.jl
+++ b/test/test_optigraph.jl
@@ -293,6 +293,11 @@ function test_variable_constraints()
     set_start_value(n2[:x], 3.0)
     @test start_value(n2[:x]) == 3.0
 
+    # variable not owned
+    @test_throws ErrorException("Variable $(n1[:x]) does not belong to node $n2") @constraint(
+        n2, n1[:x] >= 0
+    )
+
     # bounds
     @test has_lower_bound(n1[:x]) == true
     @test has_upper_bound(n1[:x]) == false