OffsetArray support for cat/vcat/hcat

Author: johnnychen94

base/abstractarray.jl

@@ -1442,13 +1442,13 @@ function _typed_vcat(::Type{T}, V::AbstractVecOrTuple{AbstractVector}) where T
     for Vk in V
         n += Int(length(Vk))::Int
     end
-    a = similar(V[1], T, n)
-    pos = 1
-    for k=1:Int(length(V))::Int
+    a = similar(first(V), T, n)
+    pos = first(axes(a, 1))
+    for k = eachindex(V)
         Vk = V[k]
-        p1 = pos + Int(length(Vk))::Int - 1
-        a[pos:p1] = Vk
-        pos = p1+1
+        n = length(Vk)
+        copyto!(a, pos, Vk, first(axes(Vk, 1)), n)
+        pos += n
     end
     a
 end
@@ -1459,11 +1459,10 @@ hcat(A::AbstractVecOrMat...) = typed_hcat(promote_eltype(A...), A...)
 hcat(A::AbstractVecOrMat{T}...) where {T} = typed_hcat(T, A...)
 
 function _typed_hcat(::Type{T}, A::AbstractVecOrTuple{AbstractVecOrMat}) where T
-    nargs = length(A)
-    nrows = size(A[1], 1)
+    nrows = size(first(A), 1)
     ncols = 0
     dense = true
-    for j = 1:nargs
+    for j = eachindex(A)
         Aj = A[j]
         if size(Aj, 1) != nrows
             throw(ArgumentError("number of rows of each array must match (got $(map(x->size(x,1), A)))"))
@@ -1472,17 +1471,17 @@ function _typed_hcat(::Type{T}, A::AbstractVecOrTuple{AbstractVecOrMat}) where T
         nd = ndims(Aj)
         ncols += (nd==2 ? size(Aj,2) : 1)
     end
-    B = similar(A[1], T, nrows, ncols)
-    pos = 1
+    B = similar(first(A), T, nrows, ncols)
+    pos = first(axes(B, 1))
     if dense
-        for k=1:nargs
+        for k=eachindex(A)
             Ak = A[k]
             n = length(Ak)
             copyto!(B, pos, Ak, 1, n)
             pos += n
         end
     else
-        for k=1:nargs
+        for k=eachindex(A)
             Ak = A[k]
             p1 = pos+(isa(Ak,AbstractMatrix) ? size(Ak, 2) : 1)-1
             B[:, pos:p1] = Ak
@@ -1496,17 +1495,16 @@ vcat(A::AbstractVecOrMat...) = typed_vcat(promote_eltype(A...), A...)
 vcat(A::AbstractVecOrMat{T}...) where {T} = typed_vcat(T, A...)
 
 function _typed_vcat(::Type{T}, A::AbstractVecOrTuple{AbstractVecOrMat}) where T
-    nargs = length(A)
     nrows = sum(a->size(a, 1), A)::Int
-    ncols = size(A[1], 2)
-    for j = 2:nargs
+    ncols = size(first(A), 2)
+    for j = first(axes(A))[2:end]
         if size(A[j], 2) != ncols
             throw(ArgumentError("number of columns of each array must match (got $(map(x->size(x,2), A)))"))
         end
     end
-    B = similar(A[1], T, nrows, ncols)
-    pos = 1
-    for k=1:nargs
+    B = similar(first(A), T, nrows, ncols)
+    pos = first(axes(B, 1))
+    for k=eachindex(A)
         Ak = A[k]
         p1 = pos+size(Ak,1)::Int-1
         B[pos:p1, :] = Ak
@@ -1589,7 +1587,7 @@ _cat(dims, X...) = cat_t(promote_eltypeof(X...), X...; dims=dims)
 @inline function _cat_t(dims, ::Type{T}, X...) where {T}
     catdims = dims2cat(dims)
     shape = cat_shape(catdims, map(cat_size, X)::Tuple{Vararg{Union{Int,Dims}}})::Dims
