Merge pull request #17137 from JuliaLang/teh/safer_indices

Refactor API for unconventionally-indexed arrays

Author: timholy

base/abstractarray.jl

@@ -211,17 +211,17 @@ function flipdim{T}(A::Array{T}, d::Integer)
 end
 
 function rotl90(A::AbstractMatrix)
-    B = similar_transpose(A)
-    ind2 = indices(A,2)
+    ind1, ind2 = indices(A)
+    B = similar(A, (ind2,ind1))
     n = first(ind2)+last(ind2)
     for i=indices(A,1), j=ind2
         B[n-j,i] = A[i,j]
     end
     return B
 end
 function rotr90(A::AbstractMatrix)
-    B = similar_transpose(A)
-    ind1 = indices(A,1)
+    ind1, ind2 = indices(A)
+    B = similar(A, (ind2,ind1))
     m = first(ind1)+last(ind1)
     for i=ind1, j=indices(A,2)
         B[j,m-i] = A[i,j]
@@ -246,10 +246,6 @@ end
 rotr90(A::AbstractMatrix, k::Integer) = rotl90(A,-k)
 rot180(A::AbstractMatrix, k::Integer) = mod(k, 2) == 1 ? rot180(A) : copy(A)
 
-similar_transpose(A::AbstractMatrix) = similar_transpose(indicesbehavior(A), A)
-similar_transpose(::IndicesStartAt1, A::AbstractMatrix) = similar(A, (size(A,2), size(A,1)))
-similar_transpose(::IndicesBehavior, A::AbstractMatrix) = similar(A, (indices(A,2), indices(A,1)))
-
 ## Transpose ##
 transpose!(B::AbstractMatrix, A::AbstractMatrix) = transpose_f!(transpose, B, A)
 ctranspose!(B::AbstractMatrix, A::AbstractMatrix) = transpose_f!(ctranspose, B, A)
@@ -315,11 +311,13 @@ function ccopy!(B, A)
 end
 
 function transpose(A::AbstractMatrix)
-    B = similar_transpose(A)
+    ind1, ind2 = indices(A)
+    B = similar(A, (ind2, ind1))
     transpose!(B, A)
 end
 function ctranspose(A::AbstractMatrix)
-    B = similar_transpose(A)
+    ind1, ind2 = indices(A)
+    B = similar(A, (ind2, ind1))
     ctranspose!(B, A)
 end
 ctranspose{T<:Real}(A::AbstractVecOrMat{T}) = transpose(A)

base/arraymath.jl

@@ -47,6 +47,8 @@ end
 
 isassigned{N}(B::BitArray{N}, i::Int) = 1 <= i <= length(B)
 
+linearindexing{A<:BitArray}(::Type{A}) = LinearFast()
+
 ## aux functions ##
 
 const _msk64 = ~UInt64(0)

base/bitarray.jl

@@ -3,7 +3,7 @@
 module Broadcast
 
 using Base.Cartesian
-using Base: promote_op, promote_eltype, promote_eltype_op, @get!, _msk_end, unsafe_bitgetindex, shape, linearindices, allocate_for, tail, dimlength
+using Base: promote_op, promote_eltype, promote_eltype_op, @get!, _msk_end, unsafe_bitgetindex, linearindices, to_shape, tail, dimlength, OneTo
 import Base: .+, .-, .*, ./, .\, .//, .==, .<, .!=, .<=, .÷, .%, .<<, .>>, .^
 export broadcast, broadcast!, bitbroadcast
 export broadcast_getindex, broadcast_setindex!
@@ -13,8 +13,8 @@ export broadcast_getindex, broadcast_setindex!
 ## Calculate the broadcast shape of the arguments, or error if incompatible
 # array inputs
 broadcast_shape() = ()
-broadcast_shape(A) = shape(A)
-@inline broadcast_shape(A, B...) = broadcast_shape((), shape(A), map(shape, B)...)
+broadcast_shape(A) = indices(A)
+@inline broadcast_shape(A, B...) = broadcast_shape((), indices(A), map(indices, B)...)
 # shape inputs
 broadcast_shape(shape::Tuple) = shape
 @inline broadcast_shape(shape::Tuple, shape1::Tuple, shapes::Tuple...) = broadcast_shape(_bcs((), shape, shape1), shapes...)
@@ -39,8 +39,10 @@ _bcsm(a::Number, b::Number) = a == b || b == 1
 ## Check that all arguments are broadcast compatible with shape
 ## Check that all arguments are broadcast compatible with shape
 # comparing one input against a shape
