Make unvetted size throw an error for arrays with non-1 indexing

Author: timholy

base/abstractarray.jl

@@ -23,14 +23,44 @@ end
 
 size{T,N}(t::AbstractArray{T,N}, d) = d <= N ? size(t)[d] : 1
 size{N}(x, d1::Integer, d2::Integer, dx::Vararg{Integer, N}) = (size(x, d1), size(x, d2), ntuple(k->size(x, dx[k]), Val{N})...)
+
+# trait to indicate "vetted" usage for indexing that may not start at 1
+abstract IndicesSafety
+immutable SafeIndices <: IndicesSafety end
+immutable UnsafeIndices <: IndicesSafety end
+
+macro safeindices(ex)
+    esc(_safeindices(ex))
+end
+
+function _safeindices(ex::Expr)
+    if ex.head == :call
+        f = ex.args[1]
+        if f == :size || f == :length
+            return Expr(:call, f, :(Base.SafeIndices()), ex.args[2:end]...)
+        end
+    end
+    return Expr(ex.head, map(_safeindices, ex.args)...)
+end
+
+_safeindices(arg) = arg
+
+# The default is that size is safe, but array types that use non-1
+# indexing should specialize this to make it unsafe without
+# SafeIndices().
+size( ::IndicesSafety, A)                   = size(A)
+size( ::IndicesSafety, A::AbstractArray, d) = size(A,d)  # fixme
+# size(s::IndicesSafety, A::AbstractArray, d) = d <= ndims(A) ? size(s, A)[d] : 1
+size{N}(s::IndicesSafety, A::AbstractArray, d1::Integer, d2::Integer, dx::Vararg{Integer, N}) = (size(s, A, d1), size(s, A, d2), ntuple(k->size(s, A, dx[k]), Val{N})...)
+
 """
     indices(A, d)
 
 Returns the valid range of indices for array `A` along dimension `d`.
 """
 function indices(A::AbstractArray, d)
     @_inline_meta
-    1:size(A,d)
+    1:size(SafeIndices(),A,d)
 end
 """
     indices(A)
@@ -61,15 +91,17 @@ is `indices(A, 1)`.
 Calling this function is the "safe" way to write algorithms that
 exploit linear indexing.
 """
-linearindices(A) = 1:length(A)
+linearindices(A) = 1:length(SafeIndices(), A)
 linearindices(A::AbstractVector) = indices1(A)
 eltype{T}(::Type{AbstractArray{T}}) = T
 eltype{T,N}(::Type{AbstractArray{T,N}}) = T
 elsize{T}(::AbstractArray{T}) = sizeof(T)
 ndims{T,N}(::AbstractArray{T,N}) = N
 ndims{T,N}(::Type{AbstractArray{T,N}}) = N
 ndims{T<:AbstractArray}(::Type{T}) = ndims(supertype(T))
-length(t::AbstractArray) = prod(size(t))::Int
+length(s::IndicesSafety, t::AbstractArray) = prod(size(s,t))
+length(s::IndicesSafety, t) = length(t)
+length(t::AbstractArray) = length(UnsafeIndices(), t)
 endof(a::AbstractArray) = length(a)
 first(a::AbstractArray) = a[first(eachindex(a))]
 
@@ -120,13 +152,14 @@ function isassigned(a::AbstractArray, i::Int...)
 end
 
 # used to compute "end" for last index
-function trailingsize(A, n)
-    s = 1
+function trailingsize(s::IndicesSafety, A, n)
+    sz = 1
     for i=n:ndims(A)
-        s *= size(A,i)
+        sz *= size(s,A,i)
     end
-    return s
+    return sz
 end
+trailingsize(A, n) = trailingsize(UnsafeIndices(), A, n)
 
 ## Traits for array types ##
 
@@ -178,21 +211,22 @@ start at something different from 1), it is equivalent to `indices(A,
 d)`.
 """
 shape(a, d) = shape(indicesbehavior(a), a, d)
-shape(::IndicesStartAt1, a) = size(a)
-shape(::IndicesStartAt1, a, d) = size(a, d)
+shape(::IndicesStartAt1, a) = size(SafeIndices(), a)
+shape(::IndicesStartAt1, a, d) = size(SafeIndices(), a, d)
 shape(::IndicesBehavior, a) = indices(a)
 shape(::IndicesBehavior, a, d) = indices(a, d)
 
 ## Bounds checking ##
-@generated function trailingsize{T,N,n}(A::AbstractArray{T,N}, ::Type{Val{n}})
+@generated function trailingsize{T,N,n}(s::IndicesSafety, A::AbstractArray{T,N}, ::Type{Val{n}})
     (isa(n, Int) && isa(N, Int)) || error("Must have concrete type")
     n > N && return 1
-    ex = :(size(A, $n))
+    ex = :(size(s, A, $n))
     for m = n+1:N
-        ex = :($ex * size(A, $m))
+        ex = :($ex * size(s, A, $m))
     end
     Expr(:block, Expr(:meta, :inline), ex)
 end
+trailingsize{n}(A::AbstractArray, ::Type{Val{n}}) = trailingsize(UnsafeIndices(), A, Val{n})
 
 # check along a single dimension
 """
