Make OffsetArray's axes offset themselves (#46)

* Make OffsetArray's axes offset themselves

This is the companion PR to JuliaLang/julia#27038.

* Display slices as simple ranges in show -- we allow constructing OffsetArrays in this manner

* Add compatibility for 0.6

Author: mbauman

src/OffsetArrays.jl

@@ -8,6 +8,9 @@ using Compat: axes, CartesianIndices
 
 export OffsetArray, OffsetVector, @unsafe
 
+## Major change in 0.7: OffsetArray now uses Offset axes
+const AxisType = VERSION < v"0.7.0-DEV.5242" ? identity : Base.Slice
+
 # TODO: just use .+
 # See https://github.com/JuliaLang/julia/pull/22932#issuecomment-330711997
 if VERSION < v"0.7.0-DEV.1759"
@@ -74,8 +77,8 @@ end
 OffsetArray(A::AbstractArray{T,0}, inds::Tuple{}) where {T} = OffsetArray{T,0,typeof(A)}(A, ())
 OffsetArray(A::AbstractArray{T,N}, inds::Tuple{}) where {T,N} = error("this should never be called")
 function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,AbstractUnitRange}) where {T,N}
-    lA = map(length, axes(A))
-    lI = map(length, inds)
+    lA = map(indexlength, axes(A))
+    lI = map(indexlength, inds)
     lA == lI || throw(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
     OffsetArray(A, map(indexoffset, inds))
 end
@@ -98,38 +101,48 @@ Base.eachindex(::IndexLinear, A::OffsetVector)   = axes(A, 1)
 # performance-critical and relies on axes, these are usually worth
 # optimizing thoroughly.
 @inline Compat.axes(A::OffsetArray, d) =
-    1 <= d <= length(A.offsets) ? plus(axes(parent(A))[d], A.offsets[d]) : (1:1)
+    1 <= d <= length(A.offsets) ? AxisType(plus(axes(parent(A))[d], A.offsets[d])) : (1:1)
 @inline Compat.axes(A::OffsetArray) =
     _axes(axes(parent(A)), A.offsets)  # would rather use ntuple, but see #15276
 @inline _axes(inds, offsets) =
-    (plus(inds[1], offsets[1]), _axes(tail(inds), tail(offsets))...)
+    (AxisType(plus(inds[1], offsets[1])), _axes(tail(inds), tail(offsets))...)
 _axes(::Tuple{}, ::Tuple{}) = ()
 Base.indices1(A::OffsetArray{T,0}) where {T} = 1:1  # we only need to specialize this one
 
+
+const OffsetAxis = Union{Integer, UnitRange, Base.Slice{<:UnitRange}, Base.OneTo}
 function Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T
     B = similar(parent(A), T, dims)
 end
-function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{UnitRange,Vararg{UnitRange}}) where T
-    B = similar(A, T, map(length, inds))
+function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where T
+    B = similar(A, T, map(indexlength, inds))
     OffsetArray(B, map(indexoffset, inds))
 end
 
-Base.similar(::Type{T}, shape::Tuple{UnitRange,Vararg{UnitRange}}) where {T<:OffsetArray} =
-    OffsetArray(T(map(length, shape)), map(indexoffset, shape))
-Base.similar(::Type{T}, shape::Tuple{UnitRange,Vararg{UnitRange}}) where {T<:Array} =
-    OffsetArray(T(undef, map(length, shape)), map(indexoffset, shape))
-Base.similar(::Type{T}, shape::Tuple{UnitRange,Vararg{UnitRange}}) where {T<:BitArray} =
-    OffsetArray(T(undef, map(length, shape)), map(indexoffset, shape))
+Base.similar(::Type{T}, shape::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T<:AbstractArray} =
+    OffsetArray(T(undef, map(indexlength, shape)), map(indexoffset, shape))
 