-    A = cat_similar(X[1], T, shape)
+    A = cat_similar(first(X), T, shape)
     if count(!iszero, catdims)::Int > 1
         fill!(A, zero(T))
     end
@@ -1604,7 +1602,7 @@ function __cat(A, shape::NTuple{M,Int}, catdims, X...) where M
     for x in X
         for i = 1:N
             if concat[i]
-                inds[i] = offsets[i] .+ cat_indices(x, i)
+                inds[i] = offsets[i] .+ parent(cat_indices(x, i))
                 offsets[i] += cat_size(x, i)
             else
                 inds[i] = 1:shape[i]

test/abstractarray.jl

@@ -616,6 +616,52 @@ function test_cat(::Type{TestAbstractArray})
     @test cat([1], [1], dims=[1, 2]) == I(2)
 end
 
+module TestOffsetArraysCats
+    using Test
+    isdefined(Main, :OffsetArrays) || @eval Main include(joinpath(@__DIR__, "testhelpers", "OffsetArrays.jl"))
+    using .Main.OffsetArrays
+
+    # `cat`s on OffsetArrays ignore their offsets and treat them as normal list
+
+    # 1d
+    v1 = collect(1:4)
+    v2 = collect(5:8)
+    ov1 = OffsetArray(v1, -1)
+    ov2 = OffsetArray(v2, 1)
+    @test hcat(ov1, v1, ov2, v2) == hcat(v1, v1, v2, v2)
+    @test vcat(ov1, v1, ov2, v2) == vcat(v1, v1, v2, v2)
+    @test hvcat((2, 2), ov1, v2, v1, ov2) == hvcat((2, 2), v1, v2, v1, v2)
+    # 37628
+    @test reduce(hcat, (v1, v2)) == hcat(v1, v2)
+    @test reduce(vcat, (v1, v2)) == vcat(v1, v2)
+    @test reduce(hcat, OffsetVector([1:2, 1:2],10)) == [1 1;2 2]
+
+    # 2d
+    a1 = reshape(collect(1:6), 2, 3)
+    a2 = reshape(collect(7:12), 2, 3)
+    oa1 = OffsetArray(a1, -1, -1)
+    oa2 = OffsetArray(a2, 1, 1)
+    @test hcat(oa1, a1, oa2, a2) == hcat(a1, a1, a2, a2)
+    @test vcat(oa1, a1, oa2, a2) == vcat(a1, a1, a2, a2)
+    @test hvcat((2, 2), oa1, a2, a1, oa2) == hvcat((2, 2), a1, a2, a1, a2)
+
+    # 3d
+    a1 = reshape(collect(1:12), 2, 3, 2)
+    a2 = reshape(collect(13:24), 2, 3, 2)
+    oa1 = OffsetArray(a1, -1, -1, -1)
+    oa2 = OffsetArray(a2, 1, 1, 1)
+    @test hcat(oa1, a1, oa2, a2) == hcat(a1, a1, a2, a2)
+    @test vcat(oa1, a1, oa2, a2) == vcat(a1, a1, a2, a2)
+    @test hvcat((2, 2), oa1, a2, a1, oa2) == hvcat((2, 2), a1, a2, a1, a2)
+    # https://github.com/JuliaArrays/OffsetArrays.jl/issues/63
+    form=OffsetArray(reshape(zeros(Int8,0),0,0,2),0:-1,0:-1,0:1)
+    exp=OffsetArray(reshape(zeros(Int8,0),0,16,2),0:-1,0:15,0:1)
+    @test size(hcat(form,exp)) == (0, 16, 2)
+    # 37493
+    @test hcat(zeros(2, 1:1, 2), zeros(2, 2:3, 2)) == zeros(2, 3, 2)
+    @test vcat(zeros(1:1, 2, 2), zeros(2:3, 2, 2)) == zeros(3, 2, 2)
+end
+
 function test_ind2sub(::Type{TestAbstractArray})
     n = rand(2:5)
     dims = tuple(rand(1:5, n)...)