fix and test mixed CartesianIndex

src/DiskArrays.jl

@@ -2,6 +2,8 @@ module DiskArrays
 
 using LRUCache: LRUCache, LRU
 
+using Base: tail
+
 # Use the README as the module docs
 @doc let
     path = joinpath(dirname(@__DIR__), "README.md")

src/diskindex.jl

@@ -52,7 +52,13 @@ DiskIndex(a, i::Tuple{<:AbstractVector{<:Integer}}, batchstrategy) =
 function _resolve_indices(chunks, i, indices_pre::DiskIndex, strategy::BatchStrategy)
     inow = first(i)
     indices_new, chunksrem = process_index(inow, chunks, strategy)
-    _resolve_indices(chunksrem, Base.tail(i), merge_index(indices_pre, indices_new), strategy)
+    _resolve_indices(chunksrem, tail(i), merge_index(indices_pre, indices_new), strategy)
+end
+# Splat out CartesianIndex as regular indices
+function _resolve_indices(
+    chunks, i::Tuple{<:CartesianIndex}, indices_pre::DiskIndex, strategy::BatchStrategy
+)
+    _resolve_indices(chunks, (Tuple(i[1])..., tail(i)...), indices_pre, strategy)
 end
 _resolve_indices(::Tuple{}, ::Tuple{}, indices::DiskIndex, strategy::BatchStrategy) = indices
 # No dimension left in array, only singular indices allowed