+check_broadcast_shape(::Tuple{}) = nothing
+check_broadcast_shape(::Tuple{}, A::Union{AbstractArray,Number}) = check_broadcast_shape((), indices(A))
 check_broadcast_shape(shp) = nothing
-check_broadcast_shape(shp, A) = check_broadcast_shape(shp, shape(A))
+check_broadcast_shape(shp, A) = check_broadcast_shape(shp, indices(A))
 check_broadcast_shape(::Tuple{}, ::Tuple{}) = nothing
 check_broadcast_shape(shp, ::Tuple{}) = nothing
 check_broadcast_shape(::Tuple{}, Ashp::Tuple) = throw(DimensionMismatch("cannot broadcast array to have fewer dimensions"))
@@ -133,7 +135,7 @@ end
 end
 
 @inline function broadcast!{nargs}(f, B::AbstractArray, As::Vararg{Any,nargs})
-    check_broadcast_shape(shape(B), As...)
+    check_broadcast_shape(indices(B), As...)
     sz = size(B)
     mapindex = map(x->newindexer(sz, x), As)
     _broadcast!(f, B, mapindex, As, Val{nargs})
@@ -179,20 +181,20 @@ function broadcast_t(f, ::Type{Any}, As...)
     shp = broadcast_shape(As...)
     iter = CartesianRange(shp)
     if isempty(iter)
-        return allocate_for(Array{Union{}}, As, shp)
+        return similar(Array{Union{}}, shp)
     end
     nargs = length(As)
     sz = size(iter)
     indexmaps = map(x->newindexer(sz, x), As)
     st = start(iter)
     I, st = next(iter, st)
     val = f([ As[i][newindex(I, indexmaps[i])] for i=1:nargs ]...)
-    B = allocate_for(Array{typeof(val)}, As, shp)
+    B = similar(Array{typeof(val)}, shp)
     B[I] = val
     return _broadcast!(f, B, indexmaps, As, Val{nargs}, iter, st, 1)
 end
 
-@inline broadcast_t(f, T, As...) = broadcast!(f, allocate_for(Array{T}, As, broadcast_shape(As...)), As...)
+@inline broadcast_t(f, T, As...) = broadcast!(f, similar(Array{T}, broadcast_shape(As...)), As...)
 
 @inline broadcast(f, As...) = broadcast_t(f, promote_eltype_op(f, As...), As...)
 
@@ -215,15 +217,15 @@ function broadcast(f, As...)
 end
 =#
 
-@inline bitbroadcast(f, As...) = broadcast!(f, allocate_for(BitArray, As, broadcast_shape(As...)), As...)
+@inline bitbroadcast(f, As...) = broadcast!(f, similar(BitArray, broadcast_shape(As...)), As...)
 
-broadcast_getindex(src::AbstractArray, I::AbstractArray...) = broadcast_getindex!(Array{eltype(src)}(broadcast_shape(I...)), src, I...)
+broadcast_getindex(src::AbstractArray, I::AbstractArray...) = broadcast_getindex!(Array{eltype(src)}(to_shape(broadcast_shape(I...))), src, I...)
 @generated function broadcast_getindex!(dest::AbstractArray, src::AbstractArray, I::AbstractArray...)
     N = length(I)
     Isplat = Expr[:(I[$d]) for d = 1:N]
     quote
         @nexprs $N d->(I_d = I[d])
-        check_broadcast_shape(size(dest), $(Isplat...))  # unnecessary if this function is never called directly
+        check_broadcast_shape(indices(dest), $(Isplat...))  # unnecessary if this function is never called directly
         checkbounds(src, $(Isplat...))
         @nloops $N i dest d->(@nexprs $N k->(j_d_k = size(I_k, d) == 1 ? 1 : i_d)) begin
             @nexprs $N k->(@inbounds J_k = @nref $N I_k d->j_d_k)
@@ -240,22 +242,22 @@ end
         @nexprs $N d->(I_d = I[d])
         checkbounds(A, $(Isplat...))
         shape = broadcast_shape($(Isplat...))
-        @nextract $N shape d->(length(shape) < d ? 1 : shape[d])
+        @nextract $N shape d->(length(shape) < d ? OneTo(1) : shape[d])
         if !isa(x, AbstractArray)
-            @nloops $N i d->(1:shape_d) d->(@nexprs $N k->(j_d_k = size(I_k, d) == 1 ? 1 : i_d)) begin
+            xA = convert(eltype(A), x)
+            @nloops $N i d->shape_d d->(@nexprs $N k->(j_d_k = size(I_k, d) == 1 ? 1 : i_d)) begin
                 @nexprs $N k->(@inbounds J_k = @nref $N I_k d->j_d_k)
