Refactor SubArray indexing to use Varargs

mbauman · mbauman · commit 7f2d8454807d · 2016-05-07T19:06:12.000-04:00
diff --git a/base/multidimensional.jl b/base/multidimensional.jl
@@ -458,95 +458,6 @@ for (f, fmod, op) = ((:cummin, :_cummin!, :min), (:cummax, :_cummax!, :max))
     @eval ($f)(A::AbstractArray) = ($f)(A, 1)
 end
 
-## SubArray index merging
-# A view created like V = A[2:3:8, 5:2:17] can later be indexed as V[2:7],
-# creating a new 1d view.
-# In such cases we have to collapse the 2d space spanned by the ranges.
-#
-# API:
-#    merge_indexes(V, indexes::NTuple, index)
-# indexes encodes the view's trailing indexes into the parent array,
-# and index encodes the subset of these elements that we'll select.
-#
-# It returns a CartesianIndex or array of CartesianIndexes.
-
-# Checking 'in' a range is fast -- so check all possibilities and keep the good ones
-@generated function merge_indexes{N}(V, indexes::NTuple{N}, index::Union{Colon, Range})
-    # There may be a vector of cartesian indices in the passed indexes... which
-    # makes the number of indices more than N. Since we pre-allocate the array
-    # of CartesianIndexes, we need to figure out how big to make it
-    M = 0
-    for T in indexes.parameters
-        T <: CartesianIndex ? (M += length(T)) : (M += 1)
-    end
-    index_length_expr = index <: Colon ? Symbol("Istride_", N+1) : :(length(index))
-    quote
-        Cartesian.@nexprs $N d->(I_d = indexes[d])
-        dimlengths = Cartesian.@ncall $N index_lengths_dim V.parent length(V.indexes)-N+1 I
-        Istride_1 = 1   # strides of the indexes to merge
-        Cartesian.@nexprs $N d->(Istride_{d+1} = Istride_d*dimlengths[d])
-        idx_len = $(index_length_expr)
-        if idx_len < 0.1*$(Symbol("Istride_", N+1))   # this has not been carefully tuned
-            return merge_indexes_div(V, indexes, index, dimlengths)
-        end
-        Cartesian.@nexprs $N d->(counter_d = 1) # counter_0 is the linear index
-        k = 0
-        merged = Array(CartesianIndex{$M}, idx_len)
-        Cartesian.@nloops $N i d->(1:dimlengths[d]) d->(counter_{d-1} = counter_d + (i_d-1)*Istride_d; @inbounds idx_d = I_d[i_d]) begin
-            if counter_0 in index # this branch is elided for ::Colon
-                @inbounds merged[k+=1] = Cartesian.@ncall $N CartesianIndex{$M} idx
-            end
-        end
-        merged
-    end
-end
-
-# mapping getindex across the parent and subindices rapidly gets too big to
-# automatically inline, but it is crucial that it does so to avoid allocations
-# Unlike SubArray's reindex, merge_indexes doesn't drop any indices.
-@inline inlinemap(f, t::Tuple, s::Tuple) = (f(t[1], s[1]), inlinemap(f, tail(t), tail(s))...)
-inlinemap(f, t::Tuple{}, s::Tuple{}) = ()
-inlinemap(f, t::Tuple{}, s::Tuple) = ()
-inlinemap(f, t::Tuple, s::Tuple{}) = ()
-
-# Otherwise, we fall back to the slow div/rem method, using ind2sub.
-@inline merge_indexes{N}(V, indexes::NTuple{N}, index) =
-    merge_indexes_div(V, indexes, index, index_lengths_dim(V.parent, length(V.indexes)-N+1, indexes...))
-
-@inline merge_indexes_div{N}(V, indexes::NTuple{N}, index::Real, dimlengths) =
-    CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, index)))
-merge_indexes_div{N}(V, indexes::NTuple{N}, index::AbstractArray, dimlengths) =
-    reshape([CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, i))) for i in index], size(index))
-merge_indexes_div{N}(V, indexes::NTuple{N}, index::Colon, dimlengths) =
-    [CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, i))) for i in 1:prod(dimlengths)]
-
-# Merging indices is particularly difficult in the case where we partially linearly
-# index through a multidimensional array. It's easiest if we can simply reduce the
-# partial indices to a single linear index into the parent index array.
-function merge_indexes{N}(V, indexes::NTuple{N}, index::Tuple{Colon, Vararg{Colon}})
-    shape = index_shape(indexes[1], index...)
-    reshape(merge_indexes(V, indexes, :), (shape[1:end-1]..., shape[end]*prod(index_lengths_dim(V.parent, length(V.indexes)-length(indexes)+2, tail(indexes)...))))
-end
-@inline merge_indexes{N}(V, indexes::NTuple{N}, index::Tuple{Real, Vararg{Real}}) = merge_indexes(V, indexes, sub2ind(size(indexes[1]), index...))
-# In general, it's a little trickier, but we can use the product iterator
-# if we replace colons with ranges. This can be optimized further.
-function merge_indexes{N}(V, indexes::NTuple{N}, index::Tuple)
-    I = replace_colons(V, indexes, index)
-    shp = index_shape(indexes[1], I...) # index_shape does no bounds checking
-    dimlengths = index_lengths_dim(V.parent, length(V.indexes)-N+1, indexes...)
-    sz = size(indexes[1])
-    reshape([CartesianIndex(inlinemap(getindex, indexes, ind2sub(dimlengths, sub2ind(sz, i...)))) for i in product(I...)], shp)
-end
-@inline replace_colons(V, indexes, I) = replace_colons_dim(V, indexes, 1, I)
-@inline replace_colons_dim(V, indexes, dim, I::Tuple{}) = ()
-@inline replace_colons_dim(V, indexes, dim, I::Tuple{Colon}) =
-    (1:trailingsize(indexes[1], dim)*prod(index_lengths_dim(V.parent, length(V.indexes)-length(indexes)+2, tail(indexes)...)),)
-@inline replace_colons_dim(V, indexes, dim, I::Tuple{Colon, Vararg{Any}}) =
-    (1:size(indexes[1], dim), replace_colons_dim(V, indexes, dim+1, tail(I))...)
-@inline replace_colons_dim(V, indexes, dim, I::Tuple{Any, Vararg{Any}}) =
-    (I[1], replace_colons_dim(V, indexes, dim+1, tail(I))...)
-
-
  cumsum(A::AbstractArray, axis::Integer=1) =  cumsum!(similar(A, Base._cumsum_type(A)), A, axis)
 cumsum!(B, A::AbstractArray) = cumsum!(B, A, 1)
 cumprod(A::AbstractArray, axis::Integer=1) = cumprod!(similar(A), A, axis)