@@ -61,17 +67,25 @@ function _resolve_indices(::Tuple{}, i, indices_pre::DiskIndex, strategy::BatchS
     (length(inow) == 1 && only(inow) == 1) || throw(ArgumentError("Trailing indices must be 1"))
     indices_new = DiskIndex(size(inow), (), size(inow), (), ())
     indices = merge_index(indices_pre, indices_new)
-    _resolve_indices((), Base.tail(i), indices, strategy)
+    _resolve_indices((), tail(i), indices, strategy)
+end
+# Splat out CartesianIndex as regular trailing indices
+function _resolve_indices(
+    ::Tuple{}, i::Tuple{<:CartesianIndex}, indices_pre::DiskIndex, strategy::BatchStrategy
+)
+    _resolve_indices((), (Tuple(i[1])..., tail(i)...), indices_pre, strategy)
 end
 # Still dimensions left, but no indices available
 function _resolve_indices(chunks, ::Tuple{}, indices_pre::DiskIndex, strategy::BatchStrategy)
     chunksnow = first(chunks)
-    arraysize_from_chunksize(chunksnow) == 1 || throw(ArgumentError("Indices can only be omitted for trailing singleton dimensions"))
+    checktrailing(arraysize_from_chunksize(chunksnow))
     indices_new = add_dimension_index(strategy)
     indices = merge_index(indices_pre, indices_new)
-    _resolve_indices(Base.tail(chunks), (), indices, strategy)
+    _resolve_indices(tail(chunks), (), indices, strategy)
 end
 
+checktrailing(i) = i == 1 || throw(ArgumentError("Indices can only be omitted for trailing singleton dimensions"))
+
 add_dimension_index(::NoBatch) = DiskIndex((), (1,), (), (1,), (1:1,))
 add_dimension_index(::Union{ChunkRead,SubRanges}) = DiskIndex((), (1,), ([()],), ([(1,)],), ([(1:1,)],))
 
@@ -98,18 +112,23 @@ Calculate indices for `i` the first chunk/s in `chunks`
 Returns a [`DiskIndex`](@ref), and the remaining chunks.
 """
 process_index(i, chunks, ::NoBatch) = process_index(i, chunks)
-process_index(inow::Integer, chunks) = DiskIndex((), (1,), (), (1,), (inow:inow,)), Base.tail(chunks)
+function process_index(I::CartesianIndex, chunks, bs::NoBatch)
+    _, chunksrem = splitchunks(I, chunks)
+    DiskIndex((), (1, 1), (), (1, 1), map(i -> i:i, Tuple(I))), chunksrem
+end
+process_index(inow::Integer, chunks) = 
+    DiskIndex((), (1,), (), (1,), (inow:inow,)), tail(chunks)
 function process_index(::Colon, chunks)
     s = arraysize_from_chunksize(first(chunks))
-    DiskIndex((s,), (s,), (Colon(),), (Colon(),), (1:s,),), Base.tail(chunks)
+    DiskIndex((s,), (s,), (Colon(),), (Colon(),), (1:s,),), tail(chunks)
 end
 function process_index(i::AbstractUnitRange{<:Integer}, chunks, ::NoBatch)
-    DiskIndex((length(i),), (length(i),), (Colon(),), (Colon(),), (i,)), Base.tail(chunks)
+    DiskIndex((length(i),), (length(i),), (Colon(),), (Colon(),), (i,)), tail(chunks)
 end
 function process_index(i::AbstractArray{<:Integer}, chunks, ::NoBatch)
     indmin, indmax = isempty(i) ? (1, 0) : extrema(i)
     di = DiskIndex(size(i), ((indmax - indmin + 1),), map(_ -> Colon(), size(i)), ((i .- (indmin - 1)),), (indmin:indmax,))
-    return di, Base.tail(chunks)
+    return di, tail(chunks)
 end
 function process_index(i::AbstractArray{Bool,N}, chunks, ::NoBatch) where {N}
     chunksnow, chunksrem = splitchunks(i, chunks)
@@ -159,10 +178,15 @@ splitchunks(i::AbstractArray{<:CartesianIndex}, chunks) =
 splitchunks(i::AbstractArray{Bool}, chunks) = splitchunks(size(i), (), chunks)
 splitchunks(i::CartesianIndices, chunks) = splitchunks(i.indices, (), chunks)
 splitchunks(i::CartesianIndex, chunks) = splitchunks(i.I, (), chunks)
-splitchunks(_, chunks) = (first(chunks),), Base.tail(chunks)
+splitchunks(_, chunks) = (first(chunks),), tail(chunks)
 splitchunks(si, chunksnow, chunksrem) =
-    splitchunks(Base.tail(si), (chunksnow..., first(chunksrem)), Base.tail(chunksrem))
+    splitchunks(tail(si), (chunksnow..., first(chunksrem)), tail(chunksrem))
+function splitchunks(si, chunksnow, chunksrem::Tuple{})
+    checktrailing(si[1])
+    splitchunks(tail(si), chunksnow, chunksrem)
+end
 splitchunks(::Tuple{}, chunksnow, chunksrem) = (chunksnow, chunksrem)
+splitchunks(::Tuple{}, chunksnow, chunksrem::Tuple{}) = (chunksnow, chunksrem)
 
 """
     output_aliasing(di::DiskIndex, ndims_dest, ndims_source)

src/indexing.jl

@@ -202,7 +202,7 @@ end
 Generate an `Array` to pass to `readblock!`
 """
 function create_outputarray(out::AbstractArray, a::AbstractArray, output_size::Tuple)
-    size(out) == output_size || throw(ArgumentError("Expected output array size of $output_size"))
+    size(out) == output_size || throw(ArgumentError("Expected output array size of $output_size, got $(size(out))"))
     return out
 end
 create_outputarray(::Nothing, a::AbstractArray, output_size::Tuple) =

test/runtests.jl

@@ -37,6 +37,7 @@ end
     @test a[CartesianIndex(1, 2), 3] == 15
     @test a[CartesianIndex(1, 2, 3)] == 15
 end
+
 @testset "isdisk" begin
     a = reshape(1:24, 2, 3, 4)
     da = AccessCountDiskArray(a; chunksize=(2, 2, 2))
@@ -51,16 +52,22 @@ end
 
 function test_getindex(a)
     @test a[2, 3, 1] == 10
+    @test a[CartesianIndex(2, 3), 1] == 10
+    @test a[2, CartesianIndex(3,), 1] == 10
+    @test a[CartesianIndex(2, 3, 1)] == 10
+    @test a[2, CartesianIndex(3, 1)] == 10
     @test a[2, 3] == 10
+    @test a[CartesianIndex(2, 3)] == 10
     @test a[2, 3, 1, 1] == 10
     @test a[:, 1] == [1, 2, 3, 4]
     @test a[1:2, 1:2, 1, 1] == [1 5; 2 6]
     @test a[end:-1:1, 1, 1] == [4, 3, 2, 1]
     @test a[2, 3, 1, 1:1] == [10]
     @test a[2, 3, 1, [1], [1]] == fill(10, 1, 1)
     @test a[:, 3, 1, [1]] == reshape(9:12, 4, 1)
+    @test a[:, CartesianIndex(3, 1), [1]] == reshape(9:12, 4, 1)
     @test a[CartesianIndices((1:2, 1:2)), 1] == [1 5; 2 6]
-    @test getindex_count(a) == 10
+    @test getindex_count(a) == 16
     # Test bitmask indexing
     m = falses(4, 5, 1)
     m[2, [1, 2, 3, 5], 1] .= true
@@ -73,6 +80,7 @@ function test_getindex(a)
     @test a[2:4:14] == [2, 6, 10, 14]
     # Test that readblock was called exactly onces for every getindex
     @test a[2:2:4, 1:2:5] == [2 10 18; 4 12 20]
+    @test a[2:2:4, 1:2:5] == [2 10 18; 4 12 20]
     @test a[[1, 3, 4], [1, 3], 1] == [1 9; 3 11; 4 12]
     @testset "allowscalar" begin
         DiskArrays.allowscalar(false)
@@ -85,7 +93,7 @@ function test_getindex(a)
 end
 
 function test_setindex(a)
-    a[1, 1, 1] = 1
+    a[CartesianIndex(1, 1), 1] = 1
     a[1, 2] = 2
     a[1, 3, 1, 1] = 3
     a[2:2, :] = [1, 2, 3, 4, 5]
@@ -117,10 +125,12 @@ function test_view(a)
     v[1:2, 1] = [1, 2]
     v[1:2, 2:3] = [4 4; 4 4]
     @test v[1:2, 1] == [1, 2]
+    @test v[1:2, CartesianIndex(1,)] == [1, 2]
+    @test v[1:2, CartesianIndex(1, 1)] == [1, 2]
     @test v[1:2, 2:3] == [4 4; 4 4]
     @test trueparent(a)[2:3, 2] == [1, 2]
     @test trueparent(a)[2:3, 3:4] == [4 4; 4 4]
-    @test getindex_count(a) == 2
+    @test getindex_count(a) == 4
     @test setindex_count(a) == 2
 
     v2 = view(a, 2:3, 2:4, Int[])