improve type inference of Base.aligned_sizeof

[aviatesk]

This commit includes a bit of refactoring of Base.aligned_sizeof to make it more inference-friendly, especially in cases like Base.aligned_sizeof(::Union{DataType,Union}).

In particular, it eliminates the chance of inference accounting for a method error of datatype_alignment(::Union) in the second branch.

xref: aviatesk/JET.jl#512

[Tokazama]

Would it be helpful if datatype_alignment was in native Julia code here, cus I think I did something like that last year when learning the c library and probably still have the code

[aviatesk]

Maybe you meant allocatedinline? It's implemented in C, but it's also used for the other parts of code base including the codegen part.

[Tokazama]

I just meant that datatype_alignment is a thin wrapper for union_isinlineable. Since allocatedinline also uses union_isinlineable through jl_stored_inline. Anyway, this is the code I dug up in case it's of any use.

function datatype_inline_layout(t::DataType, ptrfree::Bool, asfield::Bool)
    ((t.name.flags & 0x04) === 0x04) || return (0, 0, 0)
    lyt = unsafe_load(convert(Ptr{DataTypeLayout}, t.layout))
    (lyt.npointers > 0 && ptrfree) && return (0, 0, 0)
    (t.name.n_uninitialized != 0) && return (0, 0, 0)
    (((lyt.flags >> 1) & 3) > 1) && return (0, 0, 0)
    al = Int(lyt.alignment)
    # primitive types in struct slots need their sizes aligned. issue #37974
    sz = (asfield && isprimitivetype(t)) ? LLT_ALIGN(Int(lyt.size), al) : Int(lyt.size)
    return (1, sz, al)
end

function union_isinlinable(@nospecialize(t::Type), ptrfree::Bool, asfield::Bool)
    @_foldable_meta
    if isa(t, DataType)
        return datatype_inline_layout(t, ptrfree, asfield)
    elseif isa(t, Union)
        ta = t.a
        if isa(ta, DataType)
            cnt, nb, align = datatype_inline_layout(ta, ptrfree, asfield)
            cnt === 0 && return (0, 0, 0)
            t = t.b
            while true
                if isa(t, Union)
                    ta = t.a
                    isa(ta, DataType) || return (0, 0, 0)
                    n, sz, al = datatype_inline_layout(ta, ptrfree, asfield)
                    n === 0 && return (0, 0, 0)
                    nb = nb < sz ? sz : nb
                    align = align < al ? al : align
                    cnt += 1
                    t = t.b
                elseif isa(t, DataType)
                    n, sz, al = datatype_inline_layout(t, ptrfree, asfield)
                    n === 0 && return (0, 0, 0)
                    nb = nb < sz ? sz : nb
                    align = align < al ? al : align
                    return (cnt + 1, nb, align)
                else
                    return (0, 0, 0)
                end
            end
        else
            return (0, 0, 0)
        end
    else
        return (0, 0, 0)
    end
end

function aligned_sizeof(@nospecialize T::Type)
    n, sz, al = union_isinlinable(T, false, false)
    if n === 0
        return Core.sizeof(Ptr{Cvoid})
    else
        return LLT_ALIGN(sz, al)
    end
end