diff --git a/src/Bridges/Constraint/map.jl b/src/Bridges/Constraint/map.jl index b962c75ddb..95a6d63e16 100644 --- a/src/Bridges/Constraint/map.jl +++ b/src/Bridges/Constraint/map.jl @@ -181,7 +181,7 @@ 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 @@ -189,7 +189,8 @@ 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 -> diff --git a/src/Bridges/Variable/map.jl b/src/Bridges/Variable/map.jl index a02588d6f9..7e664cc5b9 100644 --- a/src/Bridges/Variable/map.jl +++ b/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), ), diff --git a/src/Utilities/lazy_iterators.jl b/src/Utilities/lazy_iterators.jl index 2d434e1788..1075fefa04 100644 --- a/src/Utilities/lazy_iterators.jl +++ b/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 diff --git a/test/Utilities/lazy_iterators.jl b/test/Utilities/lazy_iterators.jl index 92cde24d3d..fc13871447 100644 --- a/test/Utilities/lazy_iterators.jl +++ b/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