From a049bfe997f9c3b2b31e6bdac40bce27d96534e0 Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Fri, 4 Aug 2023 16:10:23 +1200
Subject: [PATCH] [Bridges] various updates based on VectorNonlinearFunction
 changes (#2244)

---
 src/Bridges/Constraint/bridges/scalarize.jl  | 35 ++++++--------------
 src/Bridges/Constraint/bridges/soc_to_psd.jl |  2 +-
 src/Bridges/Constraint/bridges/vectorize.jl  |  7 ++--
 3 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/src/Bridges/Constraint/bridges/scalarize.jl b/src/Bridges/Constraint/bridges/scalarize.jl
index cd9f539685..0038cb4553 100644
--- a/src/Bridges/Constraint/bridges/scalarize.jl
+++ b/src/Bridges/Constraint/bridges/scalarize.jl
@@ -43,14 +43,11 @@ function bridge_constraint(
     f::MOI.AbstractVectorFunction,
     set::MOI.Utilities.VectorLinearSet,
 ) where {T,F,S}
-    dimension = MOI.output_dimension(f)
-    constants = MOI.constant(f, T)
-    new_f = MOI.Utilities.scalarize(f, true)
-    constraints = Vector{MOI.ConstraintIndex{F,S}}(undef, dimension)
-    for i in 1:dimension
-        constraints[i] = MOI.add_constraint(model, new_f[i], S(-constants[i]))
-    end
-    return ScalarizeBridge{T,F,S}(constraints, constants)
+    constraints = MOI.ConstraintIndex{F,S}[
+        MOI.Utilities.normalize_and_add_constraint(model, fi, S(zero(T)))
+        for fi in MOI.Utilities.eachscalar(f)
+    ]
+    return ScalarizeBridge{T,F,S}(constraints, MOI.constant(f, T))
 end
 
 function MOI.supports_constraint(
@@ -270,24 +267,14 @@ function MOI.set(
     bridge::ScalarizeBridge{T,F,S},
     func,
 ) where {T,F,S}
-    old_constants = bridge.constants
     bridge.constants = MOI.constant(func, T)
-    new_func = MOI.Utilities.scalarize(func, true)
-    MOI.set.(
-        model,
-        MOI.ConstraintFunction(),
-        bridge.scalar_constraints,
-        new_func,
-    )
-    for i in eachindex(bridge.constants)
-        if bridge.constants[i] != old_constants[i]
-            MOI.set(
-                model,
-                MOI.ConstraintSet(),
-                bridge.scalar_constraints[i],
-                S(-bridge.constants[i]),
-            )
+    for (i, fi) in enumerate(MOI.Utilities.eachscalar(func))
+        ci = bridge.scalar_constraints[i]
+        if !iszero(bridge.constants[i])
+            fi = MOI.Utilities.operate!(-, T, fi, bridge.constants[i])
         end
+        MOI.set(model, MOI.ConstraintFunction(), ci, fi)
+        MOI.set(model, MOI.ConstraintSet(), ci, S(-bridge.constants[i]))
     end
     return
 end
diff --git a/src/Bridges/Constraint/bridges/soc_to_psd.jl b/src/Bridges/Constraint/bridges/soc_to_psd.jl
index a8482e27ee..b9cdb70aa7 100644
--- a/src/Bridges/Constraint/bridges/soc_to_psd.jl
+++ b/src/Bridges/Constraint/bridges/soc_to_psd.jl
@@ -19,7 +19,7 @@ of the matrix:
 """
 function _soc_to_psd_matrix(
     ::Type{T},
-    f::Union{MOI.AbstractVectorFunction,AbstractVector{T}},
+    f::Union{MOI.AbstractVectorFunction,AbstractVector},
     g::Union{MOI.AbstractScalarFunction,T},
 ) where {T}
     F = MOI.Utilities.promote_operation(vcat, T, typeof(g), T)
diff --git a/src/Bridges/Constraint/bridges/vectorize.jl b/src/Bridges/Constraint/bridges/vectorize.jl
index 9f60285e28..2a26007c5f 100644
--- a/src/Bridges/Constraint/bridges/vectorize.jl
+++ b/src/Bridges/Constraint/bridges/vectorize.jl
@@ -49,9 +49,12 @@ function bridge_constraint(
     if !iszero(scalar_const)
         throw(MOI.ScalarFunctionConstantNotZero{T,G,typeof(set)}(scalar_const))
     end
-    vector_f = convert(F, scalar_f)
     set_const = MOI.constant(set)
-    MOI.Utilities.operate_output_index!(-, T, 1, vector_f, set_const)
+    vector_f = MOI.Utilities.operate(
+        vcat,
+        T,
+        MOI.Utilities.operate(-, T, scalar_f, set_const),
+    )
     vector_constraint = MOI.add_constraint(model, vector_f, S(1))
     return VectorizeBridge{T,F,S,G}(vector_constraint, set_const)
 end