-                @inbounds (@nref $N A J) = x
+                @inbounds (@nref $N A J) = xA
             end
         else
             X = x
-            # To call setindex_shape_check, we need to create fake 1-d indexes of the proper size
-            @nexprs $N d->(fakeI_d = 1:shape_d)
-            @ncall $N Base.setindex_shape_check X shape
-            k = 1
-            @nloops $N i d->(1:shape_d) d->(@nexprs $N k->(j_d_k = size(I_k, d) == 1 ? 1 : i_d)) begin
-                @nexprs $N k->(@inbounds J_k = @nref $N I_k d->j_d_k)
-                @inbounds (@nref $N A J) = X[k]
-                k += 1
+            @nexprs $N d->(shapelen_d = dimlength(shape_d))
+            @ncall $N Base.setindex_shape_check X shapelen
+            Xstate = start(X)
+            @inbounds @nloops $N i d->shape_d d->(@nexprs $N k->(j_d_k = size(I_k, d) == 1 ? 1 : i_d)) begin
+                @nexprs $N k->(J_k = @nref $N I_k d->j_d_k)
+                x_el, Xstate = next(X, Xstate)
+                (@nref $N A J) = x_el
             end
         end
         A
@@ -271,22 +273,22 @@ end
 
 eltype_plus(As::AbstractArray...) = promote_eltype_op(+, As...)
 
-.+(As::AbstractArray...) = broadcast!(+, Array{eltype_plus(As...)}(broadcast_shape(As...)), As...)
+.+(As::AbstractArray...) = broadcast!(+, Array{eltype_plus(As...)}(to_shape(broadcast_shape(As...))), As...)
 
 function .-(A::AbstractArray, B::AbstractArray)
-    broadcast!(-, Array{promote_op(-, eltype(A), eltype(B))}(broadcast_shape(A,B)), A, B)
+    broadcast!(-, Array{promote_op(-, eltype(A), eltype(B))}(to_shape(broadcast_shape(A,B))), A, B)
 end
 
 eltype_mul(As::AbstractArray...) = promote_eltype_op(*, As...)
 
-.*(As::AbstractArray...) = broadcast!(*, Array{eltype_mul(As...)}(broadcast_shape(As...)), As...)
+.*(As::AbstractArray...) = broadcast!(*, Array{eltype_mul(As...)}(to_shape(broadcast_shape(As...))), As...)
 
 function ./(A::AbstractArray, B::AbstractArray)
