Fix some edge cases for mangling (#672)

Author: Keluaa

src/mangling.jl

@@ -29,31 +29,33 @@ safe_name(x) = safe_name(repr(x))
 # we generate function names that look like C++ functions, because many tools, like NVIDIA's
 # profilers, support them (grouping different instantiations of the same kernel together).
 
-function mangle_param(t, substitutions=Any[])
+function mangle_param(t, substitutions = Any[], top = false)
     t == Nothing && return "v"
 
     function find_substitution(x)
         sub = findfirst(isequal(x), substitutions)
-        if sub === nothing
+        res = if sub === nothing
             nothing
         elseif sub == 1
             "S_"
         else
             seq_id = uppercase(string(sub-2; base=36))
             "S$(seq_id)_"
         end
+        return res
+    end
+
+    # check if we already know this type
+    str = find_substitution(t)
+    if str !== nothing
+        return str
     end
 
     if isa(t, DataType) && t <: Ptr
         tn = mangle_param(eltype(t), substitutions)
+        push!(substitutions, t)
         "P$tn"
     elseif isa(t, DataType)
-        # check if we already know this type
-        str = find_substitution(t)
-        if str !== nothing
-            return str
-        end
-
         # check if we already know this base type
         str = find_substitution(t.name.wrapper)
         if str === nothing
@@ -62,47 +64,66 @@ function mangle_param(t, substitutions=Any[])
             push!(substitutions, t.name.wrapper)
         end
 
-        # encode typevars as template parameters
-        if !isempty(t.parameters)
-            str *= "I"
-            for t in t.parameters
-                str *= mangle_param(t, substitutions)
+        if t.name.wrapper == t && !isa(t.name.wrapper, UnionAll)
+            # a type with no typevars
+            str
+        else
+            # encode typevars as template parameters
+            if isempty(t.parameters)
+                w_types = t.name.wrapper.types
+                if !isempty(w_types) && !isempty(w_types[end] isa Core.TypeofVararg)
+                    # If the type accepts a variable amount of parameters,
+                    # e.g. `Tuple{}`, then we mark it as empty: "Tuple<>"
+                    str *= "IJEE"
+                end
+            else
+                str *= "I"
+                for tp in t.parameters
+                    str *= mangle_param(tp, substitutions)
+                end
+                str *= "E"
             end
-            str *= "E"
-
             push!(substitutions, t)
+            str
         end
-
-        str
     elseif isa(t, Union)
-        # check if we already know this union type
-        str = find_substitution(t)
-        if str !== nothing
-            return str
-        end
-
         # check if we already know the Union name
         str = find_substitution(Union)
         if str === nothing
             tn = "Union"
             str = "$(length(tn))$tn"
-            push!(substitutions, tn)
+            push!(substitutions, Union)
         end
 
         # encode union types as template parameters
-        if !isempty(Base.uniontypes(t))
-            str *= "I"
-            for t in Base.uniontypes(t)
-                str *= mangle_param(t, substitutions)
-            end
-            str *= "E"
-
-            push!(substitutions, t)
+        str *= "I"
+        for tp in Base.uniontypes(t)  # cannot be empty as `Union{}` is not a `Union`
+            str *= mangle_param(tp, substitutions)
         end
+        str *= "E"
 
+        push!(substitutions, t)
         str
     elseif isa(t, UnionAll)
-        mangle_param(t.body, substitutions)
+        mangle_param(Base.unwrap_unionall(t), substitutions)
+    elseif isa(t, Core.TypeofVararg)
+        T = isdefined(t, :T) ? t.T : Any
+        if isdefined(t, :N)
+            # For NTuple, repeat the type as needed
+            str = ""
+            for _ in 1:t.N
+                str *= mangle_param(T, substitutions)
+            end
+            str
+        elseif top
+            # Variadic arguments only make sense for function arguments
+            mangle_param(T, substitutions) * "z"  # T...
+        else
+            # Treat variadic arguments for a type as no arguments
+            ""
+        end
+    elseif isa(t, Char)
+        mangle_param(UInt32(t), substitutions)
     elseif isa(t, Union{Bool, Cchar, Cuchar, Cshort, Cushort, Cint, Cuint, Clong, Culong, Clonglong, Culonglong, Int128, UInt128})
         ts = t isa Bool       ? 'b' : # bool
              t isa Cchar      ? 'a' : # signed char
@@ -153,7 +174,7 @@ function mangle_sig(sig)
     # mangle each parameter
     substitutions = []
     for t in tt
-        str *= mangle_param(t, substitutions)
+        str *= mangle_param(t, substitutions, true)
     end
 
     return str

test/utils.jl

@@ -29,19 +29,39 @@ end
     @test mangle(identity, Val{Cshort(1)}) == "identity(Val<(short)1>)"
     @test mangle(identity, Val{1.0}) == "identity(Val<0x1p+0>)"
     @test mangle(identity, Val{1f0}) == "identity(Val<0x1p+0f>)"
+    @test mangle(identity, Val{'a'}) == "identity(Val<97u>)"
+    @test mangle(identity, Val{'∅'}) == "identity(Val<8709u>)"
 
     # unions
     @test mangle(identity, Union{Int32, Int64}) == "identity(Union<Int32, Int64>)"
+    @test mangle(identity, Union, Int, Union{Int64, Int32}) == "identity(Union, Int64, Union<Int32, Int64>)"
 
     # union alls
     @test mangle(identity, Array) == "identity(Array<T, N>)"
+    @test mangle(identity, Tuple) == "identity(Tuple)"
+    @test mangle(identity, Vector) == "identity(Array<T, 1>)"
+    @test mangle(identity, NTuple{2, I} where {I <: Integer}) == "identity(Tuple<I__Integer, I__Integer>)"
+
+    # Vararg
+    @test mangle(identity, Vararg{Int}) == "identity(Int64, ...)"
+    @test mangle(identity, Vararg{Int, 2}) == "identity(Int64, Int64)"
+    @test mangle(identity, Tuple{1, 2}, Tuple{}, Tuple) == "identity(Tuple<1, 2>, Tuple<>, Tuple)"
+    @test mangle(identity, NTuple{2, Int}) == "identity(Tuple<Int64, Int64>)"
+    @test mangle(identity, Tuple{Vararg{Int}}) == "identity(Tuple<>)"
 
     # many substitutions
     @test mangle(identity, Val{1}, Val{2}, Val{3}, Val{4}, Val{5}, Val{6}, Val{7}, Val{8},
                            Val{9}, Val{10}, Val{11}, Val{12}, Val{13}, Val{14}, Val{15},
                            Val{16}, Val{16}) ==
           "identity(Val<1>, Val<2>, Val<3>, Val<4>, Val<5>, Val<6>, Val<7>, Val<8>, Val<9>, Val<10>, Val<11>, Val<12>, Val<13>, Val<14>, Val<15>, Val<16>, Val<16>)"
 
+    # intertwined substitutions
+    @test mangle(
+        identity, Val{1}, Ptr{Tuple{Ptr{Int}}}, Ptr{Int}, Val{1}, Val{2},
+        Tuple{Ptr{Int}}, Tuple{Int8}, Int64, Int8
+    ) ==
+        "identity(Val<1>, Tuple<Int64*>*, Int64*, Val<1>, Val<2>, Tuple<Int64*>, Tuple<Int8>, Int64, Int8)"
+
     # problematic examples
     @test mangle(identity, String, Matrix{Float32}, Broadcast.Broadcasted{Broadcast.ArrayStyle{Matrix{Float32}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, typeof(Base.literal_pow), Tuple{Base.RefValue{typeof(sin)}, Broadcast.Extruded{Matrix{Float32}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}}}) == "identity(String, Array<Float32, 2>, Broadcasted<ArrayStyle<Array<Float32, 2>>, Tuple<OneTo<Int64>, OneTo<Int64>>, literal_pow, Tuple<RefValue<sin>, Extruded<Array<Float32, 2>, Tuple<Bool, Bool>, Tuple<Int64, Int64>>>>)"
 end