Specialize map for eltype conversion (#231)

* specialize map for eltype conversion

* use unwrap and parent_call

Author: jishnub

src/OffsetArrays.jl

@@ -306,6 +306,14 @@ if VERSION < v"1.6"
     end
 end
 
+# Utils to translate a function to the parent while preserving offsets
+unwrap(x) = x, identity
+unwrap(x::OffsetArray) = parent(x), data -> OffsetArray(data, x.offsets)
+function parent_call(f, x)
+    parent, wrap_offset = unwrap(x)
+    wrap_offset(f(parent))
+end
+
 Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
     similar(parent(A), T, dims)
 function Base.similar(A::AbstractArray, ::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T
@@ -385,7 +393,7 @@ parentindex(r::IdOffsetRange, i) = i - r.offset
 end
 
 @propagate_inbounds Base.getindex(A::OffsetArray{<:Any,N}, c::Vararg{Colon,N}) where N =
-    OffsetArray(A.parent[c...], A.offsets)
+    parent_call(x -> getindex(x, c...), A)
 
 # With one Colon we use linear indexing.
 # In this case we may forward the index to the parent, as the information about the axes is lost
@@ -417,7 +425,7 @@ end
 end
 
 Base.in(x, A::OffsetArray) = in(x, parent(A))
-Base.copy(A::OffsetArray) = OffsetArray(copy(A.parent), A.offsets)
+Base.copy(A::OffsetArray) = parent_call(copy, A)
 
 Base.strides(A::OffsetArray) = strides(parent(A))
 Base.elsize(::Type{OffsetArray{T,N,A}}) where {T,N,A} = Base.elsize(A)
@@ -512,6 +520,10 @@ for OR in [:IIUR, :IdOffsetRange]
     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)
+
 # mapreduce is faster with an IdOffsetRange than with an OffsetUnitRange
 # We therefore convert OffsetUnitRanges to IdOffsetRanges with the same values and axes
 function Base.mapreduce(f, op, As::OffsetUnitRange{<:Integer}...; kw...)
@@ -692,7 +704,7 @@ end
 # Adapt allows for automatic conversion of CPU OffsetArrays to GPU OffsetArrays
 ##
 import Adapt
-Adapt.adapt_structure(to, x::OffsetArray) = OffsetArray(Adapt.adapt(to, parent(x)), x.offsets)
+Adapt.adapt_structure(to, O::OffsetArray) = parent_call(x -> Adapt.adapt(to, x), O)
 
 if Base.VERSION >= v"1.4.2"
     include("precompile.jl")

test/runtests.jl

@@ -78,11 +78,6 @@ for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
     end
 end
 
-Base.@propagate_inbounds function Base.getindex(A::Array, r::ZeroBasedUnitRange)
-    B = A[parent(r)]
-    OffsetArrays._maybewrapoffset(B, axes(r))
-end
-
 function same_value(r1, r2)
     length(r1) == length(r2) || return false
     for (v1, v2) in zip(r1, r2)
@@ -782,6 +777,13 @@ end
     @test axes(A) == (IdentityUnitRange(-1:1),)
 end
 
+@testset "unwrap" begin
+    for A in [ones(2, 2), ones(2:3, 2:3), ZeroBasedRange(1:4)]
+        p, f = OffsetArrays.unwrap(A)
+        @test f(map(y -> y^2, p)) == A.^2
+    end
+end
+
 @testset "Traits" begin
     A0 = [1 3; 2 4]
     A = OffsetArray(A0, (-1,2))                   # IndexLinear
@@ -1750,6 +1752,14 @@ end
     @test dest[1,7] == 2
     @test dest[1,8] == 4
     @test dest[1,9] == -2
+
+    @testset "eltype conversion" begin
+        a = OffsetArray(1:2, 1)
+        b = map(BigInt, a)
+        @test eltype(b) == BigInt
+        @test b == a
+        @test b isa OffsetArrays.OffsetRange
+    end
 end
 
 @testset "reductions" begin