-    broadcast!(/, Array{promote_op(/, eltype(A), eltype(B))}(broadcast_shape(A, B)), A, B)
+    broadcast!(/, Array{promote_op(/, eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
 end
 
 function .\(A::AbstractArray, B::AbstractArray)
-    broadcast!(\, Array{promote_op(\, eltype(A), eltype(B))}(broadcast_shape(A, B)), A, B)
+    broadcast!(\, Array{promote_op(\, eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
 end
 
 typealias RatIntT{T<:Integer} Union{Type{Rational{T}},Type{T}}
@@ -296,11 +298,11 @@ type_rdiv{T<:Integer,S<:Integer}(::RatIntT{T}, ::RatIntT{S}) =
 type_rdiv{T<:Integer,S<:Integer}(::CRatIntT{T}, ::CRatIntT{S}) =
     Complex{Rational{promote_type(T,S)}}
 function .//(A::AbstractArray, B::AbstractArray)
-    broadcast!(//, Array{type_rdiv(eltype(A), eltype(B))}(broadcast_shape(A, B)), A, B)
+    broadcast!(//, Array{type_rdiv(eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
 end
 
 function .^(A::AbstractArray, B::AbstractArray)
-    broadcast!(^, Array{promote_op(^, eltype(A), eltype(B))}(broadcast_shape(A, B)), A, B)
+    broadcast!(^, Array{promote_op(^, eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
 end
 
 # ## element-wise comparison operators returning BitArray ##
@@ -363,10 +365,10 @@ for (f, scalarf) in ((:.==, :(==)),
                                               :((A,ind)->A), :((B,ind)->B[ind])),
                                              (:AbstractArray, :Any, :A,
                                               :((A,ind)->A[ind]), :((B,ind)->B)))
-        shape = :(shape($active))
+        shape = :(indices($active))
         @eval begin
             function ($f)(A::$sigA, B::$sigB)
-                P = allocate_for(BitArray, $active, $shape)
+                P = similar(BitArray, $shape)
                 F = parent(P)
                 l = length(F)
                 l == 0 && return F

base/broadcast.jl

@@ -766,6 +766,12 @@ function first(::Colon)
     1
 end
 
+# Not exported, but may be useful just in case
+function Broadcast.check_broadcast_shape(sz::Dims, As::Union{AbstractArray,Number}...)
+    depwarn("check_broadcast_shape(size(A), B...) should be replaced with check_broadcast_shape(indices(A), B...)", :check_broadcast_shape)
+    Broadcast.check_broadcast_shape(map(OneTo, sz), As...)
+end
+
 @deprecate slice view
 @deprecate sub view
 

base/deprecated.jl

@@ -171,7 +171,7 @@ start(v::SimpleVector) = 1
 next(v::SimpleVector,i) = (v[i],i+1)
 done(v::SimpleVector,i) = (i > v.length)
 isempty(v::SimpleVector) = (v.length == 0)
-indices(v::SimpleVector, d) = d == 1 ? (1:length(v)) : (1:1)
+indices(v::SimpleVector) = (OneTo(length(v)),)
 linearindices(v::SimpleVector) = indices(v, 1)
 
 function ==(v1::SimpleVector, v2::SimpleVector)

base/essentials.jl

@@ -28,6 +28,7 @@ export
 # Types
     AbstractChannel,
     AbstractMatrix,
+    AbstractUnitRange,
     AbstractVector,
     AbstractVecOrMat,
     Array,
@@ -485,7 +486,6 @@ export
     zeta,
 
 # arrays
-    allocate_for,
     bitbroadcast,
     broadcast!,
     broadcast,
@@ -585,7 +585,6 @@ export
     searchsortedlast,
     select!,
     select,
-    shape,
     shuffle,
     shuffle!,
     size,

base/exports.jl

@@ -10,9 +10,6 @@ using Base: LinearFast, LinearSlow, AbstractCartesianIndex, fill_to_length, tail
 
 export CartesianIndex, CartesianRange
 
-# Traits for linear indexing
-linearindexing{A<:BitArray}(::Type{A}) = LinearFast()
-
 # CartesianIndex
 immutable CartesianIndex{N} <: AbstractCartesianIndex{N}
     I::NTuple{N,Int}
@@ -78,12 +75,12 @@ end
 CartesianRange{N}(index::CartesianIndex{N}) = CartesianRange(one(index), index)
 CartesianRange(::Tuple{}) = CartesianRange{CartesianIndex{0}}(CartesianIndex{0}(()),CartesianIndex{0}(()))
 CartesianRange{N}(sz::NTuple{N,Int}) = CartesianRange(CartesianIndex(sz))
-CartesianRange{N}(rngs::NTuple{N,Union{Int,UnitRange{Int}}}) = CartesianRange(CartesianIndex(map(r->first(r), rngs)), CartesianIndex(map(r->last(r), rngs)))
+CartesianRange{N}(rngs::NTuple{N,Union{Integer,AbstractUnitRange}}) = CartesianRange(CartesianIndex(map(r->first(r), rngs)), CartesianIndex(map(r->last(r), rngs)))
 
 ndims(R::CartesianRange) = length(R.start)
 ndims{I<:CartesianIndex}(::Type{CartesianRange{I}}) = length(I)
 
-eachindex(::LinearSlow, A::AbstractArray) = CartesianRange(size(A))
+eachindex(::LinearSlow, A::AbstractArray) = CartesianRange(indices(A))
 
 @inline eachindex(::LinearSlow, A::AbstractArray, B::AbstractArray...) = CartesianRange(maxsize((), A, B...))
 maxsize(sz) = sz
@@ -178,18 +175,18 @@ index_lengths_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
 # whose length is equal to the dimension we're to process next. This
 # allows us to dispatch, which is important for the type-stability of
 # the lines involving Colon as the final index.
-index_shape(A::AbstractVector, I::Colon) = shape(A)
+index_shape(A::AbstractVector, I::Colon) = indices(A)
 index_shape(A::AbstractArray,  I::Colon) = (length(A),)
 @inline index_shape(A::AbstractArray, I...) = index_shape_dim(A, (true,), I...)
 @inline index_shape_dim(A, dim, ::Colon) = (trailingsize(A, length(dim)),)
-@inline index_shape_dim{T,N}(A::AbstractArray{T,N}, dim::NTuple{N}, ::Colon) = (shape(A, N),)
+@inline index_shape_dim{T,N}(A::AbstractArray{T,N}, dim::NTuple{N}, ::Colon) = (indices(A, N),)
 @inline index_shape_dim(A, dim, I::Real...) = ()
-@inline index_shape_dim(A, dim, ::Colon, i, I...) = (shape(A, length(dim)), index_shape_dim(A, (dim...,true), i, I...)...)
+@inline index_shape_dim(A, dim, ::Colon, i, I...) = (indices(A, length(dim)), index_shape_dim(A, (dim...,true), i, I...)...)
 @inline index_shape_dim(A, dim, ::Real, I...) = (index_shape_dim(A, (dim...,true), I...)...)
 @inline index_shape_dim{N}(A, dim, ::CartesianIndex{N}, I...) = (index_shape_dim(A, (dim...,ntuple(d->true,Val{N})...), I...)...)
-@inline index_shape_dim(A, dim, i::AbstractArray, I...) = (shape(i)..., index_shape_dim(A, (dim...,true), I...)...)
+@inline index_shape_dim(A, dim, i::AbstractArray, I...) = (indices(i)..., index_shape_dim(A, (dim...,true), I...)...)
 @inline index_shape_dim(A, dim, i::AbstractArray{Bool}, I...) = (sum(i), index_shape_dim(A, (dim...,true), I...)...)
-@inline index_shape_dim{N}(A, dim, i::AbstractArray{CartesianIndex{N}}, I...) = (shape(i)..., index_shape_dim(A, (dim...,ntuple(d->true,Val{N})...), I...)...)
+@inline index_shape_dim{N}(A, dim, i::AbstractArray{CartesianIndex{N}}, I...) = (indices(i)..., index_shape_dim(A, (dim...,ntuple(d->true,Val{N})...), I...)...)
 
 @inline decolon(A::AbstractVector, ::Colon) = (indices(A,1),)
 @inline decolon(A::AbstractArray,  ::Colon) = (1:length(A),)
@@ -291,8 +288,6 @@ end
     end
 end
 
-dimlength(r::Range) = length(r)
-dimlength(i::Integer) = i
 @noinline throw_checksize_error(A, sz) = throw(DimensionMismatch("output array is the wrong size; expected $sz, got $(size(A))"))
 
 ## setindex! ##
@@ -787,7 +782,7 @@ If `dim` is specified, returns unique regions of the array `itr` along `dim`.
 @generated function unique{T,N}(A::AbstractArray{T,N}, dim::Int)
     quote
         1 <= dim <= $N || return copy(A)
-        hashes = allocate_for(inds->zeros(UInt, inds), A, shape(A, dim))
+        hashes = similar(inds->zeros(UInt, inds), indices(A, dim))
 
         # Compute hash for each row
         k = 0
@@ -796,15 +791,15 @@ If `dim` is specified, returns unique regions of the array `itr` along `dim`.
         end
 
         # Collect index of first row for each hash
-        uniquerow = allocate_for(Array{Int}, A, shape(A, dim))
+        uniquerow = similar(Array{Int}, indices(A, dim))
         firstrow = Dict{Prehashed,Int}()
         for k = indices(A, dim)
             uniquerow[k] = get!(firstrow, Prehashed(hashes[k]), k)
         end
         uniquerows = collect(values(firstrow))
 
         # Check for collisions
-        collided = allocate_for(falses, A, shape(A, dim))
+        collided = similar(falses, indices(A, dim))
         @inbounds begin
             @nloops $N i A d->(if d == dim
                 k = i_d
@@ -819,7 +814,7 @@ If `dim` is specified, returns unique regions of the array `itr` along `dim`.
         end
 
         if any(collided)
-            nowcollided = allocate_for(BitArray, A, shape(A, dim))
+            nowcollided = similar(BitArray, indices(A, dim))
             while any(collided)
                 # Collect index of first row for each collided hash
                 empty!(firstrow)

base/multidimensional.jl

@@ -7,7 +7,7 @@ isinteger(x::Integer) = true
 size(x::Number) = ()
 size(x::Number,d) = convert(Int,d)<1 ? throw(BoundsError()) : 1
 indices(x::Number) = ()
-indices(x::Number,d) = convert(Int,d)<1 ? throw(BoundsError()) : (1:1)
+indices(x::Number,d) = convert(Int,d)<1 ? throw(BoundsError()) : OneTo(1)
 eltype{T<:Number}(::Type{T}) = T
 ndims(x::Number) = 0
 ndims{T<:Number}(::Type{T}) = 0