@@ -265,13 +299,9 @@ _chkbnds(A::AbstractArray, ::NTuple{1,Bool}, I::AbstractVector{Bool}) = length(A
 _chkbnds(A::AbstractVector, ::NTuple{1,Bool}, I::AbstractArray{Bool}) = length(A) == length(I)
 _chkbnds(A::AbstractVector, ::NTuple{1,Bool}, I::AbstractVector{Bool}) = indices(A) == indices(I)
 # Linear indexing:
-function _chkbnds(A::AbstractVector, ::NTuple{1,Bool}, I)
-    @_inline_meta
-    checkindex(Bool, indices1(A), I)
-end
 function _chkbnds(A::AbstractArray, ::NTuple{1,Bool}, I)
     @_inline_meta
-    checkindex(Bool, 1:length(A), I)
+    checkindex(Bool, linearindices(A), I)
 end
 # When all indices have been checked:
 _chkbnds{M}(A, checked::NTuple{M,Bool}) = checked[M]
@@ -359,7 +389,7 @@ similar(   a::AbstractArray, T::Type, dims::DimsInteger) = similar(a, T, convert
 # similar creates an Array by default
 similar(   a::AbstractArray, T::Type, dims::Dims)        = Array(T, dims)
 
-_similar(::IndicesStartAt1, a::AbstractArray, T::Type)   = similar(a, T, size(a))
+_similar(::IndicesStartAt1, a::AbstractArray, T::Type)   = similar(a, T, size(SafeIndices(), a))
 _similar(::IndicesBehavior, a::AbstractArray, T::Type)   = similar(a, T, indices(a))
 
 """
@@ -525,7 +555,7 @@ function copy!(::LinearIndexing, dest::AbstractArray, ::LinearSlow, src::Abstrac
 end
 
 function copy!(dest::AbstractArray, dstart::Integer, src::AbstractArray)
-    copy!(dest, dstart, src, first(linearindices(src)), length(src))
+    copy!(dest, dstart, src, first(linearindices(src)), length(SafeIndices(), src))
 end
 
 function copy!(dest::AbstractArray, dstart::Integer, src::AbstractArray, sstart::Integer)
@@ -650,7 +680,7 @@ function _maxlength(A, B, C...)
     max(length(A), _maxlength(B, C...))
 end
 
-isempty(a::AbstractArray) = (length(a) == 0)
+isempty(a::AbstractArray) = (length(SafeIndices(), a) == 0)
 
 ## Conversions ##
 
@@ -1417,7 +1447,7 @@ function mapslices(f, A::AbstractArray, dims::AbstractVector)
         idx[d] = Colon()
     end
 
-    r1 = f(view(A, idx...))
+    r1 = f(A[idx...])
 
     # determine result size and allocate
     Rsize = copy(dimsA)
@@ -1427,7 +1457,7 @@ function mapslices(f, A::AbstractArray, dims::AbstractVector)
     end
     nextra = max(0,length(dims)-ndims(r1))
     if eltype(Rsize) == Int
-        Rsize[dims] = [size(r1)..., ntuple(d->1, nextra)...]
+        Rsize[dims] = [size(SafeIndices(), r1)..., ntuple(d->1, nextra)...]
     else
         Rsize[dims] = [indices(r1)..., ntuple(d->1:1, nextra)...]
     end

base/abstractarraymath.jl

@@ -11,7 +11,7 @@ transpose(a::AbstractArray) = error("transpose not implemented for $(typeof(a)).
 
 ## Constructors ##
 
-vec(a::AbstractArray) = reshape(a,length(a))
+@safeindices vec(a::AbstractArray) = reshape(a,length(a))
 vec(a::AbstractVector) = a
 
 _sub(::Tuple{}, ::Tuple{}) = ()
@@ -69,7 +69,7 @@ function flipdim(A::AbstractVector, d::Integer)
     reverse(A)
 end
 
-function flipdim(A::AbstractArray, d::Integer)
+@safeindices function flipdim(A::AbstractArray, d::Integer)
     nd = ndims(A)
     if d > nd || isempty(A)
         return copy(A)
@@ -107,7 +107,7 @@ function circshift{T,N}(a::AbstractArray{T,N}, shiftamts)
 end
 
 # Uses K-B-N summation
-function cumsum_kbn{T<:AbstractFloat}(v::AbstractVector{T})
+@safeindices function cumsum_kbn{T<:AbstractFloat}(v::AbstractVector{T})
     r = similar(v)
     if isempty(v); return r; end
 

base/arraymath.jl

@@ -271,7 +271,7 @@ function ctranspose!(B::AbstractMatrix, A::AbstractVector)
 end
 
 const transposebaselength=64
-function transpose_f!(f,B::AbstractMatrix,A::AbstractMatrix)
+@safeindices function transpose_f!(f,B::AbstractMatrix,A::AbstractMatrix)
     indices(B,1) == indices(A,2) && indices(B,2) == indices(A,1) || throw(DimensionMismatch(string(f)))
 
     m, n = size(A)

base/broadcast.jl

@@ -64,7 +64,7 @@ end
 @inline _newindex(out, I, keep::Bool, indexmap...) = _newindex((out..., ifelse(keep, I[1], 1)), tail(I), indexmap...)
 
 newindexer(sz, x::Number) = ()
-@inline newindexer(sz, A) = _newindexer(sz, size(A))
+@safeindices @inline newindexer(sz, A) = _newindexer(sz, size(A))
 @inline _newindexer(sz, szA::Tuple{}) = ()
 @inline _newindexer(sz, szA) = (sz[1] == szA[1], _newindexer(tail(sz), tail(szA))...)
 
@@ -132,7 +132,7 @@ end
     end
 end
 
-@inline function broadcast!{nargs}(f, B::AbstractArray, As::Vararg{Any,nargs})
+@safeindices @inline function broadcast!{nargs}(f, B::AbstractArray, As::Vararg{Any,nargs})
     check_broadcast_shape(shape(B), As...)
     sz = size(B)
     mapindex = map(x->newindexer(sz, x), As)
@@ -175,7 +175,7 @@ end
     end
 end
 
-function broadcast_t(f, ::Type{Any}, As...)
+@safeindices function broadcast_t(f, ::Type{Any}, As...)
     shp = broadcast_shape(As...)
     iter = CartesianRange(shp)
     if isempty(iter)
@@ -218,12 +218,12 @@ end
 @inline bitbroadcast(f, As...) = broadcast!(f, allocate_for(BitArray, As, broadcast_shape(As...)), As...)
 
 broadcast_getindex(src::AbstractArray, I::AbstractArray...) = broadcast_getindex!(Array{eltype(src)}(broadcast_shape(I...)), src, I...)
-@generated function broadcast_getindex!(dest::AbstractArray, src::AbstractArray, I::AbstractArray...)
+@safeindices @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(shape(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)
@@ -233,7 +233,7 @@ broadcast_getindex(src::AbstractArray, I::AbstractArray...) = broadcast_getindex
     end
 end
 
-@generated function broadcast_setindex!(A::AbstractArray, x, I::AbstractArray...)
+@safeindices @generated function broadcast_setindex!(A::AbstractArray, x, I::AbstractArray...)
     N = length(I)
     Isplat = Expr[:(I[$d]) for d = 1:N]
     quote

base/exports.jl

@@ -577,6 +577,7 @@ export
     rot180,
     rotl90,
     rotr90,
+    @safeindices,
     searchsorted,
     searchsortedfirst,
     searchsortedlast,

base/multidimensional.jl

@@ -165,7 +165,7 @@ index_lengths(A::AbstractArray, I::Colon) = (length(A),)
 @inline index_lengths(A::AbstractArray, I...) = index_lengths_dim(A, 1, I...)
 index_lengths_dim(A, dim) = ()
 index_lengths_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
-@inline index_lengths_dim(A, dim, ::Colon, i, I...) = (size(A, dim), index_lengths_dim(A, dim+1, i, I...)...)
+@safeindices @inline index_lengths_dim(A, dim, ::Colon, i, I...) = (size(A, dim), index_lengths_dim(A, dim+1, i, I...)...)
 @inline index_lengths_dim(A, dim, ::Real, I...) = (1, index_lengths_dim(A, dim+1, I...)...)
 @inline index_lengths_dim{N}(A, dim, ::CartesianIndex{N}, I...) = (1, index_lengths_dim(A, dim+N, I...)...)
 @inline index_lengths_dim(A, dim, i::AbstractArray, I...) = (length(i), index_lengths_dim(A, dim+1, I...)...)
@@ -178,9 +178,9 @@ index_lengths_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
 # 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::AbstractArray,  I::Colon) = (length(A),)
+index_shape(A::AbstractArray,  I::Colon) = (length(SafeIndices(), 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{N}(A, dim::NTuple{N}, ::Colon) = (trailingsize(SafeIndices(), A, Val{N}),)
 @inline index_shape_dim{T,N}(A::AbstractArray{T,N}, dim::NTuple{N}, ::Colon) = (shape(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...)...)
@@ -190,8 +190,7 @@ index_shape(A::AbstractArray,  I::Colon) = (length(A),)
 @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 decolon(A::AbstractVector, ::Colon) = (indices(A,1),)
-@inline decolon(A::AbstractArray,  ::Colon) = (1:length(A),)
+@inline decolon(A::AbstractArray,  ::Colon) = (linearindices(A),)
 @inline decolon(A::AbstractArray, I...)     = decolon_dim(A, (true,), I...)
 @inline decolon_dim(A::AbstractArray, dim)  = ()
 @inline decolon_dim{T,N}(A::AbstractArray{T,N}, dim::NTuple{N}, ::Colon) = (indices(A, N),)
@@ -231,13 +230,13 @@ end
         @nexprs $N d->(I_d = to_index(I[d]))
         shape = @ncall $N index_shape A I
         dest = similar(A, shape)
-        size(dest) == map(dimlength, shape) || throw_checksize_error(dest, shape)
+        size(SafeIndices(), dest) == map(dimlength, shape) || throw_checksize_error(dest, shape)
         @ncall $N _unsafe_getindex! dest A I
     end
 end
 
 # logical indexing optimization - don't use find (within to_index)
-function _unsafe_getindex(::LinearIndexing, src::AbstractArray, I::AbstractArray{Bool})
+@safeindices function _unsafe_getindex(::LinearIndexing, src::AbstractArray, I::AbstractArray{Bool})
     shape = index_shape(src, I)
     dest = similar(src, shape)
     size(dest) == map(dimlength, shape) || throw_checksize_error(dest, shape)
@@ -254,7 +253,7 @@ function _unsafe_getindex(::LinearIndexing, src::AbstractArray, I::AbstractArray
 end
 
 # specialized form for LinearFast
-function _unsafe_getindex(::LinearFast, src::AbstractArray, I::AbstractArray{Bool})
+@safeindices function _unsafe_getindex(::LinearFast, src::AbstractArray, I::AbstractArray{Bool})
     shape = index_shape(src, I)
     dest = similar(src, shape)
     size(dest) == shape || throw_checksize_error(dest, shape)
@@ -410,22 +409,24 @@ for (f, fmod, op) = ((:cummin, :_cummin!, :min), (:cummax, :_cummax!, :max))
         return res
     end
 
-    @eval function ($f)(A::AbstractArray, axis::Integer)
+    @safeindices @eval function ($f)(A::AbstractArray, axis::Integer)
         res = similar(A)
         if size(A, axis) < 1
             return res
         end
-        R1 = CartesianRange(size(A)[1:axis-1])
-        R2 = CartesianRange(size(A)[axis+1:end])
+        R1 = CartesianRange(indices(A)[1:axis-1])
+        R2 = CartesianRange(indices(A)[axis+1:end])
         ($fmod)(res, A, R1, R2, axis)
     end
 
     @eval @noinline function ($fmod)(res, A::AbstractArray, R1::CartesianRange, R2::CartesianRange, axis::Integer)
+        inds = indices(A, axis)
+        i1 = first(inds)
         for I2 in R2
             for I1 in R1
-                res[I1, 1, I2] = A[I1, 1, I2]
+                res[I1, i1, I2] = A[I1, i1, I2]
             end
-            for i = 2:size(A, axis)
+            for i = i1+1:last(inds)
                 for I1 in R1
                     res[I1, i, I2] = ($op)(A[I1, i, I2], res[I1, i-1, I2])
                 end
@@ -445,7 +446,7 @@ cumprod!(B, A) = cumprod!(B, A, 1)
 cumsum!(B, A, axis::Integer) = cumop!(+, B, A, axis)
 cumprod!(B, A, axis::Integer) = cumop!(*, B, A, axis)
 
-function cumop!(op, B, A, axis::Integer)
+@safeindices function cumop!(op, B, A, axis::Integer)
     if size(B, axis) < 1
         return B
     end