Adds methods for getindex and setindex! with Colon

Adds optimized methods for getindex and setindex for Array, BitArray, SubArray with ::Colon. Adds a simple sparse matrix fallback methods to translate Colons into ranges Internal method changes: * Augments Base.index_shape to support ::Colon * Adds a Base.index_lengths that returns the lengths of each index (without dropping trailing scalars) * Base.setindex_shape_check now takes the index lengths (instead of the indices themselves)
pwl · Jun 3, 2015 · a07c6d3 · a07c6d3
1 parent ce9abe5
commit a07c6d3
Show file tree

Hide file tree

Showing 10 changed files with 132 additions and 59 deletions.
diff --git a/base/abstractarray.jl b/base/abstractarray.jl
@@ -119,12 +119,12 @@ linearindexing{T<:Range}(::Type{T}) = LinearFast()
 *(::LinearSlow, ::LinearSlow) = LinearSlow()
 
 ## Bounds checking ##
-checkbounds(sz::Int, ::Colon) = nothing
 checkbounds(sz::Int, i::Int) = 1 <= i <= sz || throw(BoundsError())
 checkbounds(sz::Int, i::Real) = checkbounds(sz, to_index(i))
 checkbounds(sz::Int, I::AbstractVector{Bool}) = length(I) == sz || throw(BoundsError())
 checkbounds(sz::Int, r::Range{Int}) = isempty(r) || (minimum(r) >= 1 && maximum(r) <= sz) || throw(BoundsError())
 checkbounds{T<:Real}(sz::Int, r::Range{T}) = checkbounds(sz, to_index(r))
+checkbounds(sc::Int, ::Colon) = true
 
 function checkbounds{T <: Real}(sz::Int, I::AbstractArray{T})
  for i in I
@@ -136,17 +136,17 @@ checkbounds(A::AbstractArray, I::AbstractArray{Bool}) = size(A) == size(I) || th
 
 checkbounds(A::AbstractArray, I) = checkbounds(length(A), I)
 
-function checkbounds(A::AbstractMatrix, I::Union(Real,Colon,AbstractArray), J::Union(Real,Colon,AbstractArray))
+function checkbounds(A::AbstractMatrix, I::Union(Real,AbstractArray,Colon), J::Union(Real,AbstractArray,Colon))
  checkbounds(size(A,1), I)
  checkbounds(size(A,2), J)
 end
 
-function checkbounds(A::AbstractArray, I::Union(Real,Colon,AbstractArray), J::Union(Real,Colon,AbstractArray))
+function checkbounds(A::AbstractArray, I::Union(Real,AbstractArray,Colon), J::Union(Real,AbstractArray,Colon))
  checkbounds(size(A,1), I)
  checkbounds(trailingsize(A,2), J)
 end
 
-function checkbounds(A::AbstractArray, I::Union(Real,Colon,AbstractArray)...)
+function checkbounds(A::AbstractArray, I::Union(Real,AbstractArray,Colon)...)
  n = length(I)
  if n > 0
  for dim = 1:(n-1)

diff --git a/base/array.jl b/base/array.jl
@@ -330,6 +330,9 @@ function getindex(A::Range, I::AbstractVector{Bool})
  return [ A[i] for i in to_index(I) ]
 end
 
+function getindex(A::Array, ::Colon)
+ return [ a for a in A ]
+end
 
 # logical indexing
 # (when the indexing is provided as an Array{Bool} or a BitArray we can be
@@ -409,6 +412,7 @@ function setindex!{T<:Real}(A::Array, X::AbstractArray, I::AbstractVector{T})
  return A
 end
 
+setindex!(A::Array, x, I::Colon) = setindex!(A, x, 1:length(A))
 
 # logical indexing
 # (when the indexing is provided as an Array{Bool} or a BitArray we can be

diff --git a/base/broadcast.jl b/base/broadcast.jl
@@ -273,7 +273,7 @@ end
  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)
- Base.setindex_shape_check(X, (@ntuple $N fakeI)...)
+ @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)

diff --git a/base/multidimensional.jl b/base/multidimensional.jl
@@ -231,10 +231,13 @@ using .IteratorsMD
 
 @generated function checksize(A::AbstractArray, I...)
  N = length(I)
