map preserves ranges (#232)

* map preserves ranges

* simplify map

* rename maybewrapoffset to indexedby

* add tests for Base.OneTo

* split methods for IdentityUnitRange and IdOffsetRange

* future-proof IIUR

Author: jishnub

src/OffsetArrays.jl

@@ -451,11 +451,11 @@ Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, or::OffsetRange) = Base.c
 
 # If both the arguments are offset, we may unwrap the indices to call (::OffsetArray)[::AbstractRange{Int}]
 @propagate_inbounds function Base.getindex(A::OffsetArray, r::OffsetRange{Int})
-    _maybewrapoffset(A[parent(r)], axes(r))
+    _indexedby(A[parent(r)], axes(r))
 end
 # If the indices are offset, we may unwrap them and pass the parent to getindex
 @propagate_inbounds function Base.getindex(A::AbstractRange, r::OffsetRange{Int})
-    _maybewrapoffset(A[parent(r)], axes(r))
+    _indexedby(A[parent(r)], axes(r))
 end
 
 # An OffsetUnitRange might use the rapid getindex(::Array, ::AbstractUnitRange{Int}) for contiguous indexing
@@ -469,7 +469,7 @@ end
 if VERSION <= v"1.7.0-DEV.1039"
     @propagate_inbounds function Base.getindex(A::Array, r::Union{IdOffsetRange, IIUR})
         B = A[_contiguousindexingtype(r)]
-        _maybewrapoffset(B, axes(r))
+        _indexedby(B, axes(r))
     end
 end
 
@@ -478,20 +478,20 @@ end
     @boundscheck checkbounds(A, r)
     # nD OffsetArrays do not have their linear indices shifted, so we may forward the indices provided to the parent
     @inbounds B = parent(A)[_contiguousindexingtype(r)]
-    _maybewrapoffset(B, axes(r))
+    _indexedby(B, axes(r))
 end
 @inline function Base.getindex(A::OffsetVector, r::AbstractUnitRange{Int})
     @boundscheck checkbounds(A, r)
     # OffsetVectors may have their linear indices shifted, so we subtract the offset from the indices provided
     @inbounds B = parent(A)[_subtractoffset(r, A.offsets[1])]
-    _maybewrapoffset(B, axes(r))
+    _indexedby(B, axes(r))
 end
 
 # This method added mainly to index an OffsetRange with another range
 @inline function Base.getindex(A::OffsetVector, r::AbstractRange{Int})
     @boundscheck checkbounds(A, r)
     @inbounds B = parent(A)[_subtractoffset(r, A.offsets[1])]
-    _maybewrapoffset(B, axes(r))
+    _indexedby(B, axes(r))
 end
 
 # In general we would pass through getindex(A, I...) which calls to_indices(A, I) and finally to_index(I)
@@ -508,21 +508,26 @@ for OR in [:IIUR, :IdOffsetRange]
         @eval @inline function Base.getindex(r::$R, s::$OR)
             @boundscheck checkbounds(r, s)
             @inbounds pr = r[UnitRange(s)]
-            _maybewrapoffset(pr, axes(s,1))
+            _indexedby(pr, axes(s))
         end
     end
 
     # this method is needed for ambiguity resolution
     @eval @inline function Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T
         @boundscheck checkbounds(r, s)
         @inbounds pr = r[UnitRange(s)]
-        _maybewrapoffset(pr, axes(s,1))
+        _indexedby(pr, axes(s))
     end
 end
 
 # eltype conversion
 # This may use specialized map methods for the parent
 Base.map(::Type{T}, O::OffsetArray) where {T} = parent_call(x -> map(T, x), O)
+Base.map(::Type{T}, r::IdOffsetRange) where {T<:Real} = _indexedby(map(T, UnitRange(r)), axes(r))
+if eltype(IIUR) === Int
+    # This is type-piracy, but there is no way to convert an IdentityUnitRange to a non-Int type in Base
+    Base.map(::Type{T}, r::IdentityUnitRange) where {T<:Real} = _indexedby(map(T, UnitRange(r)), axes(r))
+end
 
 # mapreduce is faster with an IdOffsetRange than with an OffsetUnitRange
 # We therefore convert OffsetUnitRanges to IdOffsetRanges with the same values and axes

src/axes.jl

@@ -179,7 +179,7 @@ end
 @inline function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
     @boundscheck checkbounds(r, s)
     @inbounds pr = r.parent[_subtractoffset(s, r.offset)] .+ r.offset