diff --git a/base/reshapedarray.jl b/base/reshapedarray.jl
@@ -38,6 +38,7 @@ reshape(parent::AbstractArray, dims::Dims) = _reshape(parent, dims)
 reshape(parent::AbstractArray, len::Integer) = reshape(parent, (Int(len),))
 reshape(parent::AbstractArray, dims::Int...) = reshape(parent, dims)
 
+reshape{T,N}(parent::AbstractArray{T,N}, ndims::Type{Val{N}}) = parent
 @generated function reshape{T,AN,N}(parent::AbstractArray{T,AN}, ndims::Type{Val{N}})
     if AN < N
         # Add trailing singleton dimensions
diff --git a/base/sharedarray.jl b/base/sharedarray.jl
@@ -19,7 +19,7 @@ type SharedArray{T,N} <: DenseArray{T,N}
     # the local partition into the array when viewed as a single dimensional array.
     # this can be removed when @parallel or its equivalent supports looping on
     # a subset of workers.
-    loc_subarr_1d::SubArray{T,1,Array{T,N},Tuple{UnitRange{Int}},true}
+    loc_subarr_1d::SubArray{T,1,Array{T,1},Tuple{UnitRange{Int}},true}
 
     SharedArray(d,p,r,sn) = new(d,p,r,sn)
 end
diff --git a/base/subarray.jl b/base/subarray.jl
@@ -77,9 +77,23 @@ parentindexes(a::AbstractArray) = ntuple(i->1:size(a,i), ndims(a))
 
 ## SubArray creation
 # Drops singleton dimensions (those indexed with a scalar)
-function slice(A::AbstractArray, I::ViewIndex...)
+function slice{T,N}(A::AbstractArray{T,N}, I::Vararg{ViewIndex,N})
     @_inline_meta
     @boundscheck checkbounds(A, I...)
