fix inefficient variable check; refactor all_variables (plasmo-dev#116)

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
 )

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

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