Specialize copy and getindex with Colons (#207)

* forward colon

* linear indexing with a single colon

* remove comment

Author: jishnub

Project.toml

@@ -7,6 +7,12 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 
 [compat]
 Adapt = "2, 3"
+Aqua = "0.5"
+CatIndices = "0.2"
+Documenter = "0.26"
+EllipsisNotation = "1"
+FillArrays = "0.11"
+StaticArrays = "1"
 julia = "0.7, 1"
 
 [extras]
@@ -15,9 +21,10 @@ CatIndices = "aafaddc9-749c-510e-ac4f-586e18779b91"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra", "EllipsisNotation", "StaticArrays"]
+test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra", "EllipsisNotation", "StaticArrays", "FillArrays"]

src/OffsetArrays.jl

@@ -320,12 +320,23 @@ Base.falses(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
 # and one obtains the result below.
 parentindex(r::IdOffsetRange, i) = i - r.offset
 
-@inline function Base.getindex(A::OffsetArray{T,N}, I::Vararg{Int,N}) where {T,N}
+@propagate_inbounds Base.getindex(A::OffsetArray{<:Any,0})  = A.parent[]
+
+@inline function Base.getindex(A::OffsetArray{<:Any,N}, I::Vararg{Int,N}) where N
     @boundscheck checkbounds(A, I...)
     J = map(parentindex, axes(A), I)
     @inbounds parent(A)[J...]
 end
 
+@propagate_inbounds Base.getindex(A::OffsetArray{<:Any,N}, c::Vararg{Colon,N}) where N =
+    OffsetArray(A.parent[c...], A.offsets)
+
+# 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
+# The exception to this is with OffsetVectors where the axis information is preserved,
+# but that case is handled by getindex(::OffsetArray{<:Any,N}, ::Vararg{Colon,N})
+@propagate_inbounds Base.getindex(A::OffsetArray, c::Colon) = A.parent[:]
+
 @inline function Base.getindex(A::OffsetVector, i::Int)
     @boundscheck checkbounds(A, i)
     @inbounds parent(A)[parentindex(Base.axes1(A), i)]
@@ -350,6 +361,7 @@ end
 end
 
 Base.in(x, A::OffsetArray) = in(x, parent(A))
+Base.copy(A::OffsetArray) = OffsetArray(copy(A.parent), A.offsets)
 
 Base.strides(A::OffsetArray) = strides(parent(A))
 Base.elsize(::Type{OffsetArray{T,N,A}}) where {T,N,A} = Base.elsize(A)
@@ -394,9 +406,6 @@ for OR in [:IIUR, :IdOffsetRange]
     end
 end
 
-# This is technically breaking, so it might be incorporated in the next major release
-# Base.getindex(a::OffsetRange, ::Colon) = OffsetArray(a.parent[:], a.offsets)
-
 # 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...)

test/runtests.jl

@@ -8,6 +8,7 @@ using CatIndices: BidirectionalVector
 using EllipsisNotation
 using Adapt
 using StaticArrays
+using FillArrays
 
 DocMeta.setdocmeta!(OffsetArrays, :DocTestSetup, :(using OffsetArrays); recursive=true)
 
@@ -887,9 +888,37 @@ end
     @test A[1, [4,3]] == S[1, [4,3]] == [4,2]
     @test A[:, :] == S[:, :] == A
 
+    # Indexing a nD OffsetArray with n colons preserves the type
     r1 = OffsetArray(IdentityUnitRange(100:1000), 3)
-    r2 = r1[:]
-    @test r2 == r1
+    @test r1[:] === r1
+
+    # In general with more colons than dimensions,
+    # the type might not be preserved but the values and the leading axes should be
+    r2 = r1[:,:]
+    @test axes(r2, 1) == axes(r1, 1)
+    @test same_value(r1, r2)
+
+    s = @SVector[i for i in 1:3]
+    so = OffsetArray(s, 3)
+    @test so[:] === so
+
+    a = Ones(3, 2, 1)
+    ao = OffsetArray(a, axes(a))
+    @test ao[:,:,:] === ao
+    @test same_value(ao[:], ao)
+    @test same_value(ao[:,:], ao)
+
+    # Indexing an nD OffsetArray with one Colon preserves only the values.
+    # This uses linear indexing
+    a = ones(2:3, 2:3)
+    b = a[:]
+    @test same_value(a, b)
+
+    vals = (1,2,3,4,5,6,7,8)
+    s = SArray{Tuple{2,2,2},Int,3,8}(vals)
+    so = OffsetArray(s, axes(s));
+    so2 = so[:]
+    @test same_value(so2, s)
 
     for r1 in Any[
         # AbstractArrays