use WeakKeyIdDict

Project.toml

@@ -5,6 +5,7 @@ version = "0.16.0"
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"

src/ComplexTape.jl

@@ -1,4 +1,4 @@
-struct ComplexTape{T}
+mutable struct ComplexTape{T}
     real_tape::SparseTape{T}
     imag_tape::SparseTape{T}
 

src/Context.jl

@@ -1,4 +1,4 @@
-struct Context{T,M}
+mutable struct Context{T,M}
     # MOI model
     model::M
 
@@ -18,7 +18,7 @@ struct Context{T,M}
     detected_infeasible_during_formulation::Ref{Bool}
 
     # Cache
-    conic_form_cache::IdDict{Any, Any}
+    conic_form_cache::DataStructures.WeakKeyIdDict{Any, Any}
 end
 
 function Context{T}(optimizer) where {T}

src/Convex.jl

@@ -7,6 +7,7 @@ using SparseArrays
 using LDLFactorizations
 using AbstractTrees: AbstractTrees, children
 using SuiteSparseGraphBLAS
+using DataStructures
 
 using MathOptInterface
 const MOI = MathOptInterface

src/SparseTape.jl

@@ -32,7 +32,7 @@ end
 # end
 # SparseAffineOperation(A, b) = SparseAffineOperation(sparse(A), b)
 
-struct SparseTape{T}
+mutable struct SparseTape{T}
     operations::Vector{SparseAffineOperation{T}}
     variables::Vector{MOI.VariableIndex}
 end

src/VectorAffineFunctionAsMatrix.jl

@@ -24,17 +24,15 @@ end
 # convert to a usual VAF
 function to_vaf(vaf_as_matrix::VectorAffineFunctionAsMatrix{T}) where {T}
     I, J, V = findnz(vaf_as_matrix.aff.matrix)
-    vats = MOI.VectorAffineTerm{T}[]
-    for n in eachindex(I, J, V)
+    vats = Vector{MOI.VectorAffineTerm{T}}(undef, length(I))
+    for (idx, n) in enumerate(eachindex(I, J, V))
         i = I[n]
         j = J[n]
         v = V[n]
-        push!(
-            vats,
-            MOI.VectorAffineTerm{T}(
-                i,
-                MOI.ScalarAffineTerm{T}(v, vaf_as_matrix.variables[j]),
-            ),
+
+        vats[idx] = MOI.VectorAffineTerm{T}(
+            i,
+            MOI.ScalarAffineTerm{T}(v, vaf_as_matrix.variables[j]),
         )
     end
 

src/atoms/affine/index.jl

@@ -71,9 +71,9 @@ function _conic_form!(context::Context{T}, x::IndexAtom) where {T}
             end
         end
 
-        index_matrix = sparse(1:sz, J, one(T), m, n)
+        index_matrix = GBMatrix{T,T}(collect(1:sz), J, one(T), m, n)
     else
-        index_matrix = sparse(1:length(x.inds), x.inds, one(T), m, n)
+        index_matrix = GBMatrix{T,T}(collect(1:length(x.inds)), collect(x.inds), one(T), m, n)
     end
 
     return operate(add_operation, T, sign(x), index_matrix, obj)

src/constant.jl

@@ -23,7 +23,7 @@ function _sign(x::Value)
     end
 end
 
-struct Constant{T<:Real} <: AbstractExpr
+mutable struct Constant{T<:Real} <: AbstractExpr
     head::Symbol
     id_hash::UInt64
     value::Matrix{T}
@@ -46,7 +46,7 @@ struct Constant{T<:Real} <: AbstractExpr
 end
 # Constant(x::Constant) = x
 
-struct ComplexConstant{T<:Real} <: AbstractExpr
+mutable struct ComplexConstant{T<:Real} <: AbstractExpr
     head::Symbol
     id_hash::UInt64
     size::Tuple{Int,Int}
@@ -74,7 +74,7 @@ function evaluate(c::ComplexConstant)
     return evaluate(c.real_constant) + im * evaluate(c.imag_constant)
 end
 
-struct ComplexStructOfVec{T<:Real}
+mutable struct ComplexStructOfVec{T<:Real}
     real_vec::Vector{T}
     imag_vec::Vector{T}
 end

src/variable_template.jl

@@ -75,6 +75,18 @@ function _conic_form!(context::Context, a::AbstractVariable)
 end
 
 function conic_form!(context::Context, a::AbstractExpr)
-    d = context.conic_form_cache
-    return get!(() -> _conic_form!(context, a), d, a)
+
+    # Nicer implementation
+    # d = context.conic_form_cache
+    # return get!(() -> _conic_form!(context, a), d, a)
+
+    # Avoid closure
+    wkh = context.conic_form_cache
+    default = () -> _conic_form!(context, a)
+    key = a
+    return Base.@lock wkh.lock begin
+        get!(default, wkh.ht, DataStructures.WeakRefForWeakDict(key))
+    end
+
+    return
 end