Row and Column Iterator for Matrices (#29749)

* first pass, sans tests

* add test

* refactor a bit

* fix eltype to show views

* fix whitespace

* fix typo

* rename [ci skip]

* typofix [ci skip]

* add eachslice() [ci skip]

* eachcolumn => eachcol and a bit of error handling [ci skip]

* feedback

* fix whitespace

* refactor eachslice; type-stable now?

* git feedback

* more git feedback

* remove default for eachslice

* add tests

* prune eachrow eachcol exports and newline fix?

* typo in exports

* update docs for AbstractArray [ci skip]

* Update test/arrayops.jl

Co-Authored-By: arnavs <[email protected]>

* Update base/abstractarraymath.jl

Co-Authored-By: arnavs <[email protected]>

* bikeshedding

* Add missing backticks and improve message

[ci skip]

* add exports for eachcol, eachrow

[ci skip]

Author: Arnav Sood

base/abstractarraymath.jl

@@ -407,3 +407,40 @@ _reperr(s, n, N) = throw(ArgumentError("number of " * s * " repetitions " *
 
     return R
 end
+
+"""
+    eachrow(A::AbstractVecOrMat)
+
+Create a generator that iterates over the first dimension of vector or matrix `A`,
+returning the rows as views.
+See also [`eachcol`](@ref) and [`eachslice`](@ref).
+"""
+eachrow(A::AbstractVecOrMat) = (view(A, i, :) for i in axes(A, 1))
+
+
+"""
+    eachcol(A::AbstractVecOrMat)
+
+Create a generator that iterates over the second dimension of matrix `A`, returning the
+columns as views.
+See also [`eachrow`](@ref) and [`eachslice`](@ref).
+"""
+eachcol(A::AbstractVecOrMat) = (view(A, :, i) for i in axes(A, 2))
+
+"""
+    eachslice(A::AbstractArray; dims)
+
+Create a generator that iterates over dimensions `dims` of `A`, returning views that select all
+the data from the other dimensions in `A`.
+
+Only a single dimension in `dims` is currently supported. Equivalent to `(view(A,:,:,...,i,:,:
+...)) for i in axes(A, dims))`, where `i` is in position `dims`.
+See also [`eachrow`](@ref), [`eachcol`](@ref), and [`selectdim`](@ref).
+"""
+@inline function eachslice(A::AbstractArray; dims)
+    length(dims) == 1 || throw(ArgumentError("only single dimensions are supported"))
+    dim = first(dims)
+    dim <= ndims(A) || throw(DimensionMismatch("A doesn't have $dim dimensions"))
+    idx1, idx2 = ntuple(d->(:), dim-1), ntuple(d->(:), ndims(A)-dim)
+    return (view(A, idx1..., i, idx2...) for i in axes(A, dim))
+end

base/exports.jl

@@ -365,7 +365,10 @@ export
     cumsum!,
     accumulate,
     accumulate!,
+    eachcol,
     eachindex,
+    eachrow,
+    eachslice,
     extrema,
     fill!,
     fill,

test/arrayops.jl

@@ -1971,6 +1971,22 @@ end
     @test IndexStyle(selectdim(A, 3, 1)) == IndexStyle(view(A, :, :, 1)) == IndexLinear()
 end
 
+# row/column/slice iterator tests
+using Base: eachrow, eachcol
+@testset "row/column/slice iterators" begin
+    # Simple ones
+    M = [1 2 3; 4 5 6; 7 8 9]
+    @test collect(eachrow(M)) == collect(eachslice(M, dims = 1)) == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+    @test collect(eachcol(M)) == collect(eachslice(M, dims = 2)) == [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
+    @test_throws DimensionMismatch eachslice(M, dims = 4)
+
+    # Higher-dimensional case
+    M = reshape([(1:16)...], 2, 2, 2, 2)
+    @test_throws MethodError collect(eachrow(M))
+    @test_throws MethodError collect(eachcol(M))
+    @test collect(eachslice(M, dims = 1))[1][:, :, 1] == [1 5; 3 7]
+end
+
 ###
 ### IndexCartesian workout
 ###