Add deepgetindex, deepsetindex! and deepview

oschulz · oschulz · commit 714bfd8bcfce · 2018-11-28T00:17:53.000+01:00
diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl
@@ -204,6 +204,58 @@ end
 
 
 
+Base.@pure _result_is_nested(idxs_outer::Tuple, idxs_inner::Tuple) =
+    Val{!(Base.index_dimsum(idxs_outer...) isa Tuple{}) && !(Base.index_dimsum(idxs_inner...) isa Tuple{})}()
+
+Base.@pure ndims_after_getindex(idxs::Tuple) = Val{length(Base.index_dimsum(idxs...))}()
+
+
+Base.@propagate_inbounds function deepgetindex(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L}
+    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
+    nested = _result_is_nested(idxs_outer, idxs_inner)
+    _deepgetindex_impl_aosa(A, idxs_outer, idxs_inner, nested)
+end
+
+Base.@propagate_inbounds _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) =
+    getindex(A.data, idxs_inner..., idxs_outer...)
+
+Base.@propagate_inbounds function _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true})
+    new_data = getindex(A.data, idxs_inner..., idxs_outer...)
+    nestedview(new_data, ndims_after_getindex(idxs_inner))
+end
+
+
+Base.@propagate_inbounds function deepsetindex!(A::ArrayOfSimilarArrays{T,M,N,L}, x, idxs::Vararg{Any,L}) where {T,M,N,L}
+    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
+    _deepsetindex_impl_aosa!(A, x, idxs_outer, idxs_inner, _result_is_nested(idxs_outer, idxs_inner))
+    A
+end
+
+Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) =
+    setindex!(A.data, x, idxs_inner..., idxs_outer...)
+
+Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) =
+    _deepsetindex_impl!(A, x, idxs_outer, idxs_inner)
+
+# TODO: Specialized method _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) =
+
+
+Base.@propagate_inbounds function deepview(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L}
+    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
+    nested = _result_is_nested(idxs_outer, idxs_inner)
+    _deepview_impl_aosa(A, idxs_outer, idxs_inner, nested)
+end
+
+Base.@propagate_inbounds _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) =
+    view(A.data, idxs_inner..., idxs_outer...)
+
+Base.@propagate_inbounds function _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true})
+    new_data = view(A.data, idxs_inner..., idxs_outer...)
+    nestedview(new_data, ndims_after_getindex(idxs_inner))
+end
+
+
+
 const VectorOfSimilarArrays{
     T, M, L,
     P<:AbstractArray{T,L}
diff --git a/src/functions.jl b/src/functions.jl
@@ -44,6 +44,8 @@ export flatview
 
 @inline flatview(A::AbstractArray) = A
 
+# TODO: Implement flatview on generic nested arrays via new `FlatView`, using
+# deepgetindex to implement getindex, etc.
 @inline flatview(A::AbstractArray{<:AbstractArray}) = Base.Iterators.flatten(A)
 
 
@@ -92,3 +94,107 @@ end
 
 @inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) =
     innersize(A)[dim]
