From 8f44aa30a5d5e7f9aeb8558cb5a5dcfe40c53a32 Mon Sep 17 00:00:00 2001
From: jalving <jhjalving@gmail.com>
Date: Mon, 9 Oct 2023 21:29:55 -0700
Subject: [PATCH] support lazy constraints in plasmo

---
 src/moi_backend_graph.jl   |  5 ++++-
 src/optigraph.jl           | 40 ++++++++++++++++++++++++++++++++++++++
 src/optimizer_interface.jl |  8 ++++++++
 3 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/src/moi_backend_graph.jl b/src/moi_backend_graph.jl
index 10f8cfb..d20b6a6 100644
--- a/src/moi_backend_graph.jl
+++ b/src/moi_backend_graph.jl
@@ -103,15 +103,18 @@ function MOI.set(graph_backend::GraphBackend, attr::MOI.AbstractOptimizerAttribu
     return nothing
 end
 
-# TODO: properly support variable and constraint attributes
 function MOI.get(
     graph_backend::GraphBackend, attr::MOI.VariablePrimalStart, idx::MOI.VariableIndex
 )
     return MOI.get(graph_backend.model_cache, attr, idx)
 end
 
+#TODO: properly support variable and constraint attributes
 #MOI.set(graph_backend::GraphBackend,attr::MOI.AnyAttribute,args...) = MOI.set(graph_backend.optimizer,attr,args...)
+
 MOIU.state(graph_backend::GraphBackend) = graph_backend.state
+
+# TODO: decide whether we support graph modes
 #MOIU.mode(graph_backend::GraphBackend) = graph_backend.mode
 
 """
diff --git a/src/optigraph.jl b/src/optigraph.jl
index 84c9cb2..62ffcb1 100644
--- a/src/optigraph.jl
+++ b/src/optigraph.jl
@@ -223,6 +223,13 @@ function add_node!(graph::OptiGraph, optinode::OptiNode)
     return optinode
 end
 
+function add_node!(graph::OptiGraph, m::JuMP.Model, label::String)
+    optinode = add_node!(graph)
+    set_model(optinode, m)
+    optinode.label = label
+    return optinode
+end
+
 """
     optinodes(graph::OptiGraph)::Vector{OptiNode}
 
@@ -927,6 +934,19 @@ function JuMP.start_value(graph::OptiGraph, variable::JuMP.VariableRef)
     return MOI.get(backend(graph), MOI.VariablePrimalStart(), var_idx)
 end
 
+function JuMP.set_attribute(
+    graph::OptiGraph,
+    attr::MOI.AbstractModelAttribute,
+    value
+)
+    MOI.set(graph, attr, value)
+    return
+end
+
+function JuMP.solver_name(graph::OptiGraph)
+    return MOI.get(graph.moi_backend, MOI.SolverName())
+end
+
 """
     JuMP.termination_status(graph::OptiGraph)
 
@@ -936,6 +956,26 @@ function JuMP.termination_status(graph::OptiGraph)
     return MOI.get(graph.moi_backend, MOI.TerminationStatus())
 end
 
+function JuMP.primal_status(graph::OptiGraph)
+    return MOI.get(graph.moi_backend, MOI.PrimalStatus())
+end
+
+function JuMP.dual_status(graph::OptiGraph)
+    return MOI.get(graph.moi_backend, MOI.DualStatus())
+end
+
+function JuMP.callback_value(cb_data, x::GenericVariableRef, graph::OptiGraph)
+    return MOI.get(
+        graph.moi_backend.optimizer,
+        MOI.CallbackVariablePrimal(cb_data),
+        index(x),
+    )
+end
+
+function JuMP.callback_node_status(cb_data, graph::OptiGraph)
+    return MOI.get(graph.moi_backend.optimizer, MOI.CallbackNodeStatus(cb_data))
+end
+
 ####################################
 #Print Functions
 ####################################
diff --git a/src/optimizer_interface.jl b/src/optimizer_interface.jl
index 863b7ac..3945dc2 100644
--- a/src/optimizer_interface.jl
+++ b/src/optimizer_interface.jl
@@ -355,6 +355,14 @@ function MOI.set(graph::OptiGraph, attr::MOI.AbstractModelAttribute, value)
     return MOI.set(backend(graph), attr, value)
 end
 
+function MOI.submit(
+    graph::OptiGraph,
+    cb::MOI.LazyConstraint,
+    con::ScalarConstraint,
+)
+    return MOI.submit(graph.moi_backend.optimizer, cb, moi_function(con.func), con.set)
+end
+
 #######################################################
 #Optinode optimizer interface
 #######################################################