- quote
- @nexprs $N d->(size(A, d) == length(I[d]) || throw(DimensionMismatch("index $d has length $(length(I[d])), but size(A, $d) = $(size(A,d))")))
- nothing
+ ex = Expr(:block)
+ push!(ex.args, :(idxlens = index_lengths(A, I...)))
+ for d=1:N
+ push!(ex.args, :(size(A, $d) == idxlens[$d] || throw(DimensionMismatch("index ", $d, " has length ", idxlens[$d], ", but size(A, ", $d, ") = ", size(A,$d)))))
  end
+ push!(ex.args, :(nothing))
+ ex
 end
 
 @inline unsafe_getindex(v::BitArray, ind::Int) = Base.unsafe_bitgetindex(v.chunks, ind)
@@ -245,7 +248,7 @@ end
 @inline unsafe_setindex!{T}(v::AbstractArray{T}, x::T, ind::Real) = unsafe_setindex!(v, x, to_index(ind))
 
 # Version that uses cartesian indexing for src
-@generated function _getindex!(dest::Array, src::AbstractArray, I::Union(Int,AbstractVector)...)
+@generated function _getindex!(dest::Array, src::AbstractArray, I::Union(Int,AbstractVector,Colon)...)
  N = length(I)
  Isplat = Expr[:(I[$d]) for d = 1:N]
  quote
@@ -260,7 +263,7 @@ end
 end
 
 # Version that uses linear indexing for src
-@generated function _getindex!(dest::Array, src::Array, I::Union(Int,AbstractVector)...)
+@generated function _getindex!(dest::Array, src::Array, I::Union(Int,AbstractVector,Colon)...)
  N = length(I)
  Isplat = Expr[:(I[$d]) for d = 1:N]
  quote
@@ -279,12 +282,12 @@ end
 
 # It's most efficient to call checkbounds first, then to_index, and finally
 # allocate the output. Hence the different variants.
-_getindex(A, I::Tuple{Vararg{Union(Int,AbstractVector),}}) =
- _getindex!(similar(A, index_shape(I...)), A, I...)
+_getindex(A, I::Tuple{Vararg{Union(Int,AbstractVector,Colon),}}) =
+ _getindex!(similar(A, index_shape(A, I...)), A, I...)
 
 # The @generated function here is just to work around the performance hit
 # of splatting
-@generated function getindex(A::Array, I::Union(Real,AbstractVector)...)
+@generated function getindex(A::Array, I::Union(Real,AbstractVector,Colon)...)
  N = length(I)
  Isplat = Expr[:(I[$d]) for d = 1:N]
  quote
@@ -294,7 +297,7 @@ _getindex(A, I::Tuple{Vararg{Union(Int,AbstractVector),}}) =
 end
 
 # Also a safe version of getindex!
-@generated function getindex!(dest, src, I::Union(Real,AbstractVector)...)
+@generated function getindex!(dest, src, I::Union(Real,AbstractVector,Colon)...)
  N = length(I)
  Isplat = Expr[:(I[$d]) for d = 1:N]
  Jsplat = Expr[:(to_index(I[$d])) for d = 1:N]
@@ -305,21 +308,23 @@ end
 end
 
 
-@generated function setindex!(A::Array, x, J::Union(Real,AbstractArray)...)
+@generated function setindex!(A::Array, x, J::Union(Real,AbstractArray,Colon)...)
  N = length(J)
  if x<:AbstractArray
  ex=quote
  X = x
- @ncall $N setindex_shape_check X I
+ idxlens = @ncall $N index_lengths A I
+ setindex_shape_check(X, idxlens...)
  Xs = start(X)