+if VERSION < v"0.7.0-DEV.5242"
+# Reshape's methods in Base changed, so using the new definitions leads to ambiguities
 Base.reshape(A::AbstractArray, inds::Tuple{UnitRange,Vararg{UnitRange}}) =
     OffsetArray(reshape(A, map(length, inds)), map(indexoffset, inds))
-
 Base.reshape(A::OffsetArray, inds::Tuple{UnitRange,Vararg{UnitRange}}) =
     OffsetArray(reshape(parent(A), map(length, inds)), map(indexoffset, inds))
-
 function Base.reshape(A::OffsetArray, inds::Tuple{UnitRange,Vararg{Union{UnitRange,Int,Base.OneTo}}})
     throw(ArgumentError("reshape must supply UnitRange axes, got $(typeof(inds)).\n       Note that reshape(A, Val{N}) is not supported for OffsetArrays."))
 end
+else
+Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
+    OffsetArray(reshape(A, map(indexlength, inds)), map(indexoffset, inds))
+
+# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
+# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
+Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
+    OffsetArray(reshape(parent(A), map(indexlength, inds)), map(indexoffset, inds))
+# And for non-offset axes, we can just return a reshape of the parent directly
+Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
+Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
+end
 
 if VERSION < v"0.7.0-DEV.4873"
     # Julia PR #26733 removed similar(f, ...) in favor of just using method extension directly