+    unsafe_slice(A, I...)
+end
+function slice(A::AbstractArray, i::ViewIndex)
+    @_inline_meta
+    @boundscheck checkbounds(A, i)
+    unsafe_slice(reshape(A, Val{1}), i)
+end
+function slice{N}(A::AbstractArray, I::Vararg{ViewIndex,N}) # TODO: DEPRECATE FOR #14770
+    @_inline_meta
+    @boundscheck checkbounds(A, I...)
+    unsafe_slice(reshape(A, Val{N}), I...)
+end
+function unsafe_slice{T,N}(A::AbstractArray{T,N}, I::Vararg{ViewIndex,N})
+    @_inline_meta
     J = to_indexes(I...)
     SubArray(A, J, index_shape(A, J...))
 end
@@ -89,9 +103,23 @@ keep_leading_scalars(T::Tuple{Real, Vararg{Real}}) = T
 keep_leading_scalars(T::Tuple{Real, Vararg{Any}}) = (@_inline_meta; (NoSlice(T[1]), keep_leading_scalars(tail(T))...))
 keep_leading_scalars(T::Tuple{Any, Vararg{Any}}) = (@_inline_meta; (T[1], keep_leading_scalars(tail(T))...))
 
-function sub(A::AbstractArray, I::ViewIndex...)
+function sub{T,N}(A::AbstractArray{T,N}, I::Vararg{ViewIndex,N})
     @_inline_meta
     @boundscheck checkbounds(A, I...)
+    unsafe_sub(A, I...)
+end
+function sub(A::AbstractArray, i::ViewIndex)
+    @_inline_meta
+    @boundscheck checkbounds(A, i)
+    unsafe_sub(reshape(A, Val{1}), i)
+end
+function sub{N}(A::AbstractArray, I::Vararg{ViewIndex,N}) # TODO: DEPRECATE FOR #14770
+    @_inline_meta
+    @boundscheck checkbounds(A, I...)
+    unsafe_sub(reshape(A, Val{N}), I...)
+end
+function unsafe_sub{T,N}(A::AbstractArray{T,N}, I::Vararg{ViewIndex,N})
+    @_inline_meta
     J = keep_leading_scalars(to_indexes(I...))
     SubArray(A, J, index_shape(A, J...))
 end
@@ -100,90 +128,61 @@ end
 #
 # Recursively look through the heads of the parent- and sub-indexes, considering
 # the following cases:
-# * Parent index is empty  -> ignore trailing scalars, but preserve added dimensions
-# * Parent index is Any    -> re-index that with the sub-index
-# * Parent index is Scalar -> that dimension was dropped, so skip the sub-index and use the index as is
+# * Parent index is array  -> re-index that with one or more sub-indexes (one per dimension)
+# * Parent index is Colon  -> just use the sub-index as provided
+# * Parent index is scalar -> that dimension was dropped, so skip the sub-index and use the index as is
 #
-# Furthermore, we must specially consider the case with one final sub-index,
-# as it may be a linear index that spans multiple parent indexes.
 
+typealias AbstractZeroDimArray{T} AbstractArray{T, 0}
 typealias DroppedScalar Union{Real, AbstractCartesianIndex}
