Add VectorLazyMap (#2254)

src/Bridges/Constraint/map.jl

@@ -181,15 +181,16 @@ were created with `add_key_for_bridge`.
 function keys_of_type(map::Map, C::Type{MOI.ConstraintIndex{F,S}}) where {F,S}
     return Base.Iterators.Filter(
         ci -> haskey(map, ci),
-        MOI.Utilities.LazyMap{C}(i -> _index(i, F, S), eachindex(map.bridges)),
+        MOI.Utilities.lazy_map(C, i -> _index(i, F, S), eachindex(map.bridges)),
     )
 end
 
 function keys_of_type(
     map::Map,
     C::Type{MOI.ConstraintIndex{MOI.VariableIndex,S}},
 ) where {S}
-    return MOI.Utilities.LazyMap{C}(
+    return MOI.Utilities.lazy_map(
+        C,
         key -> C(key[1]),
         Base.Iterators.Filter(
             key -> key[2] == S,
@@ -239,7 +240,8 @@ Return the list of all keys that correspond to
 [`MOI.VectorOfVariables`](@ref) constraints.
 """
 function vector_of_variables_constraints(map::Map)
-    return MOI.Utilities.LazyMap{MOI.ConstraintIndex{MOI.VectorOfVariables}}(
+    return MOI.Utilities.lazy_map(
+        MOI.ConstraintIndex{MOI.VectorOfVariables},
         i -> MOI.ConstraintIndex{map.constraint_types[i]...}(-i),
         Base.Iterators.Filter(
             i ->

src/Bridges/Variable/map.jl

@@ -142,7 +142,8 @@ end
 function Base.keys(map::Map)
     return Base.Iterators.Filter(
         vi -> haskey(map, vi),
-        MOI.Utilities.LazyMap{MOI.VariableIndex}(
+        MOI.Utilities.lazy_map(
+            MOI.VariableIndex,
             i -> MOI.VariableIndex(-i),
             eachindex(map.bridges),
         ),

src/Utilities/lazy_iterators.jl

@@ -14,15 +14,19 @@ Base.eltype(::EmptyVector{T}) where {T} = T
 
 Base.iterate(::EmptyVector) = nothing
 
+abstract type AbstractLazyMap{T} end
+
 """
-    struct LazyMap{T, VT}
-        f::Function
+    struct LazyMap{T,VT,F}
+        f::F
         data::VT
     end
 
-Iterator over the elements of `data` mapped by `f`. This is similar to
-`Base.Generator(f, data)` except that the `eltype` of a `LazyMap` is given at
-construction while the `eltype` of `Base.Generator(f, data)` is `Any`.
+Iterator over the elements of `data` mapped by `f`.  This is similar to
+[`Utilities.LazyMap`](@ref) except that `VT` should be a subtype of
+`AbstractVector{T}` and `VectorLazyMap` is a subtype of `AbstractVector{T}` as
+well.  Use [`Utilities.lazy_map`](@ref) to create a `VectorLazyMap` or `LazyMap`
+depending on the type of `data`.
 """
 struct LazyMap{T,VT,F}
     f::F
@@ -33,11 +37,48 @@ function LazyMap{T}(f, data) where {T}
     return LazyMap{T,typeof(data),typeof(f)}(f, data)
 end
 
-Base.size(it::LazyMap) = size(it.data)
+"""
+    lazy_map(::Type{T}, f, data)
+
+Return an [`Utilities.VectorLazyMap`](@ref) is an `AbstractVector` and a
+[`Utilities.LazyMap`](@ref) otherwise.  These are similar to
+`Base.Generator(f, data)` except that their `eltype` is fixed to `T`
+while the `eltype` of `Base.Generator(f, data)` is `Any`.
+"""
+lazy_map(::Type{T}, f, data) where {T} = LazyMap{T}(f, data)
 
-Base.length(it::LazyMap) = length(it.data)
+"""
+    struct VectorLazyMap{T,VT<:AbstractVector{T},F} <: AbstractVector{T}
+        f::F
+        data::VT
+    end
 
-function Base.iterate(it::LazyMap, args...)
+Iterator over the elements of `data` mapped by `f`. This is similar to
+[`Utilities.LazyMap`](@ref) except that `VT` should be a subtype of
+`AbstractVector{T}` and `VectorLazyMap` is a subtype of `AbstractVector{T}` as
+well.  Use [`Utilities.lazy_map`](@ref) to create a `VectorLazyMap` or `LazyMap`
+depending on the type of `data`.
+"""
+struct VectorLazyMap{T,VT<:AbstractVector,F} <: AbstractVector{T}
+    f::F
+    data::VT
+end
+
+function VectorLazyMap{T}(f, data) where {T}
+    return VectorLazyMap{T,typeof(data),typeof(f)}(f, data)
+end
+
+function lazy_map(::Type{T}, f, data::AbstractVector) where {T}
+    return VectorLazyMap{T}(f, data)
+end
+
+const AnyLazyMap{T} = Union{LazyMap{T},VectorLazyMap{T}}
+
+Base.size(it::AnyLazyMap) = size(it.data)
+
+Base.length(it::AnyLazyMap) = length(it.data)
+
+function Base.iterate(it::AnyLazyMap, args...)
     elem_state = iterate(it.data, args...)
     if elem_state === nothing
         return
@@ -46,10 +87,14 @@ function Base.iterate(it::LazyMap, args...)
     end
 end
 
-Base.IteratorSize(it::LazyMap) = Base.IteratorSize(it.data)
+Base.IteratorSize(it::AnyLazyMap) = Base.IteratorSize(it.data)
 
-Base.eltype(::LazyMap{T}) where {T} = T
+Base.eltype(::AnyLazyMap{T}) where {T} = T
+
+function Iterators.reverse(it::AnyLazyMap{T}) where {T}
+    return lazy_map(T, it.f, Iterators.reverse(it.data))
+end
 
-function Iterators.reverse(it::LazyMap{T}) where {T}
-    return LazyMap{T}(it.f, Iterators.reverse(it.data))
+function Base.getindex(it::AnyLazyMap, i)
+    return it.f(getindex(it.data, i))
 end

test/Utilities/lazy_iterators.jl

@@ -37,20 +37,31 @@ end
 test_EmptyVector_Int() = _test_EmptyVector(Int)
 test_EmptyVector_Float64() = _test_EmptyVector(Float64)
 
-function _test_LazyMap(T)
-    v = MOI.Utilities.LazyMap{T}(x -> x^2, [2, 3])
-    @test size(v) == (2,)
-    @test length(v) == 2
-    @test !isempty(v)
-    @test eltype(v) == T
-    c = collect(v)
-    @test c isa Vector{T}
-    @test c == [4, 9]
+function _test_lazy_map(T)
+    for (M{T}, a) in [
+        (MOI.Utilities.LazyMap, Iterators.drop(1:3, 1)),
+        (MOI.Utilities.VectorLazyMap, [2, 3]),
+    ]
+        v = MOI.Utilities.lazy_map(T, x -> x^2, a)
+        @test v isa M
+        @test length(v) == 2
+        @test !isempty(v)
+        @test eltype(v) == T
+        c = collect(v)
+        @test c isa Vector{T}
+        @test c == [4, 9]
+        if a isa AbstractVector
+            @test size(v) == (2,)
+            @test v[1] == 4
+            @test v[2] == 9
+            @test collect(Iterators.reverse(v)) == [9, 4]
+        end
+    end
     return
 end
 
-test_LazyMap_Int() = _test_LazyMap(Int)
-test_LazyMap_Float64() = _test_LazyMap(Float64)
+test_lazy_map_Int() = _test_lazy_map(Int)
+test_lazy_map_Float64() = _test_lazy_map(Float64)
 
 end  # module