Why not use @generated?

General Usage
7

Here is a (slightly contrived) example where @generated seems to hurt inference:

foldldiag_gen(op, acc, A::AbstractArray) =
    ndims(A) === 0 ? acc : _foldldiag_gen(op, acc, A, 1, size(A, 1))

@generated function _foldldiag_gen(op, acc::T, A::AbstractArray, i0, i1) where T
    idx = [:i for i in 1:ndims(A)]
    quote
        for i in i0:i1
            acc′ = op(acc, A[$(idx...)])
            acc′ isa T || return _foldldiag_gen(op, acc′, A, i + 1, i1)
            acc = acc′
        end
        return acc
    end
end

foldldiag_fun(op, acc, A::AbstractArray) =
    ndims(A) === 0 ? acc : _foldldiag_fun(op, acc, A, 1, size(A, 1))

function _foldldiag_fun(op, acc::T, A, i0, i1) where T
    for i in i0:i1
        idx = ntuple(_ -> i, ndims(A))
        acc′ = op(acc, A[idx...])
        acc′ isa T || return _foldldiag_fun(op, acc′, A, i + 1, i1)
        acc = acc′
    end
    return acc
end

Base.return_types(
    foldldiag_gen,
    Tuple{
        typeof(+),
        Bool,
        Matrix{T} where T <: Union{Int, Float64},
    },
)
# 1-element Array{Any,1}:
#  Any

Base.return_types(
    foldldiag_fun,
    Tuple{
        typeof(+),
        Bool,
        Matrix{T} where T <: Union{Int, Float64},
    },
)
# 1-element Array{Any,1}:
#  Union{Bool, Float64, Int64}
2 Likes