-# When indexing beyond the parent indices, drop all trailing scalars (they must be 1 to be inbounds)
-reindex(V, idxs::Tuple{}, subidxs::Tuple{Vararg{DroppedScalar}}) = ()
-# Drop any intervening scalars that are beyond the parent indices but before a nonscalar
-reindex(V, idxs::Tuple{}, subidxs::Tuple{DroppedScalar, Vararg{Any}}) =
-    (@_propagate_inbounds_meta; (reindex(V, idxs, tail(subidxs))...))
-# And keep the nonscalar index to add the dimension
-reindex(V, idxs::Tuple{}, subidxs::Tuple{Any, Vararg{Any}}) =
-    (@_propagate_inbounds_meta; (subidxs[1], reindex(V, idxs, tail(subidxs))...))
+
+reindex(V, ::Tuple{}, ::Tuple{}) = ()
 
 # Skip dropped scalars, so simply peel them off the parent indices and continue
 reindex(V, idxs::Tuple{DroppedScalar, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) =
     (@_propagate_inbounds_meta; (idxs[1], reindex(V, tail(idxs), subidxs)...))
 
 # Colons simply pass their subindexes straight through
-reindex(V, idxs::Tuple{Colon}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (subidxs[1],))
 reindex(V, idxs::Tuple{Colon, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) =
     (@_propagate_inbounds_meta; (subidxs[1], reindex(V, tail(idxs), tail(subidxs))...))
-reindex(V, idxs::Tuple{Colon, Vararg{Any}}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (merge_indexes(V, idxs, subidxs[1]),))
-# As an optimization, we don't need to merge indices if all trailing indices are dropped scalars
-reindex(V, idxs::Tuple{Colon, Vararg{DroppedScalar}}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (subidxs[1], tail(idxs)...))
 
 # Re-index into parent vectors with one subindex
-reindex(V, idxs::Tuple{AbstractVector}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (idxs[1][subidxs[1]],))
 reindex(V, idxs::Tuple{AbstractVector, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) =
     (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], reindex(V, tail(idxs), tail(subidxs))...))
-reindex(V, idxs::Tuple{AbstractVector, Vararg{Any}}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (merge_indexes(V, idxs, subidxs[1]),))
-reindex(V, idxs::Tuple{AbstractVector, Vararg{DroppedScalar}}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], tail(idxs)...))
 
 # Parent matrices are re-indexed with two sub-indices
-reindex(V, idxs::Tuple{AbstractMatrix}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (idxs[1][subidxs[1]],))
-reindex(V, idxs::Tuple{AbstractMatrix}, subidxs::Tuple{Any, Any}) =
-    (@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]],))
 reindex(V, idxs::Tuple{AbstractMatrix, Vararg{Any}}, subidxs::Tuple{Any, Any, Vararg{Any}}) =
     (@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]], reindex(V, tail(idxs), tail(tail(subidxs)))...))
-reindex(V, idxs::Tuple{AbstractMatrix, Vararg{Any}}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (merge_indexes(V, idxs, subidxs[1]),))
-reindex(V, idxs::Tuple{AbstractMatrix, Vararg{Any}}, subidxs::Tuple{Any, Any}) =
-    (@_propagate_inbounds_meta; (merge_indexes(V, idxs, subidxs),))
-reindex(V, idxs::Tuple{AbstractMatrix, Vararg{DroppedScalar}}, subidxs::Tuple{Any}) =
-    (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], tail(idxs)...))
-reindex(V, idxs::Tuple{AbstractMatrix, Vararg{DroppedScalar}}, subidxs::Tuple{Any, Any}) =
-    (@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]], tail(idxs)...))
 
 # In general, we index N-dimensional parent arrays with N indices
 @generated function reindex{T,N}(V, idxs::Tuple{AbstractArray{T,N}, Vararg{Any}}, subidxs::Tuple{Vararg{Any}})
     if length(subidxs.parameters) >= N
         subs = [:(subidxs[$d]) for d in 1:N]
         tail = [:(subidxs[$d]) for d in N+1:length(subidxs.parameters)]
         :(@_propagate_inbounds_meta; (idxs[1][$(subs...)], reindex(V, tail(idxs), ($(tail...),))...))
-    elseif length(idxs.parameters) == 1
-        :(@_propagate_inbounds_meta; (idxs[1][subidxs...],))
-    elseif all(T->T<:DroppedScalar, idxs.parameters[2:end])
-        :(@_propagate_inbounds_meta; (idxs[1][subidxs...], tail(idxs)...))
     else
-        :(@_propagate_inbounds_meta; (merge_indexes(V, idxs, subidxs),))
+        :(error("mismatched index dimensionalities"))
     end
 end
 
 # In general, we simply re-index the parent indices by the provided ones
-getindex(V::SubArray) = (@_propagate_inbounds_meta; getindex(V, 1))
-function getindex(V::SubArray, I::Real...)
+typealias SlowSubArray{T,N,P,I} SubArray{T,N,P,I,false}
+function getindex{T,N}(V::SlowSubArray{T,N}, I::Vararg{Real,N})
     @_inline_meta
     @boundscheck checkbounds(V, I...)
     @inbounds r = V.parent[reindex(V, V.indexes, to_indexes(I...))...]
     r
 end
