Support atsign-unsafe indexing with trailing 1s

Author: timholy

src/OffsetArrays.jl

@@ -121,6 +121,11 @@ offset{N}(offsets::NTuple{N,Int}, inds::NTuple{N,Int}) = _offset((), offsets, in
 _offset(out, ::Tuple{}, ::Tuple{}) = out
 @inline _offset(out, offsets, inds) = _offset((out..., inds[1]-offsets[1]), Base.tail(offsets), Base.tail(inds))
 
+# Support trailing 1s
+@inline offset(offsets::Tuple{Vararg{Int}}, inds::Tuple{Vararg{Int}}) = (offset(offsets, Base.front(inds))..., inds[end])
+offset(offsets::Tuple{}, inds::Tuple{}) = ()
+offset(offsets::Tuple{Vararg{Int}}, inds::Tuple{}) = error("inds cannot be shorter than offsets")
+
 indexoffset(r::Range) = first(r) - 1
 indexoffset(i::Integer) = 0
 
@@ -152,6 +157,34 @@ end
 @inline unsafe_getindex(a::AbstractArray, I...) = (@inbounds ret = a[I...]; ret)
 @inline unsafe_setindex!(a::AbstractArray, val, I...) = (@inbounds a[I...] = val; val)
 
+# Linear indexing
+@inline unsafe_getindex(a::OffsetArray, i::Int) = _unsafe_getindex(Base.linearindexing(a), a, i)
+@inline unsafe_setindex!(a::OffsetArray, val, i::Int) = _unsafe_setindex!(Base.linearindexing(a), a, val, i)
+for T in (LinearFast, LinearSlow)  # ambiguity-resolution requires specificity for both
+    @eval begin
+        @inline function _unsafe_getindex(::$T, a::OffsetVector, i::Int)
+            @inbounds ret = parent(a)[offset(a.offsets, (i,))[1]]
+            ret
+        end
+        @inline function _unsafe_setindex!(::$T, a::OffsetVector, val, i::Int)
+            @inbounds parent(a)[offset(a.offsets, (i,))[1]] = val
+            val
+        end
+    end
+end
+@inline function _unsafe_getindex(::LinearFast, a::OffsetArray, i::Int)
+    @inbounds ret = parent(a)[i]
+    ret
+end
+@inline _unsafe_getindex(::LinearSlow, a::OffsetArray, i::Int) =
+    unsafe_getindex(a, ind2sub(indices(a), i)...)
+@inline function _unsafe_setindex!(::LinearFast, a::OffsetArray, val, i::Int)
+    @inbounds parent(a)[i] = val
+    val
+end
+@inline _unsafe_setindex!(::LinearSlow, a::OffsetArray, val, i::Int) =
+    unsafe_setindex!(a, val, ind2sub(indices(a), i)...)
+
 @inline unsafe_getindex(a::OffsetArray, I::Int...) = unsafe_getindex(parent(a), offset(a.offsets, I)...)
 @inline unsafe_setindex!(a::OffsetArray, val, I::Int...) = unsafe_setindex!(parent(a), val, offset(a.offsets, I)...)
 @inline unsafe_getindex(a::OffsetArray, I...) = unsafe_getindex(a, Base.IteratorsMD.flatten(I)...)

test/runtests.jl

@@ -39,12 +39,25 @@ S = OffsetArray(view(A0, 1:2, 1:2), (-1,2))   # LinearSlow
 @test_throws DimensionMismatch OffsetArray(A0, 0:1, 2:4)
 
 # Scalar indexing
-@test A[0,3] == A[1] == S[0,3] == S[1] == 1
-@test A[1,3] == A[2] == S[1,3] == S[2] == 2
-@test A[0,4] == A[3] == S[0,4] == S[3] == 3
-@test A[1,4] == A[4] == S[1,4] == S[4] == 4
+@test A[0,3] == A[0,3,1] == A[1] == S[0,3] == S[0,3,1] == S[1] == 1
+@test A[1,3] == A[1,3,1] == A[2] == S[1,3] == S[1,3,1] == S[2] == 2
+@test A[0,4] == A[0,4,1] == A[3] == S[0,4] == S[0,4,1] == S[3] == 3
+@test A[1,4] == A[1,4,1] == A[4] == S[1,4] == S[1,4,1] == S[4] == 4
+@test @unsafe(A[0,3]) == @unsafe(A[0,3,1]) == @unsafe(A[1]) == @unsafe(S[0,3]) == @unsafe(S[0,3,1]) == @unsafe(S[1]) == 1
+@test @unsafe(A[1,3]) == @unsafe(A[1,3,1]) == @unsafe(A[2]) == @unsafe(S[1,3]) == @unsafe(S[1,3,1]) == @unsafe(S[2]) == 2
+@test @unsafe(A[0,4]) == @unsafe(A[0,4,1]) == @unsafe(A[3]) == @unsafe(S[0,4]) == @unsafe(S[0,4,1]) == @unsafe(S[3]) == 3
+@test @unsafe(A[1,4]) == @unsafe(A[1,4,1]) == @unsafe(A[4]) == @unsafe(S[1,4]) == @unsafe(S[1,4,1]) == @unsafe(S[4]) == 4
 @test_throws BoundsError A[1,1]
 @test_throws BoundsError S[1,1]
+@test_throws BoundsError A[0,3,2]
+@test_throws BoundsError A[0,3,0]
+Ac = copy(A)
+Ac[0,3] = 10
+@test Ac[0,3] == 10
+Ac[0,3,1] = 11
+@test Ac[0,3] == 11
+@unsafe Ac[0,3,1] = 12
+@test Ac[0,3] == 12
 
 # Vector indexing
 @test A[:, 3] == S[:, 3] == OffsetArray([1,2], (A.offsets[1],))
@@ -63,8 +76,15 @@ S = OffsetArray(view(A0, 1:2, 1:2), (-1,2))   # LinearSlow
 
 # CartesianIndexing
 @test A[CartesianIndex((0,3))] == S[CartesianIndex((0,3))] == 1
+@test A[CartesianIndex((0,3)),1] == S[CartesianIndex((0,3)),1] == 1
+@test @unsafe(A[CartesianIndex((0,3))]) == @unsafe(S[CartesianIndex((0,3))]) == 1
+@test @unsafe(A[CartesianIndex((0,3)),1]) == @unsafe(S[CartesianIndex((0,3)),1]) == 1
 @test_throws BoundsError A[CartesianIndex(1,1)]
+@test_throws BoundsError A[CartesianIndex(1,1),0]
+@test_throws BoundsError A[CartesianIndex(1,1),2]
 @test_throws BoundsError S[CartesianIndex(1,1)]
+@test_throws BoundsError S[CartesianIndex(1,1),0]
+@test_throws BoundsError S[CartesianIndex(1,1),2]
 @test eachindex(A) == 1:4
 @test eachindex(S) == CartesianRange((0:1,3:4))