- @nloops $N i d->(1:length(I_d)) d->(@inbounds offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
+ @nloops $N i d->(1:idxlens[d]) d->(@inbounds offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
  v, Xs = next(X, Xs)
  @inbounds A[offset_0] = v
  end
  end
  else
  ex=quote
- @nloops $N i d->(1:length(I_d)) d->(@inbounds offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
+ idxlens = @ncall $N index_lengths A I
+ @nloops $N i d->(1:idxlens[d]) d->(@inbounds offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
  @inbounds A[offset_0] = x
  end
  end
@@ -355,20 +360,23 @@ end
 
 ### subarray.jl
 
+# This is the code-generation block for SubArray's staged setindex! function:
+# _setindex!(V::SubArray, x, J::Union(Real,AbstractVector,Colon)...)
 function gen_setindex_body(N::Int)
  quote
  Base.Cartesian.@nexprs $N d->(J_d = J[d])
  Base.Cartesian.@ncall $N checkbounds V J
  Base.Cartesian.@nexprs $N d->(I_d = Base.to_index(J_d))
+ idxlens = @ncall $N index_lengths V I
  if !isa(x, AbstractArray)
- Base.Cartesian.@nloops $N i d->(1:length(I_d)) d->(@inbounds j_d = Base.unsafe_getindex(I_d, i_d)) begin
+ Base.Cartesian.@nloops $N i d->(1:idxlens[d]) d->(@inbounds j_d = Base.unsafe_getindex(I_d, i_d)) begin
  @inbounds (Base.Cartesian.@nref $N V j) = x
  end
  else
  X = x
- Base.Cartesian.@ncall $N Base.setindex_shape_check X I
+ setindex_shape_check(X, idxlens...)
  k = 1
- Base.Cartesian.@nloops $N i d->(1:length(I_d)) d->(@inbounds j_d = Base.unsafe_getindex(I_d, i_d)) begin
+ Base.Cartesian.@nloops $N i d->(1:idxlens[d]) d->(@inbounds j_d = Base.unsafe_getindex(I_d, i_d)) begin
  @inbounds (Base.Cartesian.@nref $N V j) = X[k]
  k += 1
  end
@@ -438,6 +446,26 @@ function merge_indexes(V, parentindexes::NTuple, parentdims::Dims, linindex, lin
  merge_indexes_div(V, parentindexes, parentdims, linindex, lindim)
 end
 
+# Even simpler is the case where the linear index is ::Colon: return all indexes
+@generated function merge_indexes(V, indexes::NTuple, dims::Dims, ::Colon)
+ N = length(indexes)
+ N > 0 || throw(ArgumentError("cannot merge empty indexes"))
+ quote
+ Base.Cartesian.@nexprs $N d->(I_d = indexes[d])
+ dimoffset = ndims(V.parent) - length(dims)
+ n = prod(map(length, indexes))
+ Pstride_1 = 1 # parent strides
+ Base.Cartesian.@nexprs $(N-1) d->(Pstride_{d+1} = Pstride_d*dims[d])
+ Base.Cartesian.@nexprs $N d->(offset_d = 1) # offset_0 is a linear index into parent
+ k = 0
+ index = Array(Int, n)
+ Base.Cartesian.@nloops $N i d->(1:dimsize(V, d+dimoffset, I_d)) d->(offset_{d-1} = offset_d + (I_d[i_d]-1)*Pstride_d) begin
+ index[k+=1] = offset_0
+ end
+ index
+ end
+end
+
 # This could be written as a regular function, but performance
 # will be better using Cartesian macros to avoid the heap and
 # an extra loop.
@@ -599,20 +627,27 @@ function getindex(B::BitArray, I0::UnitRange{Int})
  return unsafe_getindex(B, I0)
 end
 
+function getindex(B::BitArray, ::Colon)
+ X = BitArray(0)
+ X.chunks = copy(B.chunks)
+ X.len = length(B)
+ return X
+end
+
 getindex{T<:Real}(B::BitArray, I0::UnitRange{T}) = getindex(B, to_index(I0))
 
-@generated function unsafe_getindex(B::BitArray, I0::UnitRange{Int}, I::Union(Int,UnitRange{Int})...)
+@generated function unsafe_getindex(B::BitArray, I0::Union(Colon,UnitRange{Int}), I::Union(Int,UnitRange{Int},Colon)...)
  N = length(I)
  Isplat = Expr[:(I[$d]) for d = 1:N]
  quote
  @nexprs $N d->(I_d = I[d])
- X = BitArray(index_shape(I0, $(Isplat...)))
+ X = BitArray(index_shape(B, I0, $(Isplat...)))
 
  f0 = first(I0)
- l0 = length(I0)
+ l0 = size(X, 1)
 
  gap_lst_1 = 0
- @nexprs $N d->(gap_lst_{d+1} = length(I_d))
+ @nexprs $N d->(gap_lst_{d+1} = size(X, d+1))
  stride = 1
  ind = f0
  @nexprs $N d->begin
@@ -636,20 +671,21 @@ end
 
 # general multidimensional non-scalar indexing
 
-@generated function unsafe_getindex(B::BitArray, I::Union(Int,AbstractVector{Int})...)
+@generated function unsafe_getindex(B::BitArray, I::Union(Int,AbstractVector{Int},Colon)...)
  N = length(I)
  Isplat = Expr[:(I[$d]) for d = 1:N]
  quote
  @nexprs $N d->(I_d = I[d])
- X = BitArray(index_shape($(Isplat...)))
+ shape = @ncall $N index_shape B I
+ X = BitArray(shape)
  Xc = X.chunks
 
  stride_1 = 1
  @nexprs $N d->(stride_{d+1} = stride_d * size(B, d))
  @nexprs 1 d->(offset_{$N} = 1)
  ind = 1
- @nloops($N, i, d->I_d,
- d->(offset_{d-1} = offset_d + (i_d-1)*stride_d), # PRE
+ @nloops($N, i, X, d->(@inbounds j_d = unsafe_getindex(I[d], i_d);
+  offset_{d-1} = offset_d + (j_d-1)*stride_d), # PRE
  begin
  unsafe_bitsetindex!(Xc, B[offset_0], ind)
  ind += 1
@@ -660,7 +696,7 @@ end
 
 # general version with Real (or logical) indexing which dispatches on the appropriate method
 
-@generated function getindex(B::BitArray, I::Union(Real,AbstractVector)...)
+@generated function getindex(B::BitArray, I::Union(Real,AbstractVector,Colon)...)
  N = length(I)
  Isplat = Expr[:(I[$d]) for d = 1:N]
  Jsplat = Expr[:(to_index(I[$d])) for d = 1:N]
@@ -790,25 +826,27 @@ end
 
 # general multidimensional non-scalar indexing
 
-@generated function unsafe_setindex!(B::BitArray, X::AbstractArray, I::Union(Int,AbstractArray{Int})...)
+@generated function unsafe_setindex!(B::BitArray, X::AbstractArray, I::Union(Int,AbstractArray{Int},Colon)...)
  N = length(I)
  quote
  refind = 1
  @nexprs $N d->(I_d = I[d])
- @nloops $N i d->I_d @inbounds begin
- @ncall $N unsafe_setindex! B convert(Bool,X[refind]) i
+ idxlens = @ncall $N index_lengths B I
+ @nloops $N i d->(1:idxlens[d]) d->(J_d = I_d[i_d]) @inbounds begin
+ @ncall $N unsafe_setindex! B convert(Bool,X[refind]) J
  refind += 1
  end
  return B
  end
 end
 
-@generated function unsafe_setindex!(B::BitArray, x::Bool, I::Union(Int,AbstractArray{Int})...)
+@generated function unsafe_setindex!(B::BitArray, x::Bool, I::Union(Int,AbstractArray{Int},Colon)...)
  N = length(I)
  quote
  @nexprs $N d->(I_d = I[d])
- @nloops $N i d->I_d begin
- @ncall $N unsafe_setindex! B x i
+ idxlens = @ncall $N index_lengths B I
+ @nloops $N i d->(1:idxlens[d]) d->(J_d = I_d[i_d]) begin
+ @ncall $N unsafe_setindex! B x J
  end
  return B
  end
@@ -822,7 +860,7 @@ function setindex!(B::BitArray, x, i::Real)
  return unsafe_setindex!(B, convert(Bool,x), to_index(i))
 end
 
-@generated function setindex!(B::BitArray, x, I::Union(Real,AbstractArray)...)
+@generated function setindex!(B::BitArray, x, I::Union(Real,AbstractArray,Colon)...)
  N = length(I)
  quote
  checkbounds(B, I...)
@@ -837,16 +875,17 @@ end
 function setindex!(B::BitArray, X::AbstractArray, i::Real)
  checkbounds(B, i)
  j = to_index(i)
- setindex_shape_check(X, j)
+ setindex_shape_check(X, index_lengths(A, j)[1])
  return unsafe_setindex!(B, X, j)
 end
 
-@generated function setindex!(B::BitArray, X::AbstractArray, I::Union(Real,AbstractArray)...)
+@generated function setindex!(B::BitArray, X::AbstractArray, I::Union(Real,AbstractArray,Colon)...)
  N = length(I)
  quote
  checkbounds(B, I...)
  @nexprs $N d->(J_d = to_index(I[d]))
- @ncall $N setindex_shape_check X J
+ idxlens = @ncall $N index_lengths B J
+ setindex_shape_check(X, idxlens...)
  return @ncall $N unsafe_setindex! B X J
  end
 end