+
+
+"""
+    deepgetindex(A::AbstractArray, idxs...)
+    deepgetindex(A::AbstractArray{<:AbstractArray, N}, idxs...) where {N}
+
+Recursive `getindex` on flat or nested arrays. If A is an array of arrays,
+uses the first `N` entries in `idxs` on `A`, then the rest on the inner
+array(s). If A is not a nested array, `deepgetindex` is equivalent to
+`getindex`.
+
+See also [`deepsetindex!`](@ref) and [`deepview`](@ref).
+"""
+function deepgetindex end
+export deepgetindex
+
+Base.@propagate_inbounds deepgetindex(A::AbstractArray{T,N}, idxs::Vararg{Any,N}) where {T,N} = getindex(A, idxs...)
+Base.@propagate_inbounds deepgetindex(A::AbstractArray{<:AbstractArray,N}, idxs::Vararg{Any,N}) where {N} = getindex(A, idxs...)
+
+Base.@propagate_inbounds function deepgetindex(A::AbstractArray{<:AbstractArray,N}, idxs...) where {N}
+    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
+    _deepgetindex_impl(A, idxs_outer, idxs_inner)
+end
+
+Base.@propagate_inbounds _deepgetindex_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::Tuple) where {N} =
+    deepgetindex(getindex(A, idxs_outer...), idxs_inner...)
+
+Base.@propagate_inbounds _deepgetindex_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::Tuple) where {N} =
+    _deepgetindex_tupled.(view(A, idxs_outer...), (idxs_inner,))
+
+Base.@propagate_inbounds _deepgetindex_tupled(A::AbstractArray, idxs::Tuple) = deepgetindex(A, idxs...)
+
+
+"""
+    deepsetindex!(A::AbstractArray, x, idxs...)
+    deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs...) where {N}
+
+Recursive `setindex!` on flat or nested arrays. If A is an array of arrays,
+uses the first `N` entries in `idxs` on `A`, then the rest on the inner
+array(s). If A is not a nested array, `deepsetindex!` is equivalent to
+`setindex!`.
+
+See also [`deepgetindex`](@ref) and [`deepview`](@ref).
+"""
+function deepsetindex! end
+export deepsetindex!
+
+Base.@propagate_inbounds deepsetindex!(A::AbstractArray{T,N}, x, idxs::Vararg{Any,N}) where {T,N} = setindex!(A, x, idxs...)
+Base.@propagate_inbounds deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs::Vararg{Any,N}) where {N} = setindex!(A, x, idxs...)
+
+Base.@propagate_inbounds function deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs...) where {N}
+    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
+    _deepsetindex_impl!(A, x, idxs_outer, idxs_inner)
+    A
+end
+
+Base.@propagate_inbounds function _deepsetindex_impl!(A::AbstractArray{<:AbstractArray}, x, idxs_outer::NTuple{N,Real}, idxs_inner::Tuple) where {N}
+    B = getindex(A, idxs_outer...)
+    deepsetindex!(B, x, idxs_inner...)
+end
+
+Base.@propagate_inbounds function _deepsetindex_impl!(A::AbstractArray{<:AbstractArray}, x, idxs_outer::NTuple{N,Any}, idxs_inner::Tuple) where {N}
+    B = view(A, idxs_outer...)
+    for i in eachindex(B, x)
+        deepsetindex!(B[i], x[i], idxs_inner...)
+    end
+end
+
+
+
+"""
+    deepview(A::AbstractArray, idxs...)
+    deepview(A::AbstractArray{<:AbstractArray, N}, idxs...) where {N}
+
+Recursive `view` on flat or nested arrays. If A is an array of arrays,
+uses the first `N` entries in `idxs` on `A`, then the rest on the inner
+array(s). If A is not a nested array, `deepview` is equivalent to `view`.
+
+See also [`deepgetindex`](@ref) and [`deepsetindex!`](@ref).
+"""
+function deepview end
+export deepview
+
+Base.@propagate_inbounds deepview(A::AbstractArray{T,N}, idxs::Vararg{Any,N}) where {T,N} = view(A, idxs...)
+Base.@propagate_inbounds deepview(A::AbstractArray{<:AbstractArray,N}, idxs::Vararg{Any,N}) where {N} = view(A, idxs...)
+
+Base.@propagate_inbounds function deepview(A::AbstractArray{<:AbstractArray,N}, idxs...) where {N}
+    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
+    _deepview_impl(A, idxs_outer, idxs_inner)
+end
+
+Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::NTuple{M,Real}) where {N,M} =
+    deepview(getindex(A, idxs_outer...), idxs_inner...)
+
+Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::NTuple{M,Any}) where {N,M} =
+    deepview(getindex(A, idxs_outer...), idxs_inner...)
+
+Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::NTuple{M,Real}) where {N,M} =
+    throw(ArgumentError("deepview not supported yes with outer indices $idxs_outer and inner indices $idxs_inner"))
+
+Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::NTuple{M,Any}) where {N,M} =
+    _deepview_tupled.(view(A, idxs_outer...), (idxs_inner,))
+
+Base.@propagate_inbounds _deepview_tupled(A::AbstractArray, idxs::Tuple) = deepview(A, idxs...)
diff --git a/src/util.jl b/src/util.jl
@@ -17,15 +17,18 @@ end
 Base.@pure _ncolons(::Val{N}) where N = ntuple(_ -> Colon(), Val{N}())
 
 
-Base.@propagate_inbounds front_tuple(x::NTuple{N}, ::Val{M}) where {N,M} =
+Base.@propagate_inbounds front_tuple(x::NTuple{N,Any}, ::Val{M}) where {N,M} =
     Base.ntuple(i -> x[i], Val{M}())
 
-Base.@propagate_inbounds back_tuple(x::NTuple{N}, ::Val{M}) where {N,M} =
+Base.@propagate_inbounds back_tuple(x::NTuple{N,Any}, ::Val{M}) where {N,M} =
     Base.ntuple(i -> x[i + N - M], Val{M}())
 
