From 89be8840530d2ea0616f89c8ed9c80e49ca0afec Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Mon, 15 Apr 2024 09:24:53 +1200
Subject: [PATCH] [FileFormats.CBF] write some VectorOfVariables constraints in
 VAR block (#2478)

---
 src/FileFormats/CBF/write.jl | 79 ++++++++++++++++++++++++++++++++++--
 test/FileFormats/CBF/CBF.jl  | 62 ++++++++++++++++++++++++----
 2 files changed, 129 insertions(+), 12 deletions(-)

diff --git a/src/FileFormats/CBF/write.jl b/src/FileFormats/CBF/write.jl
index 2bc4e26541..c16b7c71f7 100644
--- a/src/FileFormats/CBF/write.jl
+++ b/src/FileFormats/CBF/write.jl
@@ -19,6 +19,9 @@ mutable struct _CBFDataStructure
     bcoord::Vector{Tuple{Int,Float64}}
     hcoord::Vector{Tuple{Int,Int,Int,Int,Float64}}
     dcoord::Vector{Tuple{Int,Int,Int,Float64}}
+    variables_with_domain::Set{MOI.VariableIndex}
+    variable_cones::Vector{Tuple{Vector{MOI.VariableIndex},String}}
+
     function _CBFDataStructure()
         return new(
             0,
@@ -30,6 +33,8 @@ mutable struct _CBFDataStructure
             Tuple{Int,Float64}[],
             Tuple{Int,Int,Int,Int,Float64}[],
             Tuple{Int,Int,Int,Float64}[],
+            Set{MOI.VariableIndex}(),
+            Tuple{Vector{MOI.VariableIndex},String}[],
         )
     end
 end
@@ -125,6 +130,57 @@ function _add_cones(
     return
 end
 
+function _add_cones(
+    data::_CBFDataStructure,
+    model::Model,
+    ::Type{F},
+    ::Type{S},
+) where {F<:MOI.VectorOfVariables,S}
+    for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
+        f = MOI.get(model, MOI.ConstraintFunction(), ci)
+        is_variable_cone = true
+        for (i, xi) in enumerate(f.variables)
+            if xi in data.variables_with_domain
+                is_variable_cone = false
+                break
+            elseif xi.value != f.variables[1].value + i - 1
+                is_variable_cone = false
+                break
+            end
+            push!(data.variables_with_domain, xi)
+        end
+        str = _cone_string(data, S)
+        if !is_variable_cone
+            _add_function(data, f, S)
+            set = MOI.get(model, MOI.ConstraintSet(), ci)
+            push!(data.cones, (str, MOI.dimension(set)))
+        else
+            push!(data.variable_cones, (f.variables, str))
+        end
+    end
+    return
+end
+
+function _add_cones(
+    data::_CBFDataStructure,
+    model::Model,
+    ::Type{F},
+    ::Type{S},
+) where {
+    F<:MOI.VectorOfVariables,
+    S<:Union{MOI.ExponentialCone,MOI.DualExponentialCone},
+}
+    # The Exponential cone in MOI and CBF are reversed. Instead of dealing with
+    # this complexity, just write them out as `Ax + b in K` constraints.
+    # TODO(odow): we should support this at some point. See #2478
+    for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
+        f = MOI.get(model, MOI.ConstraintFunction(), ci)
+        _add_function(data, f, S)
+        push!(data.cones, (_cone_string(data, S), 3))
+    end
+    return
+end
+
 function _add_cones(
     data::_CBFDataStructure,
     model::Model,
@@ -252,11 +308,26 @@ function _write_POWCONES(io::IO, model::Model, S, keyword)
     return
 end
 
-function _write_VAR(io::IO, model::Model)
+function _write_VAR(io::IO, model::Model, data)
     num_var = MOI.get(model, MOI.NumberOfVariables())
+    cones = Tuple{String,Int}[]
+    current_variable = 0
+    for (f, str) in sort!(data.variable_cones; by = x -> first(x[1]).value)
+        offset = first(f).value - current_variable - 1
+        if offset > 0
+            push!(cones, ("F", offset))
+        end
+        push!(cones, (str, length(f)))
+        current_variable = last(f).value
+    end
+    if current_variable < num_var
+        push!(cones, ("F", num_var - current_variable))
+    end
     println(io, "VAR")
-    println(io, num_var, " 1")
-    println(io, "F ", num_var)
+    println(io, num_var, " ", length(cones))
+    for (K, n) in cones
+        println(io, K, " ", n)
+    end
     println(io)
     return
 end
@@ -417,7 +488,7 @@ function Base.write(io::IO, model::Model)
     ###
     _write_OBJSENSE(io, model)
     # _write_PSDVAR
-    _write_VAR(io, model)
+    _write_VAR(io, model, data)
     _write_INT(io, model)
     _write_PSDCON(io, data)
     _write_CON(io, data)
diff --git a/test/FileFormats/CBF/CBF.jl b/test/FileFormats/CBF/CBF.jl
index 340c496807..f769273391 100644
--- a/test/FileFormats/CBF/CBF.jl
+++ b/test/FileFormats/CBF/CBF.jl
@@ -218,14 +218,6 @@ const _WRITE_READ_MODELS = [
     variables: x, y
     minobjective: x
     c1: [x, y] in Zeros(2)
-""",
-    ),
-    (
-        "VectorOfVariables in Reals",
-        """
-    variables: x, y
-    minobjective: x
-    c1: [x, y] in Reals(2)
 """,
     ),
     (
@@ -454,6 +446,60 @@ function test_example_models()
     end
 end
 
+function test_write_variable_cones()
+    model = CBF.Model()
+    for set in (
+        MOI.Zeros(2),
+        MOI.Nonnegatives(3),
+        MOI.PowerCone(0.74),
+        MOI.Nonpositives(1),
+        MOI.DualPowerCone(0.25),
+        MOI.Nonnegatives(3),
+        MOI.PowerCone(0.5),
+        MOI.SecondOrderCone(3),
+    )
+        _ = MOI.add_constrained_variables(model, set)
+    end
+    io = IOBuffer()
+    write(io, model)
+    seekstart(io)
+    @test read(io, String) == """
+    VER
+    3
+
+    POWCONES
+    2 4
+    2
+    0.74
+    0.26
+    2
+    0.5
+    0.5
+
+    POW*CONES
+    1 2
+    2
+    0.25
+    0.75
+
+    OBJSENSE
+    MIN
+
+    VAR
+    21 8
+    L= 2
+    L+ 3
+    @0:POW 3
+    L- 1
+    @0:POW* 3
+    L+ 3
+    @1:POW 3
+    Q 3
+
+    """
+    return
+end
+
 function runtests()
     for name in names(@__MODULE__, all = true)
         if startswith("$(name)", "test_")