Safe co-iteration across an axis for 1+ arrays

Author: Tokazama

README.md

@@ -20,6 +20,15 @@ Returns the parent array that `x` wraps.
 Returns `true` if the size of `T` can change, in which case operations
 such as `pop!` and `popfirst!` are available for collections of type `T`.
 
+## indices(x[, d])
+
+Given an array `x`, this returns the indices along dimension `d`. If `x` is a tuple
+of arrays then the indices corresponding to dimension `d` of all arrays in `x` are
+returned. If any indices are not equal along dimension `d` an error is thrown. A
+tuple may be used to specify a different dimension for each array. If `d` is not
+specified then indices for visiting each index of `x` is returned.
+
+
 ## ismutable(x)
 
 A trait function for whether `x` is a mutable or immutable array. Used for

src/ArrayInterface.jl

@@ -542,6 +542,46 @@ known_step(x) = known_step(typeof(x))
 known_step(::Type{T}) where {T} = nothing
 known_step(::Type{<:AbstractUnitRange{T}}) where {T} = one(T)
 
+
+"""
+    indices(x[, d]) -> AbstractRange
+
+Given an array `x`, this returns the indices along dimension `d`. If `x` is a tuple
+of arrays then the indices corresponding to dimension `d` of all arrays in `x` are
+returned. If any indices are not equal along dimension `d` an error is thrown. A
+tuple may be used to specify a different dimension for each array. If `d` is not
+specified then indices for visiting each index of `x` is returned.
+"""
+@inline indices(x) = eachindex(x)
+
+indices(x, d) = indices(axes(x, d))
+
+@inline function indices(x::NTuple{N,<:Any}, dim) where {N}
+  inds = indices(first(x), dim)
+  @assert _check_indices(inds, Base.tail(x), dim) "The indices along dimension $dim are not equal for all $x"
+  return inds
+end
+
+@inline function indices(x::NTuple{N,<:Any}, dim::NTuple{N,<:Any}) where {N}
+  ind = indices(first(x), first(dim))
+  @assert _check_indices(ind, Base.tail(x), Base.tail(dim)) "The indices along dimension $dim are not equal for all $x"
+  return ind
+end
+
+@inline function _check_indices(ind, x::Tuple, dim::Tuple)
+  for (x_i, d_i) in zip(x, dim)
+    ind == indices(x_i, d_i) || return false
+  end
+  return true
+end
+
+@inline function _check_indices(ind, x::Tuple, d)
+  for x_i in x
+    ind == indices(x_i, d) || return false
+  end
+  return true
+end
+
 function __init__()
 
   @require SuiteSparse="4607b0f0-06f3-5cda-b6b1-a6196a1729e9" begin

test/runtests.jl

@@ -196,3 +196,12 @@ end
     @test !ArrayInterface.can_change_size(Tuple{})
 end
 
+@testset "indices" begin
+    @test @inferred(ArrayInterface.indices(ones(2, 3))) == 1:6
+    @test @inferred(ArrayInterface.indices(ones(2, 3), 1)) == 1:2
+    @test @inferred(ArrayInterface.indices((ones(2, 3), ones(3, 2)), (1, 2))) == 1:2
+    @test @inferred(ArrayInterface.indices((ones(2, 3), ones(2, 3)), 1)) == 1:2
+    @test_throws AssertionError ArrayInterface.indices((ones(2, 3), ones(3, 3)), 1)
+    @test_throws AssertionError ArrayInterface.indices((ones(2, 3), ones(3, 3)), (1, 2))
+end
+