Fix calls to Varargs functions

Author: mhauru

src/copyable_task.jl

@@ -63,6 +63,31 @@ function callable_ret_type(sig, produce_types)
     return Union{Base.code_ircode_by_type(sig)[1][2],produce_type}
 end
 
+"""
+    check_varargs(sig, ir)
+
+For a call signature `sig` and the IR code for it, check whether this is a Varargs call.
+
+The need for this arises because the output of `Base.code_ircode_by_type` does not
+distinguish between varargs and tuples, see https://github.com/JuliaLang/julia/issues/58753. Thus we have to go back to the signature that created the IR to check. There are two cases
+that signal that this is indeed a varargs call:
+1. The last argument in `sig` is a `Vararg` object.
+2. The last argument of `ir` is a `Tuple` of the types of the last arguments of `sig`. For
+instance `sig` may end in `Symbol, Tuple{Int, Int}` and the last argument of `ir` would be
+`Tuple{Symbol, Tuple{Int, Int}}`.
+
+That there are these two cases, and only these two cases, is not based on a good
+understanding of anything, but rather on observing which cases arise in our test suite. This solution is thus a hack and should be rewritten by someone who actually understands how IR
+handles `Varargs`.
+"""
+function check_varargs(sig, ir)
+    sig.parameters[end] isa Core.TypeofVararg && return true
+    (ir.argtypes[end] isa Type && ir.argtypes[end] <: Tuple) || return false
+    ir_last_arg_types = ir.argtypes[end].parameters
+    sig_last_arg_types = sig.parameters[(end - length(ir_last_arg_types) + 1):end]
+    return sig_last_arg_types == ir_last_arg_types
+end
+
 """
     build_callable(sig::Type{<:Tuple})
 
@@ -84,11 +109,12 @@ function build_callable(sig::Type{<:Tuple})
         return fresh_copy(mc_cache[key])
     else
         ir = Base.code_ircode_by_type(sig)[1][1]
+        isva = check_varargs(sig, ir)
         bb, refs, types = derive_copyable_task_ir(BBCode(ir))
         unoptimised_ir = IRCode(bb)
         optimised_ir = optimise_ir!(unoptimised_ir)
         mc_ret_type = callable_ret_type(sig, types)
-        mc = misty_closure(mc_ret_type, optimised_ir, refs...; do_compile=true)
+        mc = misty_closure(mc_ret_type, optimised_ir, refs...; isva=isva, do_compile=true)
         mc_cache[key] = mc
         return mc, refs[end]
     end
@@ -315,7 +341,7 @@ end
 """
     set_taped_globals!(t::TapedTask, new_taped_globals)::Nothing
 
-Set the `taped_globals` of `t` to `new_taped_globals`. Any calls to 
+Set the `taped_globals` of `t` to `new_taped_globals`. Any calls to
 [`get_taped_globals`](@ref) in future calls to `consume(t)` (either directly, or implicitly
 via iteration) will see this new value.
 """
@@ -573,7 +599,7 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple,Vector{Any}}
     # We enforced above the condition that the final statement in a basic block must not
     # produce. This ensures that the final split does not produce. While not strictly
     # necessary, this simplifies the implementation (see below).
-    # 
+    #
     # As a result of the above, a basic block will be associated to exactly one split if it
     # does not contain any statements which may produce.
     #
@@ -595,7 +621,7 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple,Vector{Any}}
     # Owing to splitting blocks up, we will need to re-label some `GotoNode`s and
     # `GotoIfNot`s. To understand this, consider the following block, whose original `ID`
     # we assume to be `ID(old_id)`.
-    # ID(new_id) - %1 = φ(ID(3) => ...) 
+    # ID(new_id) - %1 = φ(ID(3) => ...)
     # ID(new_id) - %3 = call_which_must_not_produce(...)
     # ID(new_id) - %4 = produce(%3)
     # ID(old_id) - GotoNode(ID(5))

src/test_utils.jl

@@ -170,6 +170,38 @@ function test_cases()
             [true, 1],
             none,
         ),
+        Testcase(
+            "nested with args (static)",
+            nothing,
+            (static_nested_outer_args,),
+            nothing,
+            [:a, :b, false],
+            none,
+        ),
+        Testcase(
+            "nested with args (static + used)",
+            nothing,
+            (static_nested_outer_use_produced_args,),
+            nothing,
+            [:a, :b, 1],
+            none,
+        ),
+        Testcase(
+            "nested with args (dynamic)",
+            nothing,
+            (dynamic_nested_outer_args, Ref{Any}(nested_inner_args)),
+            nothing,
+            [:a, :b, false],
+            none,
+        ),
+        Testcase(
+            "nested with args (dynamic + used)",
+            nothing,
+            (dynamic_nested_outer_use_produced_args, Ref{Any}(nested_inner_args)),
+            nothing,
+            [:a, :b, 1],
+            none,
+        ),
         Testcase(
             "callable struct", nothing, (CallableStruct(5), 4), nothing, [5, 4, 9], allocs
         ),
@@ -334,6 +366,40 @@ function dynamic_nested_outer_use_produced(f::Ref{Any})
     return nothing
 end
 
+@noinline function nested_inner_args(xs...)
+    for x in xs
+        produce(x)
+    end
+    return 1
+end
+
+Libtask.might_produce(::Type{<:Tuple{typeof(nested_inner_args),Any}}) = true
+Libtask.might_produce(::Type{<:Tuple{typeof(nested_inner_args),Any,Vararg}}) = true
+
+function static_nested_outer_args()
+    nested_inner_args(:a, :b)
+    produce(false)
+    return nothing
+end
+
+function static_nested_outer_use_produced_args()
+    y = nested_inner_args(:a, :b)
+    produce(y)
+    return nothing
+end
+
+function dynamic_nested_outer_args(f::Ref{Any})
+    f[](:a, :b)
+    produce(false)
+    return nothing
+end
+
+function dynamic_nested_outer_use_produced_args(f::Ref{Any})
+    y = f[](:a, :b)
+    produce(y)
+    return nothing
+end
+
 struct CallableStruct{T}
     x::T
 end