-    _maybewrapoffset(pr, axes(s,1))
+    _indexedby(pr, axes(s))
 end
 # The following method is required to avoid falling back to getindex(::AbstractUnitRange, ::StepRange{<:Integer})
 @inline function Base.getindex(r::IdOffsetRange, s::StepRange{<:Integer})

src/utils.jl

@@ -77,14 +77,16 @@ function _checkindices(N::Integer, indices, label)
     N == length(indices) || throw_argumenterror(N, indices, label)
 end
 
-@inline _maybewrapoffset(r::AbstractVector, ax::Tuple{Any}) = _maybewrapoffset(r, ax[1])
-@inline _maybewrapoffset(r::AbstractUnitRange{<:Integer}, ::Base.OneTo) = no_offset_view(r)
-@inline _maybewrapoffset(r::AbstractVector, ::Base.OneTo) = no_offset_view(r)
-@inline function _maybewrapoffset(r::AbstractUnitRange{<:Integer}, ax::AbstractUnitRange)
-	of = first(ax) - 1
+@inline _indexedby(r::AbstractVector, ax::Tuple{Any}) = _indexedby(r, ax[1])
+@inline _indexedby(r::AbstractUnitRange{<:Integer}, ::Base.OneTo) = no_offset_view(r)
+@inline _indexedby(r::AbstractUnitRange{Bool}, ::Base.OneTo) = no_offset_view(r)
+@inline _indexedby(r::AbstractVector, ::Base.OneTo) = no_offset_view(r)
+@inline function _indexedby(r::AbstractUnitRange{<:Integer}, ax::AbstractUnitRange)
+	of = convert(eltype(r), first(ax) - 1)
 	IdOffsetRange(_subtractoffset(r, of), of)
 end
-@inline _maybewrapoffset(r::AbstractVector, ax::AbstractUnitRange) = OffsetArray(r, ax)
+@inline _indexedby(r::AbstractUnitRange{Bool}, ax::AbstractUnitRange) = OffsetArray(r, ax)
+@inline _indexedby(r::AbstractVector, ax::AbstractUnitRange) = OffsetArray(r, ax)
 
 # These functions are equivalent to the broadcasted operation r .- of
 # However these ensure that the result is an AbstractRange even if a specific

test/runtests.jl

@@ -63,18 +63,18 @@ for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
     for R in [:AbstractRange, :AbstractUnitRange, :StepRange]
         @eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
             @boundscheck checkbounds(A, r)
-            OffsetArrays._maybewrapoffset(A.a[r .+ 1], axes(r))
+            OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
         end
     end
     for R in [:UnitRange, :StepRange, :StepRangeLen, :LinRange]
         @eval @inline function Base.getindex(A::$R, r::$Z)
             @boundscheck checkbounds(A, r)
-            OffsetArrays._maybewrapoffset(A[r.a], axes(r))
+            OffsetArrays._indexedby(A[r.a], axes(r))
         end
     end
     @eval @inline function Base.getindex(A::StepRangeLen{<:Any,<:Base.TwicePrecision,<:Base.TwicePrecision}, r::$Z)
         @boundscheck checkbounds(A, r)
-        OffsetArrays._maybewrapoffset(A[r.a], axes(r))
+        OffsetArrays._indexedby(A[r.a], axes(r))
     end
 end
 
@@ -1759,6 +1759,28 @@ end
         @test eltype(b) == BigInt
         @test b == a
         @test b isa OffsetArrays.OffsetRange
+
+        for ri in Any[2:3, Base.OneTo(2)]
+            for r in [IdentityUnitRange(ri), IdOffsetRange(ri), IdOffsetRange(ri, 1)]
+                for T in [Int8, Int16, Int32, Int64, Int128, BigInt, Float32, Float64, BigFloat]
+                    r2 = map(T, r)
+                    @test eltype(r2) == T
+                    @test axes(r2) == axes(r)
+                    @test all(((x,y),) -> isequal(x,y), zip(r, r2))
+                end
+            end
+        end
+
+        @testset "Bool" begin
+            for ri in Any[0:0, 0:1, 1:0, 1:1, Base.OneTo(0), Base.OneTo(1)]
+                for r = Any[IdentityUnitRange(ri), IdOffsetRange(ri), IdOffsetRange(ri .- 1, 1)]
+                    r2 = map(Bool, r)
+                    @test eltype(r2) == Bool
+                    @test axes(r2) == axes(r)
+                    @test all(((x,y),) -> isequal(x,y), zip(r, r2))
+                end
+            end
+        end
     end
 end