-Base.@propagate_inbounds split_tuple(x::NTuple{N}, ::Val{M}) where {N,M} =
+Base.@propagate_inbounds split_tuple(x::NTuple{N,Any}, ::Val{M}) where {N,M} =
     (front_tuple(x, Val{M}()), back_tuple(x, Val{N - M}()))
 
+Base.@propagate_inbounds swap_front_back_tuple(x::NTuple{N,Any}, ::Val{M}) where {N,M} =
+    (back_tuple(x, Val{N - M}())..., front_tuple(x, Val{M}())...)
+
 
 _convert_elype(::Type{T}, A::AbstractArray{T}) where {T} = A
 
diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl
@@ -85,6 +85,7 @@ using UnsafeArrays
         test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2))
     end
 
+
     @testset "construct/convert from nested arrays" begin
         test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3))
         test_from_nested(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3))
@@ -110,6 +111,69 @@ using UnsafeArrays
         test_from_nested(VectorOfSimilarVectors, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1))
     end
 
+    @testset "flatview" begin
+        A = rand_nested_similar_arrays(Val(3), Val(2))
+        B = ArrayOfSimilarArrays(A)
+        @inferred(flatview(B))[:] == collect(flatview(A))
+    end
+
+
+    @testset "deepgetindex" begin
+        A = rand_nested_similar_arrays(Val(3), Val(2))
+        B = ArrayOfSimilarArrays(A)
+
+        @test deepgetindex(A, 3, 4, 2, 1, 2) == @inferred deepgetindex(B, 3, 4, 2, 1, 2)
+        @test deepgetindex(A, 2:3, 4, 2, 1, 2) == @inferred deepgetindex(B, 2:3, 4, 2, 1, 2)
+        @test deepgetindex(A, 2:3, 2:4, 2, 1, 2) == @inferred deepgetindex(B, 2:3, 2:4, 2, 1, 2)
+        @test deepgetindex(A, 2, 4, :, 1, 2) == @inferred deepgetindex(B, 2, 4, :, 1, 2)
+        @test deepgetindex(A, 2, 4, :, 1, 1:2) == @inferred deepgetindex(B, 2, 4, :, 1, 1:2)
+        @test deepgetindex(A, 2:3, 4, :, 1, 2) == @inferred deepgetindex(B, 2:3, 4, :, 1, 2)
+        @test deepgetindex(A, 2:3, 4, :, 1, 1:2) == @inferred deepgetindex(B, 2:3, 4, :, 1, 1:2)
+    end
+
+
+    @testset "deepsetindex!" begin
+        function testdata()
+            A = rand_nested_similar_arrays(Val(3), Val(2))
+            B = ArrayOfSimilarArrays(A)
+            A, B
+        end
+
+        A, B = testdata()
+        @test deepsetindex!(A, 42, 3, 4, 2, 1, 2) == @inferred deepsetindex!(B, 42, 3, 4, 2, 1, 2)
+        @test deepgetindex(B, 3, 4, 2, 1, 2) == 42
+
+        A, B = testdata()
+        X1 = rand(2)
+        @test deepsetindex!(A, X1, 2:3, 4, 2, 1, 2) == @inferred deepsetindex!(B, X1, 2:3, 4, 2, 1, 2)
+        @test deepgetindex(B, 2:3, 4, 2, 1, 2) == X1
+
+        A, B = testdata()
+        X2 = rand(2,2)
+        @test deepsetindex!(A, X2, 2, 4, :, 1, 1:2) == @inferred deepsetindex!(B, X2, 2, 4, :, 1, 1:2)
+        @test deepgetindex(B, 2, 4, :, 1, 1:2) == X2
+
+        A, B = testdata()
+        X3 = [rand(2,2), rand(2,2)]
+        @test deepsetindex!(A, X3, 2:3, 4, :, 1, 1:2) == @inferred deepsetindex!(B, X3, 2:3, 4, :, 1, 1:2)
+        @test deepgetindex(B, 2:3, 4, :, 1, 1:2) == X3
+    end
+
+
+    @testset "deepview" begin
+        A = rand_nested_similar_arrays(Val(3), Val(2))
+        B = ArrayOfSimilarArrays(A)
+
+        @test deepview(A, 3, 4, 2, 1, 2) == @inferred deepview(B, 3, 4, 2, 1, 2)
+        @test deepgetindex(A, 2:3, 4, 2, 1, 2) == @inferred deepview(B, 2:3, 4, 2, 1, 2)
+        @test deepgetindex(A, 2:3, 2:4, 2, 1, 2) == @inferred deepview(B, 2:3, 2:4, 2, 1, 2)
+        @test deepview(A, 2, 4, :, 1, 2) == @inferred deepview(B, 2, 4, :, 1, 2)
+        @test deepview(A, 2, 4, :, 1, 1:2) == @inferred deepview(B, 2, 4, :, 1, 1:2)
+        @test deepview(A, 2:3, 4, :, 1, 2) == @inferred deepview(B, 2:3, 4, :, 1, 2)
+        @test deepview(A, 2:3, 4, :, 1, 1:2) == @inferred deepview(B, 2:3, 4, :, 1, 1:2)
+    end
+
+
     @testset "examples" begin
         A_flat = rand(2,3,4,5,6)
         A_nested = nestedview(A_flat, 2)