@@ -201,7 +214,7 @@ offset(offsets::Tuple{Vararg{Int}}, inds::Tuple{}) = error("inds cannot be short
 
 indexoffset(r::AbstractRange) = first(r) - 1
 indexoffset(i::Integer) = 0
-indexlength(r::AbstractRange) = length(r)
+indexlength(r::AbstractRange) = Base._length(r)
 indexlength(i::Integer) = i
 
 macro unsafe(ex)
@@ -290,8 +303,10 @@ if VERSION >= v"0.7.0-DEV.1790"
         toplevel && print(io, " with eltype ", eltype(a))
     end
     printindices(io::IO, ind1, inds...) =
-        (print(io, ind1, ", "); printindices(io, inds...))
-    printindices(io::IO, ind1) = print(io, ind1)
+        (print(io, _unslice(ind1), ", "); printindices(io, inds...))
+    printindices(io::IO, ind1) = print(io, _unslice(ind1))
+    _unslice(x) = x
+    _unslice(x::Base.Slice) = x.indices
 end
 
 end # module

test/runtests.jl

@@ -94,21 +94,22 @@ Ac[0,3,1] = 11
 @test_throws BoundsError S[CartesianIndex(1,1),0]
 @test_throws BoundsError S[CartesianIndex(1,1),2]
 @test eachindex(A) == 1:4
-@test eachindex(S) == CartesianIndices((0:1,3:4))
+@test eachindex(S) == CartesianIndices(Base.Slice.((0:1,3:4)))
 
 # view
+const AxisType = VERSION < v"0.7.0-DEV.5242" ? identity : Base.Slice
 S = view(A, :, 3)
 @test S == OffsetArray([1,2], (A.offsets[1],))
 @test S[0] == 1
 @test S[1] == 2
 @test_throws BoundsError S[2]
-@test axes(S) === (0:1,)
+@test axes(S) === (AxisType(0:1),)
 S = view(A, 0, :)
 @test S == OffsetArray([1,3], (A.offsets[2],))
 @test S[3] == 1
 @test S[4] == 3
 @test_throws BoundsError S[1]
-@test axes(S) === (3:4,)
+@test axes(S) === (AxisType(3:4),)
 S = view(A, 0:0, 4)
 @test S == [3]
 @test S[1] == 3
@@ -127,7 +128,7 @@ S = view(A, :, :)
 @test S[0,4] == S[3] == 3
 @test S[1,4] == S[4] == 4
 @test_throws BoundsError S[1,1]
-@test axes(S) === (0:1, 3:4)
+@test axes(S) === AxisType.((0:1, 3:4))
 
 # iteration
 let a
@@ -168,43 +169,53 @@ B = similar(A, (3,4))
 @test axes(B) === (Base.OneTo(3), Base.OneTo(4))
 B = similar(A, (-3:3,1:4))
 @test isa(B, OffsetArray{Int,2})
-@test axes(B) === (-3:3, 1:4)
+@test axes(B) === AxisType.((-3:3, 1:4))
 B = similar(parent(A), (-3:3,1:4))
 @test isa(B, OffsetArray{Int,2})
-@test axes(B) === (-3:3, 1:4)
+@test axes(B) === AxisType.((-3:3, 1:4))
 
 # Reshape
 B = reshape(A0, -10:-9, 9:10)
 @test isa(B, OffsetArray{Int,2})
 @test parent(B) === A0
-@test axes(B) == (-10:-9, 9:10)
+@test axes(B) == AxisType.((-10:-9, 9:10))
 B = reshape(A, -10:-9, 9:10)
 @test isa(B, OffsetArray{Int,2})
-@test parent(B) === A0
-@test axes(B) == (-10:-9, 9:10)
+@test pointer(parent(B)) === pointer(A0)
+@test axes(B) == AxisType.((-10:-9, 9:10))
 b = reshape(A, -7:-4)
-@test axes(b) == (-7:-4,)
+@test axes(b) == (AxisType(-7:-4),)
 @test isa(parent(b), Vector{Int})
+@test pointer(parent(b)) === pointer(parent(A))
 @test parent(b) == A0[:]
 a = OffsetArray(rand(3,3,3), -1:1, 0:2, 3:5)
-@test_throws ArgumentError reshape(a, Val(2))
-@test_throws ArgumentError reshape(a, Val(4))
+if VERSION >= v"0.7.0-DEV.5242"
+# Offset axes are required for reshape(::OffsetArray, ::Val) support
+b = reshape(a, Val(2))
+@test isa(b, OffsetArray{Float64,2})
+@test pointer(parent(b)) === pointer(parent(a))
+@test axes(b) == AxisType.((-1:1, 1:9))
+b = reshape(a, Val(4))
+@test isa(b, OffsetArray{Float64,4})
+@test pointer(parent(b)) === pointer(parent(a))
+@test axes(b) == (axes(a)..., AxisType(1:1))
+end
 
 # Indexing with OffsetArray axes
 i1 = OffsetArray([2,1], (-5,))
 i1 = OffsetArray([2,1], -5)
 b = A0[i1, 1]
-@test axes(b) === (-4:-3,)
+@test axes(b) === (AxisType(-4:-3),)
 @test b[-4] == 2
 @test b[-3] == 1
 b = A0[1,i1]
-@test axes(b) === (-4:-3,)
+@test axes(b) === (AxisType(-4:-3),)
 @test b[-4] == 3
 @test b[-3] == 1
 v = view(A0, i1, 1)
-@test axes(v) === (-4:-3,)
+@test axes(v) === (AxisType(-4:-3),)
 v = view(A0, 1:1, i1)
-@test axes(v) === (Base.OneTo(1), -4:-3)
+@test axes(v) === (Base.OneTo(1), AxisType(-4:-3))
 
 # logical indexing
 @test A[A .> 2] == [3,4]
@@ -314,8 +325,10 @@ seek(io, 0)
 amin, amax = extrema(parent(A))
 @test clamp.(A, (amax+amin)/2, amax) == OffsetArray(clamp.(parent(A), (amax+amin)/2, amax), axes(A))
 
-@test unique(A, 1) == parent(A)
-@test unique(A, 2) == parent(A)
+if VERSION >= v"0.7.0-DEV.5242"
+@test unique(A, 1) == OffsetArray(parent(A), 0, first(axes(A, 2)) - 1)
+@test unique(A, 2) == OffsetArray(parent(A), first(axes(A, 1)) - 1, 0)
+end
 v = OffsetArray(rand(8), (-2,))
 @test sort(v) == OffsetArray(sort(parent(v)), v.offsets)
 @test sortrows(A) == OffsetArray(sortrows(parent(A)), A.offsets)