invperm: 0-allocation version of Base.invperm

Author: gustafsson

src/util.jl

@@ -74,10 +74,16 @@ TrivialView(a::AbstractArray{T,N}) where {T,N} = TrivialView{typeof(a),T,N}(a)
 @inline drop_sdims(a::StaticArrayLike) = TrivialView(a)
 @inline drop_sdims(a) = a
 
-Base.@propagate_inbounds function invperm(p::StaticVector)
-    # in difference to base, this does not check if p is a permutation (every value unique)
-     ip = similar(p)
-     ip[p] = 1:length(p)
-     similar_type(p)(ip)
-end
+import Base: invperm
 
+# 0-allocation version of Base.invperm
+@inline function invperm(p::StaticVector{N,T})::SVector{N,T} where {N,T<:Integer}
+    r = MVector{N,T}(undef)
+    r .= 0
+    @inbounds for i in 1:N
+        j = p[i]
+        0 < j <= N && iszero(r[j]) || throw(ArgumentError("p is not a permuation"))
+        r[j] = i
+    end
+    SVector(r)
+end

test/qr.jl

@@ -146,6 +146,30 @@ end
             @test_throws Union{BoundsError,ArgumentError} invperm(x)
         end
 
+        @testset "0 allocations" begin
+            (function()
+                x = @MVector zeros(5)
+                X = @MMatrix zeros(5,5)
+
+                           @allocated sum(invperm(F.p))
+                @test 0 == @allocated sum(invperm(F.p))
+                           @allocated sum(A * y)
+                @test 0 == @allocated sum(A * y)
+                           @allocated sum(F \ y)
+                @test 0 == @allocated sum(F \ y)
+                           @allocated sum(ldiv!(x, F, y))
+                @test 0 == @allocated sum(ldiv!(x, F, y))
+
+                # "SMAtrix * SMatrix" causes allocations, so these all allocate > 0
+                #         @allocated sum(A' * Y)
+                #a_mult = @allocated sum(A' * Y)
+                #               @allocated sum(F \ Y)
+                #@test a_mult < @allocated sum(F \ Y) < 2a_mult
+                #                   @allocated sum(ldiv!(X, F, Y))
+                #@test a_mult+32 == @allocated sum(ldiv!(X, F, Y))
+            end)()
+        end
+
         @testset "10x faster" begin
             time_to_test = @elapsed (function()
                 y2 = @SVector rand(50)