如何根据任何 splatted args 的类型进行调度？答案

【问题标题】：How to dispatch based on the type of any of the splatted args?如何根据任何 splatted args 的类型进行调度？
【发布时间】：2020-08-06 14:19:23
【问题描述】：

考虑 Base 中的一个现有函数，它接收一些抽象类型 T 的可变数量的参数。我已经定义了一个子类型S<:T，并且想编写一个方法，如果任何参数是我的子类型S，它就会分派。

例如，考虑函数Base.cat，其中T 是AbstractArray，S 是一些MyCustomArray <: AbstractArray。

期望的行为：

julia> v = [1, 2, 3];

julia> cat(v, v, v, dims=2)
3×3 Array{Int64,2}:
 1  1  1
 2  2  2
 3  3  3

julia> w = MyCustomArray([1,2,3])

julia> cat(v, v, w, dims=2)
"do something fancy"

尝试：

function Base.cat(w::MyCustomArray, a::AbstractArray...; dims)
    pritnln("do something fancy")
end

但这只有在第一个参数是 MyCustomArray 时才有效。

实现这一目标的优雅方法是什么？

【问题讨论】：

标签： julia dispatch

【解决方案1】：

我想说没有类型盗版就不可能干净地做到这一点（但如果可能的话，我也想学习如何）。

例如，考虑您询问的cat。它在 Base 中有一个非常通用的签名（实际上并不要求 A 是 AbstractArray 在你写的时候）：

julia> methods(cat)
# 1 method for generic function "cat":
[1] cat(A...; dims) in Base at abstractarray.jl:1654

你可以写一个具体的方法：

Base.cat(A::AbstractArray...; dims) = ...

并检查A 的任何元素是否是您的特殊数组，但这将是类型盗版。

现在的问题是你甚至不能写Union{S, T} 因为S <: T 它将被解决为T。

这意味着您必须在签名中明确使用S，但随后：

f(::S, ::T) = ...
f(::T, ::S) = ...

是有问题的，编译器会要求您定义f(::S, ::S)，因为上述定义会导致分派歧义。因此，即使您想将 varargs 的数量限制为某个最大数量，您也必须将 A 的所有划分的类型注释为子集以避免调度歧义（这可以使用宏来实现，但所需方法的数量呈指数增长)。

【讨论】：

【解决方案2】：

对于一般用法，我同意 Bogumił，但让我补充一点意见。如果您可以控制 cat 的调用方式，您至少可以编写某种 trait-dispatch 代码：

struct MyCustomArray{T, N} <: AbstractArray{T, N}
    x::Array{T, N}
end

HasCustom() = Val(false)
HasCustom(::MyCustomArray, rest...) = Val(true)
HasCustom(::AbstractArray, rest...) = HasCustom(rest...)

# `IsCustom` or something would be more elegant, but `Val` is quicker for now
Base.cat(::Val{true}, args...; dims) = println("something fancy")
Base.cat(::Val{false}, args...; dims) = cat(args...; dims=dims)

而且编译器足够酷，可以优化它：

julia> args = (v, v, w);

julia> @code_warntype cat(HasCustom(args...), args...; dims=2);
Variables
  #self#::Core.Compiler.Const(cat, false)
  #unused#::Core.Compiler.Const(Val{true}(), false)
  args::Tuple{Array{Int64,1},Array{Int64,1},MyCustomArray{Int64,1}}

Body::Nothing
1 ─ %1 = Main.println("something fancy")::Core.Compiler.Const(nothing, false)
└──      return %1

如果您没有可以控制对cat 的调用，我能想到的使上述技术起作用的唯一方法是使用包含此类调用的overdub 方法，以替换匹配由自定义实现调用。在这种情况下，您甚至不需要重载cat，而是可以直接用一些mycat 替换它来做您喜欢的事情。

【讨论】：