From 8079a8e932464808d60a23b26084baaa9ff66036 Mon Sep 17 00:00:00 2001
From: Dimitri Alston <123396563+DimitriAlston@users.noreply.github.com>
Date: Thu, 7 Nov 2024 09:11:56 -0500
Subject: [PATCH] Add support for `MOI.UserDefinedFunction`

---
 src/eago_optimizer/moi_wrapper.jl | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/eago_optimizer/moi_wrapper.jl b/src/eago_optimizer/moi_wrapper.jl
index 6727a4e0..a1c0dafe 100644
--- a/src/eago_optimizer/moi_wrapper.jl
+++ b/src/eago_optimizer/moi_wrapper.jl
@@ -314,4 +314,21 @@ function MOI.get(m::Optimizer, ::Type{VI}, name::String)
     else
         return index_storage
     end
-end
\ No newline at end of file
+end
+
+#####
+##### Support and set user-defined functions
+#####
+MOI.supports(m::Optimizer, ::MOI.UserDefinedFunction) = true
+
+function MOI.set(m::Optimizer, udf::MOI.UserDefinedFunction, f)
+    if isnothing(m._input_problem._nlp_data)
+        model = MOI.Nonlinear.Model()
+        backend = MOI.Nonlinear.SparseReverseMode()
+        vars = MOI.get(m, MOI.ListOfVariableIndices())
+        evaluator = MOI.Nonlinear.Evaluator(model, backend, vars) 
+        m._input_problem._nlp_data = MOI.NLPBlockData(evaluator)
+    end
+    MOI.Nonlinear.register_operator(m._input_problem._nlp_data.evaluator.model, udf.name, udf.arity, f...)
+    return nothing
+end