Fix getindex for LazyMap with Vector

JuliaAlgebra · Jul 20, 2023 · abfa2ce · abfa2ce
1 parent 0268a1c
commit abfa2ce
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/lazy_iterators.jl b/src/lazy_iterators.jl
@@ -32,7 +32,7 @@ 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`.
 """
-struct LazyMap{T,VT,F} <: AbstractVector{T}
+struct LazyMap{T,VT<:AbstractVector,F} <: AbstractVector{T}
     f::F
     data::VT
 end
@@ -58,7 +58,10 @@ Base.IteratorSize(it::LazyMap) = Base.IteratorSize(it.data)
 
 Base.eltype(::LazyMap{T}) where {T} = T
 
-Base.getindex(it::LazyMap, i) = it.f(getindex(it.data, i))
+Base.getindex(it::LazyMap, i::Integer) = it.f(getindex(it.data, i))
+function Base.getindex(it::LazyMap{T}, I::AbstractVector) where {T}
+    return LazyMap{T}(it.f, getindex(it.data, I))
+end
 
 Base.eachindex(it::LazyMap) = Base.eachindex(it.data)
 Base.lastindex(it::LazyMap) = Base.lastindex(it.data)

diff --git a/test/polynomial.jl b/test/polynomial.jl
@@ -107,12 +107,14 @@ const MP = MultivariatePolynomials
     # Doc examples
     @test collect(coefficients(4x^2 * y + x * y + 2x)) == [2, 1, 4]
     @test collect(coefficients(4x^2 * y + x * y + 2x + 3, [x, 1, x * y, y])) ==
-          [2, 3, 1, 0]
+    @test monomials(4x^2 * y + x * y + 2x + 3)[1:1] == [constant_monomial(x*y)]
 
     for p in [polynomial([4, 9], [x, x * x]), polynomial([9, 4], [x * x, x])]
         @test collect(coefficients(p)) == [4, 9]
         @test monomials(p)[1] == x
         @test monomials(p)[2] == x^2
+        @test monomials(p)[1:2][1] == x
+        @test monomials(p)[1:2][2] == x^2
         @test p == dot([4, 9], [x, x * x])
     end