change findall into for comprehension (jump-dev#1224)

* change findall into for comprehension

* change indexing for iteration on terms

* term by term iteration in getindex

* fix const out of loop

* simplify without outerloop

* sorted output

* sort indices

* sort terms by output index

* revert sorting

* use findall with fix

* make tests order-independent

* remove sort of RHSs

* white space

Author: matbesancon

src/Utilities/functions.jl

@@ -355,8 +355,7 @@ function scalar_terms_at_index(
     terms::Vector{<:Union{MOI.VectorAffineTerm,MOI.VectorQuadraticTerm}},
     i::Int,
 )
-    I = findall(t -> t.output_index == i, terms)
-    return map(i -> terms[i].scalar_term, I)
+    return [term.scalar_term for term in terms if term.output_index == i]
 end
 function Base.getindex(it::ScalarFunctionIterator{<:VAF}, i::Integer)
     return SAF(scalar_terms_at_index(it.f.terms, i), it.f.constants[i])
@@ -378,11 +377,14 @@ function Base.getindex(
     I::AbstractVector,
 ) where {T}
     terms = MOI.VectorAffineTerm{T}[]
-    constant = Vector{T}(undef, length(I))
-    for (i, j) in enumerate(I)
-        g = it[j]
-        append!(terms, map(t -> MOI.VectorAffineTerm(i, t), g.terms))
-        constant[i] = g.constant
+    # assume at least one term per index
+    sizehint!(terms, length(I))
+    constant = it.f.constants[I]
+    for term in it.f.terms
+        idx = findfirst(Base.Fix1(==, term.output_index), I)
+        if idx !== nothing
+            push!(terms, MOI.VectorAffineTerm(idx, term.scalar_term))
+        end
     end
     return VAF(terms, constant)
 end

test/Utilities/functions.jl

@@ -246,11 +246,11 @@ end
     @test g.constant == 5
     h = it[[3, 1]]
     @test h isa MOI.VectorAffineFunction
-    @test h.terms == MOI.VectorAffineTerm.([1, 1, 2, 2, 2], MOI.ScalarAffineTerm.([2, 6, 7, 1, 4], [z, x, y, z, x]))
+    @test sort(h.terms, by=t->t.output_index) == MOI.VectorAffineTerm.([1, 1, 2, 2, 2], MOI.ScalarAffineTerm.([2, 6, 7, 1, 4], [z, x, y, z, x]))
     @test MOIU.constant_vector(h) == [5, 2]
     F = MOIU.operate(vcat, Int, it[[1, 2]], it[3])
     @test F isa MOI.VectorAffineFunction{Int}
-    @test F.terms == MOI.VectorAffineTerm.([1, 1, 1, 2, 2, 2, 2, 3, 3], MOI.ScalarAffineTerm.([7, 1, 4, 1, 9, 3, 1, 2, 6], [y, z, x, x, z, y, y, z, x]))
+    @test sort(F.terms, by=t->t.output_index) == MOI.VectorAffineTerm.([1, 1, 1, 2, 2, 2, 2, 3, 3], MOI.ScalarAffineTerm.([7, 1, 4, 1, 9, 3, 1, 2, 6], [y, z, x, x, z, y, y, z, x]))
     @test MOIU.constant_vector(F) == MOIU.constant_vector(f)
 end
 @testset "Indexing on VectorQuadraticFunction" begin