+# Explicitly define scalar linear indexing -- this is needed so the nonscalar
+# indexing methods don't take precedence here
+function getindex(V::SlowSubArray, i::Real)
+    @_inline_meta
+    @boundscheck checkbounds(V, i)
+    @inbounds r = V.parent[reindex(V, V.indexes, ind2sub(size(V), to_index(i)))...]
+    r
+end
 
 typealias FastSubArray{T,N,P,I} SubArray{T,N,P,I,true}
-getindex(V::FastSubArray) = (@_propagate_inbounds_meta; getindex(V, 1))
 function getindex(V::FastSubArray, i::Real)
     @_inline_meta
     @boundscheck checkbounds(V, i)
@@ -198,40 +197,55 @@ function getindex(V::FastContiguousSubArray, i::Real)
     @inbounds r = V.parent[V.first_index + to_index(i)-1]
     r
 end
-# We need this because the ::ViewIndex... method would otherwise obscure the Base fallback
-function getindex(V::FastSubArray, I::Real...)
+# Just like the slow case, explicitly define scalar indexing at N dims, too
+function getindex{T,N}(V::FastSubArray{T,N}, I::Vararg{Real,N})
     @_inline_meta
     @boundscheck checkbounds(V, I...)
     @inbounds r = getindex(V, sub2ind(size(V), to_indexes(I...)...))
     r
 end
-getindex{T,N}(V::SubArray{T,N}, I::ViewIndex...) = (@_propagate_inbounds_meta; copy(slice(V, I...)))
 
-setindex!(V::SubArray, x) = (@_propagate_inbounds_meta; setindex!(V, x, 1))
-function setindex!{T,N}(V::SubArray{T,N}, x, I::Real...)
+# Nonscalar indexing just copies a view
+getindex{T,N}(V::SubArray{T,N}, i::ViewIndex) = (@_propagate_inbounds_meta; copy(slice(V, i)))
+getindex{T,N}(V::SubArray{T,N}, I::Vararg{ViewIndex,N}) = (@_propagate_inbounds_meta; copy(slice(V, I...)))
+
+# Setindex is similar, but since we don't specially define non-scalar methods
+# we only need to define the canonical methods. We don't need to worry about
+# e.g., linear indexing for SlowSubArray since the fallbacks can do their thing
+function setindex!{T,N}(V::SlowSubArray{T,N}, x, I::Vararg{Real,N})
     @_inline_meta
     @boundscheck checkbounds(V, I...)
     @inbounds V.parent[reindex(V, V.indexes, to_indexes(I...))...] = x
     V
 end
+function setindex!(V::FastSubArray, x, i::Real)
+    @_inline_meta
+    @boundscheck checkbounds(V, i)
+    @inbounds V.parent[V.first_index + V.stride1*(to_index(i)-1)] = x
+    V
+end
+function setindex!(V::FastContiguousSubArray, x, i::Real)
+    @_inline_meta
+    @boundscheck checkbounds(V, i)
+    @inbounds V.parent[V.first_index + to_index(i)-1] = x
+    V
+end
 # Nonscalar setindex! falls back to the defaults
 
-function slice{T,N}(V::SubArray{T,N}, I::ViewIndex...)
+function unsafe_slice{T,N}(V::SubArray{T,N}, I::Vararg{ViewIndex,N})
     @_inline_meta
-    @boundscheck checkbounds(V, I...)
     idxs = reindex(V, V.indexes, to_indexes(I...))
     SubArray(V.parent, idxs, index_shape(V.parent, idxs...))
 end
 
-function sub{T,N}(V::SubArray{T,N}, I::ViewIndex...)
+function unsafe_sub{T,N}(V::SubArray{T,N}, I::Vararg{ViewIndex,N})
     @_inline_meta
-    @boundscheck checkbounds(V, I...)
     idxs = reindex(V, V.indexes, keep_leading_scalars(to_indexes(I...)))
     SubArray(V.parent, idxs, index_shape(V.parent, idxs...))
 end
 
-linearindexing(A::FastSubArray) = LinearFast()
-linearindexing(A::SubArray) = LinearSlow()
+linearindexing{T<:FastSubArray}(::Type{T}) = LinearFast()
+linearindexing{T<:SubArray}(::Type{T}) = LinearSlow()
 
 getindex(::Colon, i) = to_index(i)
 unsafe_getindex(::Colon, i) = to_index(i)