diff --git a/test/functions.jl b/test/functions.jl
@@ -0,0 +1,65 @@
+# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).
+
+using ArraysOfArrays
+using Test
+
+
+@testset "functions" begin
+    function gen_nested()
+        A11 = [1 2; 3 4]
+        A21 = [4 5 6; 7 8 9]
+        A12 = [10 11; 12 13; 14 15]
+        A22 = [16 17; 18 19]
+        hcat(Array[A11, A21], Array[A12, A22])
+    end
+
+
+    @testset "deepgetindex" begin
+        A = gen_nested()
+        @test @inferred(deepgetindex(A, 1, 2)) === A[1, 2]
+        @test @inferred(deepgetindex(A, 1, 2, 2, 1)) == A[1, 2][2, 1]
+        @test @inferred(deepgetindex(A, 1, 2, 1:2, 1)) == A[1, 2][1:2, 1]
+        @test @inferred(deepgetindex(A, 1, 1:2, 2, 1)) == [A[1, 1][2, 1], A[1, 2][2, 1]]
+        @test @inferred(deepgetindex(A, 1, 1:2, 1:2, 2)) == [A[1, 1][1:2, 2], A[1, 2][1:2, 2]]
+        @test_throws MethodError deepgetindex(A, 1, 2, 2, 1, 2)
+
+        B = rand(3, 4, 5)
+        @test @inferred(deepgetindex(B, 2, 3, 4)) === getindex(B, 2, 3, 4)
+        @test_throws MethodError deepgetindex(B, 6)
+        @test_throws MethodError deepgetindex(B, 2, 3)
+    end
+
+
+    @testset "deepsetindex!" begin
+        B12 = [20 21 22; 23 24 25]
+        X = [[41, 42], [43, 44]]
+
+        A = gen_nested()
+        @test @inferred(deepsetindex!(A, B12, 1, 2)) === A
+        @test A[1, 2] == B12
+
+        A = gen_nested()
+        @test @inferred(deepsetindex!(A, 42, 1, 2, 2, 1)) === A
+        @test A[1, 2][2, 1] == 42
+
+        A = gen_nested()
+        @test @inferred(deepsetindex!(A, X, 1, 1:2, 1:2, 2)) === A
+        @test deepgetindex(A, 1, 1:2, 1:2, 2) == X
+    end
+
+
+    @testset "deepview" begin
+        A = gen_nested()
+        @test @inferred(deepview(A, 1, 2)) == view(A, 1, 2)
+        @test @inferred(deepview(A, 1, 2, 2, 1)) == view(A[1, 2], 2, 1)
+        @test @inferred(deepview(A, 1, 2, 1:2, 1)) == view(A[1, 2], 1:2, 1)
+        @test_throws ArgumentError deepview(A, 1, 1:2, 2, 1)
+        @test @inferred(deepview(A, 1, 1:2, 1:2, 2)) == [A[1, 1][1:2, 2], A[1, 2][1:2, 2]]
+        @test_throws MethodError deepview(A, 1, 2, 2, 1, 2)
+
+        B = rand(3, 4, 5)
+        @test @inferred(deepview(B, 2, 3, 4)) === view(B, 2, 3, 4)
+        @test_throws MethodError deepview(B, 6)
+        @test_throws MethodError deepview(B, 2, 3)
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -3,6 +3,7 @@
 using Test
 
 Test.@testset "Package ArraysOfArrays" begin
+    include("functions.jl")
     include("array_of_similar_arrays.jl")
     include("vector_of_arrays.jl")
 end
