use offset axes for offset arrays (#27038)

* Better support for non-Int axes

* Fix offset unique test:

This is an interesting case: it just falls out that unique across dimensions now preserves offset-ness of the non-uniqued dimensions

* Simplify similar(::Type{T}, ...) and require T<:AbstractArray for the default definition

* Use the same extension system for `reshape` as `similar`.

* Fixup Offset similar(::Type, ...) definition to match

base/abstractarray.jl

@@ -345,7 +345,7 @@ size_to_strides(s, d) = (s,)
 size_to_strides(s) = ()
 
 
-function isassigned(a::AbstractArray, i::Int...)
+function isassigned(a::AbstractArray, i::Integer...)
  try
  a[i...]
  true
@@ -572,7 +572,12 @@ similar(a::AbstractArray, ::Type{T}) where {T} = similar(a,
 similar(a::AbstractArray{T}, dims::Tuple) where {T} = similar(a, T, to_shape(dims))
 similar(a::AbstractArray{T}, dims::DimOrInd...) where {T} = similar(a, T, to_shape(dims))
 similar(a::AbstractArray, ::Type{T}, dims::DimOrInd...) where {T} = similar(a, T, to_shape(dims))
-similar(a::AbstractArray, ::Type{T}, dims::NeedsShaping) where {T} = similar(a, T, to_shape(dims))
+# Similar supports specifying dims as either Integers or AbstractUnitRanges or any mixed combination
+# thereof. Ideally, we'd just convert Integers to OneTos and then call a canonical method with the axes,
+# but we don't want to require all AbstractArray subtypes to dispatch on Base.OneTo. So instead we
+# define this method to convert supported axes to Ints, with the expectation that an offset array
+# package will define a method with dims::Tuple{Union{Integer, UnitRange}, Vararg{Union{Integer, UnitRange}}}
+similar(a::AbstractArray, ::Type{T}, dims::Tuple{Union{Integer, OneTo}, Vararg{Union{Integer, OneTo}}}) where {T} = similar(a, T, to_shape(dims))
 # similar creates an Array by default
 similar(a::AbstractArray, ::Type{T}, dims::Dims{N}) where {T,N} = Array{T,N}(undef, dims)
 
@@ -607,8 +612,9 @@ indices of the result will match `A`.
 would create a 1-dimensional logical array whose indices match those
 of the columns of `A`.
 """
-similar(::Type{T}, shape::Tuple) where {T} = T(to_shape(shape))
-similar(::Type{T}, dims::DimOrInd...) where {T} = similar(T, dims)
+similar(::Type{T}, dims::DimOrInd...) where {T<:AbstractArray} = similar(T, dims)
+similar(::Type{T}, shape::Tuple{Union{Integer, OneTo}, Vararg{Union{Integer, OneTo}}}) where {T<:AbstractArray} = similar(T, to_shape(shape))
+similar(::Type{T}, dims::Dims) where {T<:AbstractArray} = T(undef, dims)
 
 """
  empty(v::AbstractVector, [eltype])
@@ -1701,6 +1707,7 @@ end
 
 nextL(L, l::Integer) = L*l
 nextL(L, r::AbstractUnitRange) = L*unsafe_length(r)
+nextL(L, r::Slice) = L*unsafe_length(r.indices)
 offsetin(i, l::Integer) = i-1
 offsetin(i, r::AbstractUnitRange) = i-first(r)
 

base/abstractarraymath.jl

@@ -72,7 +72,7 @@ squeeze(A; dims) = _squeeze(A, dims)
 function _squeeze(A::AbstractArray, dims::Dims)
  for i in 1:length(dims)
  1 <= dims[i] <= ndims(A) || throw(ArgumentError("squeezed dims must be in range 1:ndims(A)"))
- length(axes(A, dims[i])) == 1 || throw(ArgumentError("squeezed dims must all be size 1"))
+ _length(axes(A, dims[i])) == 1 || throw(ArgumentError("squeezed dims must all be size 1"))
  for j = 1:i-1
  dims[j] == dims[i] && throw(ArgumentError("squeezed dims must be unique"))
  end
@@ -158,7 +158,7 @@ function reverse(A::AbstractArray; dims::Integer)
  B = similar(A)
  nnd = 0
  for i = 1:nd
- nnd += Int(length(inds[i])==1 || i==d)
+ nnd += Int(_length(inds[i])==1 || i==d)
  end
  indsd = inds[d]
  sd = first(indsd)+last(indsd)

base/array.jl

@@ -279,7 +279,6 @@ similar(a::Array{T,2}, S::Type) where {T} = Matrix{S}(undef, size(a,1)
 similar(a::Array{T}, m::Int) where {T} = Vector{T}(undef, m)
 similar(a::Array, T::Type, dims::Dims{N}) where {N} = Array{T,N}(undef, dims)
 similar(a::Array{T}, dims::Dims{N}) where {T,N} = Array{T,N}(undef, dims)
-similar(::Type{T}, shape::Tuple) where {T<:Array} = T(undef, to_shape(shape))
 
 # T[x...] constructs Array{T,1}
 """
@@ -521,9 +520,9 @@ end
 
 _collect_indices(::Tuple{}, A) = copyto!(Array{eltype(A),0}(undef), A)
 _collect_indices(indsA::Tuple{Vararg{OneTo}}, A) =
- copyto!(Array{eltype(A)}(undef, length.(indsA)), A)
+ copyto!(Array{eltype(A)}(undef, _length.(indsA)), A)
 function _collect_indices(indsA, A)
- B = Array{eltype(A)}(undef, length.(indsA))
+ B = Array{eltype(A)}(undef, _length.(indsA))
  copyto!(B, CartesianIndices(axes(B)), A, CartesianIndices(indsA))
 end
 
@@ -842,7 +841,7 @@ themselves in another collection. The result is of the preceding example is equi
 """
 function append!(a::Array{<:Any,1}, items::AbstractVector)
  itemindices = eachindex(items)
- n = length(itemindices)
+ n = _length(itemindices)
  _growend!(a, n)
  copyto!(a, length(a)-n+1, items, first(itemindices), n)
  return a
@@ -885,7 +884,7 @@ function prepend! end
 
 function prepend!(a::Array{<:Any,1}, items::AbstractVector)
  itemindices = eachindex(items)
- n = length(itemindices)
+ n = _length(itemindices)
  _growbeg!(a, n)
  if a === items
  copyto!(a, 1, items, n+1, n)

base/arrayshow.jl

@@ -78,7 +78,7 @@ function alignment(io::IO, X::AbstractVecOrMat,
  break
  end
  end
- if 1 < length(a) < length(axes(X,2))
+ if 1 < length(a) < _length(axes(X,2))
  while sum(map(sum,a)) + sep*length(a) >= cols_otherwise
  pop!(a)
  end
@@ -96,9 +96,8 @@ is specified as string sep.
 function print_matrix_row(io::IO,
  X::AbstractVecOrMat, A::Vector,
  i::Integer, cols::AbstractVector, sep::AbstractString)
- isempty(A) || first(axes(cols,1)) == 1 || throw(DimensionMismatch("indices of cols ($(axes(cols,1))) must start at 1"))
- for k = 1:length(A)
- j = cols[k]
+ for (k, j) = enumerate(cols)
+ k > length(A) && break
  if isassigned(X,Int(i),Int(j)) # isassigned accepts only `Int` indices
  x = X[i,j]
  a = alignment(io, x)
@@ -169,20 +168,20 @@ function print_matrix(io::IO, X::AbstractVecOrMat,
  @assert textwidth(hdots) == textwidth(ddots)
  sepsize = length(sep)
  rowsA, colsA = axes(X,1), axes(X,2)
- m, n = length(rowsA), length(colsA)
+ m, n = _length(rowsA), _length(colsA)
  # To figure out alignments, only need to look at as many rows as could
  # fit down screen. If screen has at least as many rows as A, look at A.
  # If not, then we only need to look at the first and last chunks of A,
  # each half a screen height in size.
  halfheight = div(screenheight,2)
  if m > screenheight
- rowsA = [rowsA[1:halfheight]; rowsA[m-div(screenheight-1,2)+1:m]]
+ rowsA = [rowsA[(0:halfheight-1) .+ firstindex(rowsA)]; rowsA[(end-div(screenheight-1,2)+1):end]]
  end
  # Similarly for columns, only necessary to get alignments for as many
  # columns as could conceivably fit across the screen
  maxpossiblecols = div(screenwidth, 1+sepsize)
  if n > maxpossiblecols
- colsA = [colsA[1:maxpossiblecols]; colsA[(n-maxpossiblecols+1):n]]
+ colsA = [colsA[(0:maxpossiblecols-1) .+ firstindex(colsA)]; colsA[(end-maxpossiblecols+1):end]]
  end
  A = alignment(io, X, rowsA, colsA, screenwidth, screenwidth, sepsize)
  # Nine-slicing is accomplished using print_matrix_row repeatedly
@@ -267,10 +266,10 @@ function show_nd(io::IO, a::AbstractArray, print_matrix::Function, label_slices:
  for i = 1:nd
  ii = idxs[i]
  ind = tailinds[i]
- if length(ind) > 10
- if ii == ind[4] && all(d->idxs[d]==first(tailinds[d]),1:i-1)
+ if _length(ind) > 10
+ if ii == ind[firstindex(ind)+3] && all(d->idxs[d]==first(tailinds[d]),1:i-1)
  for j=i+1:nd
- szj = length(axes(a, j+2))
+ szj = _length(axes(a, j+2))
  indj = tailinds[j]
  if szj>10 && first(indj)+2 < idxs[j] <= last(indj)-3
  @goto skip
@@ -280,7 +279,7 @@ function show_nd(io::IO, a::AbstractArray, print_matrix::Function, label_slices:
  print(io, "...\n\n")
  @goto skip
  end
- if ind[3] < ii <= ind[end-3]
+ if ind[firstindex(ind)+2] < ii <= ind[end-3]
  @goto skip
  end
  end
@@ -314,7 +313,7 @@ print_array(io::IO, X::AbstractArray) = show_nd(io, X, print_matrix, true)
 # implements: show(io::IO, ::MIME"text/plain", X::AbstractArray)
 function show(io::IO, ::MIME"text/plain", X::AbstractArray)
  # 0) compute new IOContext
- if !haskey(io, :compact) && length(axes(X, 2)) > 1
+ if !haskey(io, :compact) && _length(axes(X, 2)) > 1
  io = IOContext(io, :compact => true)
  end
  if get(io, :limit, false) && eltype(X) === Method
@@ -360,7 +359,7 @@ function _show_nonempty(io::IO, X::AbstractMatrix, prefix::String)
  @assert !isempty(X)
  limit = get(io, :limit, false)::Bool
  indr, indc = axes(X,1), axes(X,2)
- nr, nc = length(indr), length(indc)
+ nr, nc = _length(indr), _length(indc)
  rdots, cdots = false, false
  rr1, rr2 = UnitRange{Int}(indr), 1:0
  cr1, cr2 = UnitRange{Int}(indc), 1:0

base/bitarray.jl

@@ -339,8 +339,6 @@ similar(B::BitArray, T::Type{Bool}, dims::Dims) = BitArray(undef, dims)
 # (this triggers conversions like float(bitvector) etc.)
 similar(B::BitArray, T::Type, dims::Dims) = Array{T}(undef, dims)
 
-similar(::Type{T}, shape::Tuple) where {T<:BitArray} = T(undef, to_shape(shape))
-
 function fill!(B::BitArray, x)
  y = convert(Bool, x)
  isempty(B) && return B

base/broadcast.jl

@@ -512,7 +512,7 @@ Base.@propagate_inbounds _newindex(ax::Tuple{}, I::Tuple{}) = ()
 @inline function _newindexer(indsA::Tuple)
  ind1 = indsA[1]
  keep, Idefault = _newindexer(tail(indsA))
- (length(ind1)!=1, keep...), (first(ind1), Idefault...)
+ (Base._length(ind1)!=1, keep...), (first(ind1), Idefault...)
 end
 
 @inline function Base.getindex(bc::Broadcasted, I::Union{Int,CartesianIndex})

base/indices.jl

@@ -277,16 +277,24 @@ iterate over all the wrapped indices, even supporting offset indices.
 struct Slice{T<:AbstractUnitRange} <: AbstractUnitRange{Int}
  indices::T
 end
-axes(S::Slice) = (S.indices,)
-unsafe_indices(S::Slice) = (S.indices,)
-indices1(S::Slice) = S.indices
+Slice(S::Slice) = S
+axes(S::Slice) = (S,)
+unsafe_indices(S::Slice) = (S,)
+indices1(S::Slice) = S
+axes(S::Slice{<:OneTo}) = (S.indices,)
+unsafe_indices(S::Slice{<:OneTo}) = (S.indices,)
+indices1(S::Slice{<:OneTo}) = S.indices
+
 first(S::Slice) = first(S.indices)
 last(S::Slice) = last(S.indices)
 errmsg(A) = error("size not supported for arrays with indices $(axes(A)); see https://docs.julialang.org/en/latest/devdocs/offset-arrays/")
 size(S::Slice) = first(S.indices) == 1 ? (length(S.indices),) : errmsg(S)
 length(S::Slice) = first(S.indices) == 1 ? length(S.indices) : errmsg(S)
-unsafe_length(S::Slice) = first(S.indices) == 1 ? unsafe_length(S.indices) : errmsg(S)
+_length(S::Slice) = length(S.indices)
+unsafe_length(S::Slice) = unsafe_length(S.indices)
 getindex(S::Slice, i::Int) = (@_inline_meta; @boundscheck checkbounds(S, i); i)
+getindex(S::Slice, i::AbstractUnitRange{<:Integer}) = (@_inline_meta; @boundscheck checkbounds(S, i); i)
+getindex(S::Slice, i::StepRange{<:Integer}) = (@_inline_meta; @boundscheck checkbounds(S, i); i)
 show(io::IO, r::Slice) = print(io, "Base.Slice(", r.indices, ")")
 iterate(S::Slice, s...) = iterate(S.indices, s...)
 
@@ -352,7 +360,7 @@ LinearIndices(A::Union{AbstractArray,SimpleVector}) = LinearIndices(axes(A))
 
 # AbstractArray implementation
 IndexStyle(::Type{<:LinearIndices}) = IndexLinear()
-axes(iter::LinearIndices) = iter.indices
+axes(iter::LinearIndices) = map(indices1, iter.indices)
 size(iter::LinearIndices) = map(unsafe_length, iter.indices)
 function getindex(iter::LinearIndices, i::Int)
  @_inline_meta

base/multidimensional.jl

@@ -249,12 +249,16 @@ module IteratorsMD
  convert(Tuple{Vararg{UnitRange{Int}}}, R)
 
  # AbstractArray implementation
- Base.axes(iter::CartesianIndices{N,R}) where {N,R} = iter.indices
+ Base.axes(iter::CartesianIndices{N,R}) where {N,R} = map(Base.indices1, iter.indices)
  Base.IndexStyle(::Type{CartesianIndices{N,R}}) where {N,R} = IndexCartesian()
- @inline function Base.getindex(iter::CartesianIndices{N,R}, I::Vararg{Int, N}) where {N,R}
+ @inline function Base.getindex(iter::CartesianIndices{N,<:NTuple{N,Base.OneTo}}, I::Vararg{Int, N}) where {N}
  @boundscheck checkbounds(iter, I...)
  CartesianIndex(I)
  end
+ @inline function Base.getindex(iter::CartesianIndices{N,R}, I::Vararg{Int, N}) where {N,R}
+ @boundscheck checkbounds(iter, I...)
+ CartesianIndex(I .- first.(Base.indices1.(iter.indices)) .+ first.(iter.indices))
+ end
 
  ndims(R::CartesianIndices) = ndims(typeof(R))
  ndims(::Type{CartesianIndices{N}}) where {N} = N
@@ -319,7 +323,7 @@ module IteratorsMD
  end
 
  simd_inner_length(iter::CartesianIndices{0}, ::CartesianIndex) = 1
- simd_inner_length(iter::CartesianIndices, I::CartesianIndex) = length(iter.indices[1])
+ simd_inner_length(iter::CartesianIndices, I::CartesianIndex) = Base._length(iter.indices[1])
 
  simd_index(iter::CartesianIndices{0}, ::CartesianIndex, I1::Int) = first(iter)
  @inline function simd_index(iter::CartesianIndices, Ilast::CartesianIndex, I1::Int)
@@ -461,7 +465,7 @@ index_dimsum() = ()
 index_lengths() = ()
 @inline index_lengths(::Real, rest...) = (1, index_lengths(rest...)...)
 @inline index_lengths(A::AbstractArray, rest...) = (length(A), index_lengths(rest...)...)
-@inline index_lengths(A::Slice, rest...) = (length(indices1(A)), index_lengths(rest...)...)
+@inline index_lengths(A::Slice, rest...) = (_length(indices1(A)), index_lengths(rest...)...)
 
 # shape of array to create for getindex() with indices I, dropping scalars
 # returns a Tuple{Vararg{AbstractUnitRange}} of indices
@@ -893,7 +897,7 @@ circshift!(dest::AbstractArray, src, shiftamt) = circshift!(dest, src, (shiftamt
  inds::Tuple{AbstractUnitRange,Vararg{Any}},
  shiftamt::Tuple{Integer,Vararg{Any}})
  ind1, d = inds[1], shiftamt[1]
- s = mod(d, length(ind1))
+ s = mod(d, _length(ind1))
  sf, sl = first(ind1)+s, last(ind1)-s
  r1, r2 = first(ind1):sf-1, sf:last(ind1)
  r3, r4 = first(ind1):sl, sl+1:last(ind1)
@@ -941,7 +945,7 @@ true
 function circcopy!(dest, src)
  dest === src && throw(ArgumentError("dest and src must be separate arrays"))
  indssrc, indsdest = axes(src), axes(dest)
- if (szsrc = map(length, indssrc)) != (szdest = map(length, indsdest))
+ if (szsrc = map(_length, indssrc)) != (szdest = map(_length, indsdest))
  throw(DimensionMismatch("src and dest must have the same sizes (got $szsrc and $szdest)"))
  end
  shift = map((isrc, idest)->first(isrc)-first(idest), indssrc, indsdest)
@@ -953,7 +957,7 @@ end
 @inline function _circcopy!(dest, rdest, indsdest::Tuple{AbstractUnitRange,Vararg{Any}},
  src, rsrc, indssrc::Tuple{AbstractUnitRange,Vararg{Any}})
  indd1, inds1 = indsdest[1], indssrc[1]
- l = length(indd1)
+ l = _length(indd1)
  s = mod(first(inds1)-first(indd1), l)
  sdf = first(indd1)+s
  rd1, rd2 = first(indd1):sdf-1, sdf:last(indd1)

base/reducedim.jl

@@ -169,7 +169,7 @@ function check_reducedims(R, A)
  had_nonreduc = false
  for i = 1:ndims(A)
  Ri, Ai = axes(R, i), axes(A, i)
- sRi, sAi = length(Ri), length(Ai)
+ sRi, sAi = _length(Ri), _length(Ai)
  if sRi == 1
  if sAi > 1
  if had_nonreduc
@@ -804,4 +804,4 @@ function _findmax(A, region)
  end
 end
 
-reducedim1(R, A) = length(indices1(R)) == 1
+reducedim1(R, A) = _length(indices1(R)) == 1