faster & type stable invperm and isperm for tuples < 16 elements (#…

…35234) * Make `invperm` much faster, and type stable for tuples Co-authored-by: Takafumi Arakaki <[email protected]>
JuliaLang · Apr 8, 2020 · b1ceadc · b1ceadc
1 parent 4efb718
commit b1ceadc
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 2 deletions.
diff --git a/base/combinatorics.jl b/base/combinatorics.jl
@@ -36,6 +36,22 @@ end
 
 # Basic functions for working with permutations
 
+@inline function _foldoneto(op, acc, ::Val{N}) where N
+    @assert N::Integer > 0
+    if @generated
+        quote
+            acc_0 = acc
+            Base.Cartesian.@nexprs $N i -> acc_{i} = op(acc_{i-1}, i)
+            return $(Symbol(:acc_, N))
+        end
+    else
+        for i in 1:N
+            acc = op(acc, i)
+        end
+        return acc
+    end
+end
+
 """
     isperm(v) -> Bool
 
@@ -50,7 +66,9 @@ julia> isperm([1; 3])
 false
 ```
 """
-function isperm(A)
+isperm(A) = _isperm(A)
+
+function _isperm(A)
     n = length(A)
     used = falses(n)
     for a in A
@@ -63,6 +81,18 @@ isperm(p::Tuple{}) = true
 isperm(p::Tuple{Int}) = p[1] == 1
 isperm(p::Tuple{Int,Int}) = ((p[1] == 1) & (p[2] == 2)) | ((p[1] == 2) & (p[2] == 1))
 
+function isperm(P::Tuple)
+    valn = Val(length(P))
+    _foldoneto(true, valn) do b,i
+        s = _foldoneto(false, valn) do s, j
+            s || P[j]==i
+        end
+        b&s
+    end
+end
+
+isperm(P::Any16) = _isperm(P)
+
 # swap columns i and j of a, in-place
 function swapcols!(a::AbstractMatrix, i, j)
     i == j && return
@@ -242,7 +272,21 @@ function invperm(p::Union{Tuple{},Tuple{Int},Tuple{Int,Int}})
     isperm(p) || throw(ArgumentError("argument is not a permutation"))
     p  # in dimensions 0-2, every permutation is its own inverse
 end
-invperm(a::Tuple) = (invperm([a...])...,)
+
+function invperm(P::Tuple)
+    valn = Val(length(P))
+    ntuple(valn) do i
+        s = _foldoneto(nothing, valn) do s, j
+            s !== nothing && return s
+            P[j]==i && return j
+            nothing
+        end
+        s === nothing && throw(ArgumentError("argument is not a permutation"))
+        s
+    end
+end
+
+invperm(P::Any16) = Tuple(invperm(collect(P)))
 
 #XXX This function should be moved to Combinatorics.jl but is currently used by Base.DSP.
 """

diff --git a/test/combinatorics.jl b/test/combinatorics.jl
@@ -45,6 +45,17 @@ end
     let ai = 2:-1:1
         @test invpermute!(permute!([1, 2], ai), ai) == [1, 2]
     end
+
+    # PR 35234
+    for N in 3:1:20
+        A=randcycle(N)
+        T=Tuple(A)
+        K=Tuple(A.-1)
+        @test A[collect(invperm(T))] == 1:N
+        @test_throws ArgumentError invperm(K)
+        @test isperm(T) == true
+        @test isperm(K) == false
+    end
 end
 
 @testset "factorial" begin