trailing dimensions in eachslice (#58791)

fixes #51692

Author: adienes

base/slicearray.jl

@@ -25,7 +25,7 @@ struct Slices{P,SM,AX,S,N} <: AbstractSlices{S,N}
     """
     parent::P
     """
-    A tuple of length `ndims(parent)`, denoting how each dimension should be handled:
+    A tuple of length at least `ndims(parent)`, denoting how each dimension should be handled:
       - an integer `i`: this is the `i`th dimension of the outer `Slices` object.
       - `:`: an "inner" dimension
     """
@@ -39,34 +39,39 @@ end
 unitaxis(::AbstractArray) = Base.OneTo(1)
 
 function Slices(A::P, slicemap::SM, ax::AX) where {P,SM,AX}
+    length(slicemap) >= ndims(A) ||
+        throw(ArgumentError("Slices cannot be constructed with a slicemap of fewer elements than the parent has dimensions"))
     N = length(ax)
-    argT = map((a,l) -> l === (:) ? Colon : eltype(a), axes(A), slicemap)
+    parent_axes = ntuple(d -> axes(A, d), length(slicemap))
+    argT = map((a,l) -> l === (:) ? Colon : eltype(a), parent_axes, slicemap)
     S = Base.promote_op(view, P, argT...)
     Slices{P,SM,AX,S,N}(A, slicemap, ax)
 end
 
 _slice_check_dims(N) = nothing
 function _slice_check_dims(N, dim, dims...)
-    1 <= dim <= N || throw(DimensionMismatch("Invalid dimension $dim"))
+    1 <= dim || throw(DimensionMismatch("Invalid dimension $dim"))
     dim in dims && throw(DimensionMismatch("Dimensions $dims are not unique"))
     _slice_check_dims(N,dims...)
 end
 
 @constprop :aggressive function _eachslice(A::AbstractArray{T,N}, dims::NTuple{M,Integer}, drop::Bool) where {T,N,M}
     _slice_check_dims(N,dims...)
+    N_ = foldl(max, dims; init=N)
+
     if drop
         # if N = 4, dims = (3,1) then
         # axes = (axes(A,3), axes(A,1))
         # slicemap = (2, :, 1, :)
         ax = map(dim -> axes(A,dim), dims)
-        slicemap = ntuple(dim -> something(findfirst(isequal(dim), dims), (:)), N)
+        slicemap = ntuple(dim -> something(findfirst(isequal(dim), dims), (:)),  N_)
         return Slices(A, slicemap, ax)
     else
         # if N = 4, dims = (3,1) then
         # axes = (axes(A,1), OneTo(1), axes(A,3), OneTo(1))
         # slicemap = (1, :, 3, :)
-        ax = ntuple(dim -> dim in dims ? axes(A,dim) : unitaxis(A), N)
-        slicemap = ntuple(dim -> dim in dims ? dim : (:), N)
+        ax = ntuple(dim -> dim in dims ? axes(A,dim) : unitaxis(A), N_)
+        slicemap = ntuple(dim -> dim in dims ? dim : (:), N_)
         return Slices(A, slicemap, ax)
     end
 end

test/arrayops.jl

@@ -2402,7 +2402,7 @@ end
     M = [1 2 3; 4 5 6; 7 8 9]
     @test eachrow(M) == eachslice(M, dims = 1) == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
     @test eachcol(M) == eachslice(M, dims = 2) == [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
-    @test_throws DimensionMismatch eachslice(M, dims = 4)
+    @test eachslice(M, dims = 4) == [[1 2 3; 4 5 6; 7 8 9;;;]]
 
     SR = @inferred eachrow(M)
     @test SR[2] isa eltype(SR)
@@ -2467,6 +2467,32 @@ end
         @test_throws BoundsError A[2,3] = [4,5]
         @test_throws BoundsError A[2,3] .= [4,5]
     end
+
+    @testset "trailing dimensions" begin
+        v = collect(1:3)
+
+        S2  = eachslice(v; dims = 2, drop=true)
+        @test S2 isa AbstractSlices{<:AbstractVector, 1}
+        @test size(S2) == (1,)
+        @test S2[1] == v
+
+        S2K = eachslice(v; dims = 2, drop=false)
+        @test S2K isa AbstractSlices{<:AbstractVector, 2}
+        @test size(S2K) == (1,1)
+        @test S2K[1,1] == v
+
+        M = reshape(1:6, 2, 3)
+
+        S13 = eachslice(M; dims = (1,3))
+        @test size(S13) == (2,1)
+        @test S13[2,1] == M[2,:,1]
+
+        S13K = eachslice(M; dims = (1,3), drop=false)
+        @test size(S13K) == (2,1,1)
+        @test S13K[1,1,1] == M[1,:]
+        @test S13K[2,1,1] == M[2,:]
+    end
+
 end
 
 ###