Faster, indices-aware circshift (and non-allocating circshift!)

Fixes #16032, fixes #17581

Author: timholy

base/abstractarraymath.jl

@@ -158,8 +158,10 @@ function flipdim(A::AbstractArray, d::Integer)
     return B
 end
 
-circshift(a::AbstractArray, shiftamt::Real) = circshift(a, [Integer(shiftamt)])
-
+function circshift(a::AbstractArray, shiftamt::Real)
+    circshift!(similar(a), a, (Integer(shiftamt),))
+end
+circshift(a::AbstractArray, shiftamt::DimsInteger) = circshift!(similar(a), a, shiftamt)
 """
     circshift(A, shifts)
 
@@ -174,29 +176,25 @@ julia> b = reshape(collect(1:16), (4,4))
  3  7  11  15
  4  8  12  16
 
-julia> circshift(b, [0,2])
+julia> circshift(b, (0,2))
 4×4 Array{Int64,2}:
   9  13  1  5
  10  14  2  6
  11  15  3  7
  12  16  4  8
 
-julia> circshift(b, [-1,0])
+julia> circshift(b, (-1,0))
 4×4 Array{Int64,2}:
  2  6  10  14
  3  7  11  15
  4  8  12  16
  1  5   9  13
 ```
+
+See also `circshift!`.
 """
-function circshift{T,N}(a::AbstractArray{T,N}, shiftamts)
-    I = ()
-    for i=1:N
-        s = size(a,i)
-        d = i<=length(shiftamts) ? shiftamts[i] : 0
-        I = tuple(I..., d==0 ? [1:s;] : mod([-d:s-1-d;], s).+1)
-    end
-    a[(I::NTuple{N,Vector{Int}})...]
+function circshift(a::AbstractArray, shiftamt)
+    circshift!(similar(a), a, map(Integer, (shiftamt...,)))
 end
 
 # Uses K-B-N summation

base/exports.jl

@@ -490,6 +490,7 @@ export
     checkbounds,
     checkindex,
     circshift,
+    circshift!,
     clamp!,
     colon,
     conj!,

base/multidimensional.jl

@@ -609,6 +609,54 @@ function copy!{T,N}(dest::AbstractArray{T,N}, src::AbstractArray{T,N})
     dest
 end
 
+function copy!(dest::AbstractArray, Rdest::CartesianRange, src::AbstractArray, Rsrc::CartesianRange)
+    isempty(Rdest) && return dest
+    size(Rdest) == size(Rsrc) || throw(ArgumentError("source and destination must have same size (got $(size(Rsrc)) and $(size(Rdest)))"))
+    @boundscheck checkbounds(dest, Rdest.start)
+    @boundscheck checkbounds(dest, Rdest.stop)
+    @boundscheck checkbounds(src, Rsrc.start)
+    @boundscheck checkbounds(src, Rsrc.stop)
+    deltaI = Rdest.start - Rsrc.start
+    for I in Rsrc
+        @inbounds dest[I+deltaI] = src[I]
+    end
+    dest
+end
+
+# circshift!
+circshift!(dest::AbstractArray, src, ::Tuple{}) = copy!(dest, src)
+"""
+    circshift!(dest, src, shifts)
+
+Circularly shift the data in `src`, storing the result in
+`dest`. `shifts` specifies the amount to shift in each dimension.
+
+See also `circshift`.
+"""
+@noinline function circshift!{T,N}(dest::AbstractArray{T,N}, src, shiftamt::DimsInteger)
+    dest === src && throw(ArgumentError("dest and src must be separate arrays"))
+    inds = indices(src)
+    indices(dest) == inds || throw(ArgumentError("indices of src and dest must match (got $inds and $(indices(dest)))"))
+    _circshift!(dest, (), src, (), inds, fill_to_length(shiftamt, 0, Val{N}))
+end
+circshift!(dest::AbstractArray, src, shiftamt) = circshift!(dest, src, (shiftamt...,))
+
+@inline function _circshift!(dest, rdest, src, rsrc,
+                             inds::Tuple{AbstractUnitRange,Vararg{Any}},
+                             shiftamt::Tuple{Integer,Vararg{Any}})
+    ind1, d = inds[1], shiftamt[1]
+    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)
+    tinds, tshiftamt = tail(inds), tail(shiftamt)
+    _circshift!(dest, (rdest..., r1), src, (rsrc..., r4), tinds, tshiftamt)
+    _circshift!(dest, (rdest..., r2), src, (rsrc..., r3), tinds, tshiftamt)
+end
+# At least one of inds, shiftamt is empty
+function _circshift!(dest, rdest, src, rsrc, inds, shiftamt)
+    copy!(dest, CartesianRange(rdest), src, CartesianRange(rsrc))
+end
 
 ### BitArrays
 

doc/stdlib/arrays.rst

@@ -574,20 +574,30 @@ Indexing, Assignment, and Concatenation
         3  7  11  15
         4  8  12  16
 
-       julia> circshift(b, [0,2])
+       julia> circshift(b, (0,2))
        4×4 Array{Int64,2}:
          9  13  1  5
         10  14  2  6
         11  15  3  7
         12  16  4  8
 
-       julia> circshift(b, [-1,0])
+       julia> circshift(b, (-1,0))
        4×4 Array{Int64,2}:
         2  6  10  14
         3  7  11  15
         4  8  12  16
         1  5   9  13
 
+   See also ``circshift!``\ .
+
+.. function:: circshift!(dest, src, shifts)
+
+   .. Docstring generated from Julia source
+
+   Circularly shift the data in ``src``\ , storing the result in ``dest``\ . ``shifts`` specifies the amount to shift in each dimension.
+
+   See also ``circshift``\ .
+
 .. function:: find(A)
 
    .. Docstring generated from Julia source

test/arrayops.jl

@@ -457,6 +457,15 @@ for i = tensors
     @test isequal(i,permutedims(ipermutedims(i,perm),perm))
 end
 
+## circshift
+
+@test circshift(1:5, -1) == circshift(1:5, 4) == circshift(1:5, -6) == [2,3,4,5,1]
+@test circshift(1:5, 1) == circshift(1:5, -4) == circshift(1:5, 6)  == [5,1,2,3,4]
+a = [1:5;]
+@test_throws ArgumentError Base.circshift!(a, a, 1)
+b = copy(a)
+@test Base.circshift!(b, a, 1) == [5,1,2,3,4]
+
 ## unique across dim ##
 
 # All rows and columns unique

test/offsetarray.jl

@@ -411,6 +411,7 @@ v = OffsetArray(rand(8), (-2,))
 @test rotr90(A) == OffsetArray(rotr90(parent(A)), A.offsets[[2,1]])
 @test flipdim(A, 1) == OffsetArray(flipdim(parent(A), 1), A.offsets)
 @test flipdim(A, 2) == OffsetArray(flipdim(parent(A), 2), A.offsets)
+@test circshift(A, (-1,2)) == OffsetArray(circshift(parent(A), (-1,2)), A.offsets)
 
 @test A+1 == OffsetArray(parent(A)+1, A.offsets)
 @test 2*A == OffsetArray(2*parent(A), A.offsets)