more powerful reverse(A; dims) (#37367)

* more powerful reverse(A; dims)

* fix reverse for offsetarray

* bugfix for odd dimensions + test

* typo

* fix bootstrap failure

* explicit reverse! and dims=: tests

* NEWS

* more offsetarray tests

* compat annotation

* further simplify and speed up reverse (and reverse! is now down to 0 allocations)

* rm redundant inbounds

* whoops extra space

* use ÷ 2 instead of >> 1

* fix type output in doc

* fix type instability in slice ranges

Author: stevengj

NEWS.md

@@ -74,6 +74,9 @@ New library features
 Standard library changes
 ------------------------
 * The `nextprod` function now accepts tuples and other array types for its first argument ([#35791]).
+* The `reverse(A; dims)` function for multidimensional `A` can now reverse multiple dimensions at once
+  by passing a tuple for `dims`, and defaults to reversing all dimensions; there is also a multidimensional
+  in-place `reverse!(A; dims)` ([#37367]).
 * The function `isapprox(x,y)` now accepts the `norm` keyword argument also for numeric (i.e., non-array) arguments `x` and `y` ([#35883]).
 * `view`, `@view`, and `@views` now work on `AbstractString`s, returning a `SubString` when appropriate ([#35879]).
 * All `AbstractUnitRange{<:Integer}`s now work with `SubString`, `view`, `@view` and `@views` on strings ([#35879]).

base/abstractarraymath.jl

@@ -127,56 +127,6 @@ julia> selectdim(A, 2, 3)
     return view(A, idxs...)
 end
 
-"""
-    reverse(A; dims::Integer)
-
-Reverse `A` in dimension `dims`.
-
-# Examples
-```jldoctest
-julia> b = [1 2; 3 4]
-2×2 Matrix{Int64}:
- 1  2
- 3  4
-
-julia> reverse(b, dims=2)
-2×2 Matrix{Int64}:
- 2  1
- 4  3
-```
-"""
-function reverse(A::AbstractArray; dims::Integer)
-    nd = ndims(A); d = dims
-    1 ≤ d ≤ nd || throw(ArgumentError("dimension $d is not 1 ≤ $d ≤ $nd"))
-    if isempty(A)
-        return copy(A)
-    elseif nd == 1
-        return reverse(A)
-    end
-    inds = axes(A)
-    B = similar(A)
-    nnd = 0
-    for i = 1:nd
-        nnd += Int(length(inds[i])==1 || i==d)
-    end
-    indsd = inds[d]
-    sd = first(indsd)+last(indsd)
-    if nnd==nd
-        # reverse along the only non-singleton dimension
-        for i in indsd
-            B[i] = A[sd-i]
-        end
-        return B
-    end
-    let B=B # workaround #15276
-        alli = [ axes(B,n) for n in 1:nd ]
-        for i in indsd
-            B[[ n==d ? sd-i : alli[n] for n in 1:nd ]...] = selectdim(A, d, i)
-        end
-    end
-    return B
-end
-
 function circshift(a::AbstractArray, shiftamt::Real)
     circshift!(similar(a), a, (Integer(shiftamt),))
 end

base/array.jl

@@ -1589,22 +1589,33 @@ julia> reverse(A, 3, 5)
  3
 ```
 """
-function reverse(A::AbstractVector, s=first(LinearIndices(A)), n=last(LinearIndices(A)))
+function reverse(A::AbstractVector, start::Integer, stop::Integer=lastindex(A))
+    s, n = Int(start), Int(stop)
     B = similar(A)
-    for i = first(LinearIndices(A)):s-1
+    for i = firstindex(A):s-1
         B[i] = A[i]
     end
     for i = s:n
         B[i] = A[n+s-i]
     end
-    for i = n+1:last(LinearIndices(A))
+    for i = n+1:lastindex(A)
         B[i] = A[i]
     end
     return B
 end
 
-# to resolve ambiguity with reverse(A; dims)
-reverse(A::Vector) = invoke(reverse, Tuple{AbstractVector}, A)
+# 1d special cases of reverse(A; dims) and reverse!(A; dims):
+for (f,_f) in ((:reverse,:_reverse), (:reverse!,:_reverse!))
+    @eval begin
+        $f(A::AbstractVector; dims=:) = $_f(A, dims)
+        $_f(A::AbstractVector, ::Colon) = $f(A, firstindex(A), lastindex(A))
+        $_f(A::AbstractVector, dim::Tuple{Integer}) = $_f(A, first(dim))
+        function $_f(A::AbstractVector, dim::Integer)
+            dim == 1 || throw(ArgumentError("invalid dimension $dim ≠ 1"))
+            return $_f(A, :)
+        end
+    end
+end
 
 function reverseind(a::AbstractVector, i::Integer)
     li = LinearIndices(a)
@@ -1637,7 +1648,8 @@ julia> A
  1
 ```
 """
-function reverse!(v::AbstractVector, s=first(LinearIndices(v)), n=last(LinearIndices(v)))
+function reverse!(v::AbstractVector, start::Integer, stop::Integer=lastindex(v))
+    s, n = Int(start), Int(stop)
     liv = LinearIndices(v)
     if n <= s  # empty case; ok
     elseif !(first(liv) ≤ s ≤ last(liv))

base/arraymath.jl

@@ -58,69 +58,78 @@ end
 
 ## data movement ##
 
-reverse(A::Array; dims::Integer) = _reverse_int(A, Int(dims))
-function _reverse_int(A::Array{T}, d::Int) where T
-    nd = ndims(A)
-    1 ≤ d ≤ nd || throw(ArgumentError("dimension $d is not 1 ≤ $d ≤ $nd"))
-    sd = size(A, d)
-    if sd == 1 || isempty(A)
-        return copy(A)
-    end
+"""
+    reverse(A; dims=:)
 
-    B = similar(A)
+Reverse `A` along dimension `dims`, which can be an integer (a
+single dimension), a tuple of integers (a tuple of dimensions)
+or `:` (reverse along all the dimensions, the default).  See
+also [`reverse!`](@ref) for in-place reversal.
 
-    nnd = 0
-    for i = 1:nd
-        nnd += size(A,i)==1 || i==d
-    end
-    if nnd==nd
-        # reverse along the only non-singleton dimension
-        for i = 1:sd
-            B[i] = A[sd+1-i]
-        end
-        return B
+# Examples
+```jldoctest
+julia> b = Int64[1 2; 3 4]
+2×2 Matrix{Int64}:
+ 1  2
+ 3  4
+
+julia> reverse(b, dims=2)
+2×2 Matrix{Int64}:
+ 2  1
+ 4  3
+
+julia> reverse(b)
+2×2 Matrix{Int64}:
+ 4  3
+ 2  1
+```
+
+!!! compat "Julia 1.6"
+    Prior to Julia 1.6, only single-integer `dims` are supported in `reverse`.
+"""
+reverse(A::AbstractArray; dims=:) = _reverse(A, dims)
+_reverse(A, dims) = reverse!(copymutable(A); dims)
+
+"""
+    reverse!(A; dims=:)
+
+Like [`reverse`](@ref), but operates in-place in `A`.
+
+!!! compat "Julia 1.6"
+    Multidimensional `reverse!` requires Julia 1.6.
+"""
+reverse!(A::AbstractArray; dims=:) = _reverse!(A, dims)
+_reverse!(A::AbstractArray{<:Any,N}, ::Colon) where {N} = _reverse!(A, ntuple(identity, Val{N}()))
+_reverse!(A, dim::Integer) = _reverse!(A, (Int(dim),))
+_reverse!(A, dims::NTuple{M,Integer}) where {M} = _reverse!(A, Int.(dims))
+function _reverse!(A::AbstractArray{<:Any,N}, dims::NTuple{M,Int}) where {N,M}
+    dimrev = ntuple(k -> k in dims, Val{N}()) # boolean tuple indicating reversed dims
+
+    if N < M || M != sum(dimrev)
+        throw(ArgumentError("invalid dimensions $dims in reverse!"))
     end
+    M == 0 && return A # nothing to reverse
 
-    d_in = size(A)
-    M = 1
-    for i = 1:d-1
-        M *= d_in[i]
+    # swapping loop only needs to traverse ≈half of the array
+    halfsz = ntuple(k -> k == dims[1] ? size(A,k) ÷ 2 : size(A,k), Val{N}())
+
+    last1 = ntuple(k -> lastindex(A,k)+firstindex(A,k), Val{N}()) # offset for reversed index
+    for i in CartesianIndices(ntuple(k -> firstindex(A,k):firstindex(A,k)-1+@inbounds(halfsz[k]), Val{N}()))
+        iₜ = Tuple(i)
+        iᵣ = CartesianIndex(ifelse.(dimrev, last1 .- iₜ, iₜ))
+        @inbounds A[iᵣ], A[i] = A[i], A[iᵣ]
     end
-    N = length(A)
-    stride = M * sd
-
-    if M==1
-        for j = 0:stride:(N-stride)
-            for i = 1:sd
-                ri = sd+1-i
-                B[j + ri] = A[j + i]
-            end
-        end
-    else
-        if allocatedinline(T) && M>200
-            for i = 1:sd
-                ri = sd+1-i
-                for j=0:stride:(N-stride)
-                    offs = j + 1 + (i-1)*M
-                    boffs = j + 1 + (ri-1)*M
-                    copyto!(B, boffs, A, offs, M)
-                end
-            end
-        else
-            for i = 1:sd
-                ri = sd+1-i
-                for j=0:stride:(N-stride)
-                    offs = j + 1 + (i-1)*M
-                    boffs = j + 1 + (ri-1)*M
-                    for k=0:(M-1)
-                        B[boffs + k] = A[offs + k]
-                    end
-                end
-            end
-        end
+    if M > 1 && isodd(size(A, dims[1]))
+        # middle slice for odd dimensions must be recursively flipped
+        mid = firstindex(A, dims[1]) + (size(A, dims[1]) ÷ 2)
+        midslice = CartesianIndices(ntuple(k -> k == dims[1] ? (mid:mid) : (firstindex(A,k):lastindex(A,k)), Val{N}()))
+        _reverse!(view(A, midslice), dims[2:end])
     end
-    return B
+    return A
 end
+# fix ambiguity with array.jl:
+_reverse!(A::AbstractVector, dim::Tuple{Int}) = _reverse!(A, first(dim))
+
 
 """
     rotl90(A)

base/bitarray.jl

@@ -1163,8 +1163,8 @@ end
 
 # TODO some of this could be optimized
 
-reverse(A::BitArray; dims::Integer) = _reverse_int(A, Int(dims))
-function _reverse_int(A::BitArray, d::Int)
+_reverse(A::BitArray, d::Tuple{Integer}) = _reverse(A, d[1])
+function _reverse(A::BitArray, d::Int)
     nd = ndims(A)
     1 ≤ d ≤ nd || throw(ArgumentError("dimension $d is not 1 ≤ $d ≤ $nd"))
     sd = size(A, d)
@@ -1210,7 +1210,7 @@ function _reverse_int(A::BitArray, d::Int)
     return B
 end
 
-function reverse!(B::BitVector)
+function _reverse!(B::BitVector, ::Colon)
     # Basic idea: each chunk is divided into two blocks of size k = n % 64, and
     # h = 64 - k. Walk from either end (with indices i and j) reversing chunks
     # and separately ORing their two blocks into place.
@@ -1264,9 +1264,6 @@ function reverse!(B::BitVector)
     return B
 end
 
-reverse(v::BitVector) = reverse!(copy(v))
-
-
 function (<<)(B::BitVector, i::UInt)
     n = length(B)
     i == 0 && return copy(B)

base/range.jl

@@ -990,15 +990,15 @@ end
 Array{T,1}(r::AbstractRange{T}) where {T} = vcat(r)
 collect(r::AbstractRange) = vcat(r)
 
-reverse(r::OrdinalRange) = (:)(last(r), -step(r), first(r))
-function reverse(r::StepRangeLen)
+_reverse(r::OrdinalRange, ::Colon) = (:)(last(r), -step(r), first(r))
+function _reverse(r::StepRangeLen, ::Colon)
     # If `r` is empty, `length(r) - r.offset + 1 will be nonpositive hence
     # invalid. As `reverse(r)` is also empty, any offset would work so we keep
     # `r.offset`
     offset = isempty(r) ? r.offset : length(r)-r.offset+1
     StepRangeLen(r.ref, -r.step, length(r), offset)
 end
-reverse(r::LinRange{T}) where {T} = LinRange{T}(r.stop, r.start, length(r))
+_reverse(r::LinRange{T}, ::Colon) where {T} = LinRange{T}(r.stop, r.start, length(r))
 
 ## sorting ##
 

test/arrayops.jl

@@ -1555,6 +1555,28 @@ end
     @test reverse(["a" "b"; "c" "d"], dims = 2) == ["b" "a"; "d" "c"]
 end
 
+@testset "reverse multiple dims" begin
+    for A in (zeros(2,4), zeros(3,5))
+        A[:] .= 1:length(A) # unique-ify elements
+        @test reverse(A) == reverse!(reverse(A, dims=1), dims=2) ==
+              reverse(A, dims=(1,2)) == reverse(A, dims=(2,1))
+        @test_throws ArgumentError reverse(A, dims=(1,2,3))
+        @test_throws ArgumentError reverse(A, dims=(1,2,2))
+    end
+    for A in (zeros(2,4,6), zeros(3,5,7))
+        A[:] .= 1:length(A) # unique-ify elements
+        @test reverse(A) == reverse(A, dims=:) == reverse!(copy(A)) == reverse!(copy(A), dims=:) ==
+              reverse!(reverse!(reverse(A, dims=1), dims=2), dims=3) ==
+              reverse!(reverse(A, dims=(1,2)), dims=3) ==
+              reverse!(reverse(A, dims=(2,3)), dims=1) ==
+              reverse!(reverse(A, dims=(1,3)), dims=2) ==
+              reverse(A, dims=(1,2,3)) == reverse(A, dims=(3,2,1)) == reverse(A, dims=(2,1,3))
+        @test reverse(A, dims=(1,2)) == reverse!(reverse(A, dims=1), dims=2)
+        @test reverse(A, dims=(1,3)) == reverse!(reverse(A, dims=1), dims=3)
+        @test reverse(A, dims=(2,3)) == reverse!(reverse(A, dims=2), dims=3)
+    end
+end
+
 @testset "isdiag, istril, istriu" begin
     @test isdiag(3)
     @test istril(4)

test/offsetarray.jl

@@ -522,6 +522,12 @@ soa = OffsetArray([2,2,3], typemax(Int)-4)
 @test rotr90(A) == OffsetArray(rotr90(parent(A)), A.offsets[[2,1]])
 @test reverse(A, dims=1) == OffsetArray(reverse(parent(A), dims=1), A.offsets)
 @test reverse(A, dims=2) == OffsetArray(reverse(parent(A), dims=2), A.offsets)
+@test reverse(A) == reverse!(reverse(A, dims=1), dims=2)
+
+Aodd = OffsetArray(rand(3,5), (-3,5))
+@test reverse(Aodd, dims=1) == OffsetArray(reverse(parent(Aodd), dims=1), Aodd.offsets)
+@test reverse(Aodd, dims=2) == OffsetArray(reverse(parent(Aodd), dims=2), Aodd.offsets)
+@test reverse(Aodd) == reverse!(reverse(Aodd, dims=1), dims=2)
 
 @test A .+ 1 == OffsetArray(parent(A) .+ 1, A.offsets)
 @test 2*A == OffsetArray(2*parent(A), A.offsets)