diff --git a/Project.toml b/Project.toml index dea4f40..416deb8 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ version = "2.8.2" [deps] Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" [weakdeps] @@ -17,6 +18,7 @@ StructUtilsStaticArraysCoreExt = ["StaticArraysCore"] StructUtilsTablesExt = ["Tables"] [compat] +Preferences = "1" Measurements = "2" StaticArraysCore = "1" Tables = "1" diff --git a/src/StructUtils.jl b/src/StructUtils.jl index cfdcab2..1760fd0 100644 --- a/src/StructUtils.jl +++ b/src/StructUtils.jl @@ -1,5 +1,22 @@ module StructUtils +using Preferences + +# Compile-time preference ("trim_specialize"): `juliac --trim` builds opt into +# fully-specialized struct machinery — per-target-type `make` dispatchers, +# literal-unrolled `findfield`/tuple closures, and isa-split small-Union field +# writes — which the trim verifier can resolve statically. The default keeps +# the single-generic-instance forms: on a large nested-struct workload +# (JSON.jl benchmarks/structs.jl Root) full specialization roughly doubles +# first-call compile time, which library users pay on every fresh session, +# while trim builds pay it once at AOT-compile time. +# +# Enable from the build project with: +# Preferences.set_preferences!(StructUtils, "trim_specialize" => true) +# (a Preference, not an ENV switch, so flipping it correctly invalidates the +# precompile cache) +const TRIM_SPECIALIZE = @load_preference("trim_specialize", false)::Bool + using Dates, UUIDs export @noarg, @defaults, @tags, @kwarg, @nonstruct, Selectors @@ -600,19 +617,40 @@ function applyeach(st::StructStyle, f, x::T) where {T} end for i = 1:N fname = Meta.quot(fieldname(T, i)) + # for small-Union fields, split the closure call per component with an + # explicit isa chain: inference doesn't reliably union-split these call + # sites on its own, leaving dynamic dispatch that `--trim` can't resolve. + # `lower` can change the value's type (custom lowerings), so the final + # `else` keeps those correct — it's dead code for identity lowers. + callex = :(f(key, val)) + ft = fieldtype(T, i) + comps = ft isa Union ? Base.uniontypes(ft) : nothing + if TRIM_SPECIALIZE && comps !== nothing && length(comps) <= 4 && all(c -> isconcretetype(c) || c === Nothing, comps) + for c in reverse(comps) + # the typeassert matters: with identical `f(key, val)` calls in + # every branch, the optimizer tail-merges them back into a single + # call with a φ-union argument, undoing the split + callex = Expr(:if, :(val isa $c), :(f(key, val::$c)), callex) + end + end push!(ex.args, quote ftags = fieldtags(st, T, $fname) if !haskey(ftags, :ignore) || !ftags.ignore fname = get(ftags, :name, $fname) - ret = if isdefined(x, $i) - f(lowerkey(st, fname), lower(st, getfield(x, $i), ftags)) + # hoist the value out of all three source branches, then make the + # (single) closure call through the isa split — a call left inside + # any one branch keeps its unsplit dynamic dispatch + val = if isdefined(x, $i) + lower(st, getfield(x, $i), ftags) elseif haskey(defs, $fname) # this branch should be really rare because we should # have applied a field default in the struct constructor - f(lowerkey(st, fname), lower(st, defs[$fname], ftags)) + lower(st, defs[$fname], ftags) else - f(lowerkey(st, fname), lower(st, nothing, ftags)) + lower(st, nothing, ftags) end + key = lowerkey(st, fname) + ret = $callex ret isa EarlyReturn && return ret end end) @@ -846,7 +884,14 @@ end @inline abstractcollectionpassthrough(style::StructStyle, ::Type{T}, source) where {T} = isabstracttype(T) && source isa T && (dictlike(style, T) || arraylike(style, T)) -function make(style::StructStyle, T::Type, source, tags) +# `::Type{T} where T` (not bare `T::Type`, which julia deliberately doesn't specialize): +# this dispatcher must specialize per target type so its trait branches fold — otherwise +# one generic instance with everything dynamic ends up in `--trim` binaries +# One shared body, two signature forms (see TRIM_SPECIALIZE): `::Type{T} where T` +# specializes per target type — required for trim verification, but on big nested +# structs it roughly doubles first-call compile time versus the bare `T::Type` +# form, which julia deliberately compiles as a single generic instance. +const _MAKE4_BODY = quote if haskey(tags, :choosetype) return make(style, tags.choosetype(source), source, _delete(tags, :choosetype)) end @@ -868,33 +913,8 @@ function make(style::StructStyle, T::Type, source, tags) # we can disambiguate by checking if source is arraylike; # only applies when there's exactly one arraylike and one non-arraylike member if T isa Union - types = Base.uniontypes(T) - arr_type = nothing - scalar_type = nothing - ambiguous = false - for t in types - if arraylike(style, t) - # more than one arraylike type means we can't disambiguate - if arr_type !== nothing - ambiguous = true - break - end - arr_type = t - else - if scalar_type !== nothing - ambiguous = true - break - end - scalar_type = t - end - end - if !ambiguous && arr_type !== nothing && scalar_type !== nothing - if arraylike(style, source) - return make(style, arr_type, source, tags) - else - return make(style, scalar_type, source, tags) - end - end + r = _unionmake(style, T, source, tags) + r !== nothing && return r end end if T <: Tuple || dictlike(style, T) || arraylike(style, T) || noarg(style, T) || structlike(style, T) @@ -903,8 +923,7 @@ function make(style::StructStyle, T::Type, source, tags) return lift(style, T, source, tags) end end - -function make(style::StructStyle, T::Type, source) +const _MAKE3_BODY = quote if abstractcollectionpassthrough(style, T, source) return source, defaultstate(style) end @@ -923,37 +942,10 @@ function make(style::StructStyle, T::Type, source) return make(style, Base.nonnothingtype(T), source) end end - # for Union types like Union{T, Vector{T}} (after Nothing/Missing have been peeled), - # we can disambiguate by checking if source is arraylike; - # only applies when there's exactly one arraylike and one non-arraylike member + # see _unionmake: unrolled two-component union disambiguation if T isa Union - types = Base.uniontypes(T) - arr_type = nothing - scalar_type = nothing - ambiguous = false - for t in types - if arraylike(style, t) - # more than one arraylike type means we can't disambiguate - if arr_type !== nothing - ambiguous = true - break - end - arr_type = t - else - if scalar_type !== nothing - ambiguous = true - break - end - scalar_type = t - end - end - if !ambiguous && arr_type !== nothing && scalar_type !== nothing - if arraylike(style, source) - return make(style, arr_type, source) - else - return make(style, scalar_type, source) - end - end + r = _unionmake(style, T, source) + r !== nothing && return r end end if T <: Tuple @@ -970,6 +962,54 @@ function make(style::StructStyle, T::Type, source) return lift(style, T, source) end end +let sig4 = TRIM_SPECIALIZE ? :(make(style::StructStyle, ::Type{T}, source, tags) where {T}) : + :(make(style::StructStyle, T::Type, source, tags)), + sig3 = TRIM_SPECIALIZE ? :(make(style::StructStyle, ::Type{T}, source) where {T}) : + :(make(style::StructStyle, T::Type, source)) + @eval $(Expr(:function, sig4, _MAKE4_BODY)) + @eval $(Expr(:function, sig3, _MAKE3_BODY)) +end + +# unrolls the arraylike-vs-scalar union disambiguation with literal component types +# (`Base.uniontypes` builds a runtime Vector{Any}, which makes the recursive `make` +# calls dynamic — type-unstable and unresolvable under `--trim`). The rule can only +# ever match two-component unions ("exactly one arraylike and one non-arraylike"), +# so anything else returns `nothing` (ambiguous) and the caller falls through. +# The `arraylike` trait calls stay in the emitted code — they are style-dependent — +# but fold at inference given the literal types. +@generated function _unionmake(style::StructStyle, ::Type{T}, source, tags) where {T} + types = Base.uniontypes(T) + length(types) == 2 || return :(return nothing) + A, B = types + return quote + a_arr = arraylike(style, $A) + b_arr = arraylike(style, $B) + if a_arr && !b_arr + return arraylike(style, source) ? make(style, $A, source, tags) : make(style, $B, source, tags) + elseif b_arr && !a_arr + return arraylike(style, source) ? make(style, $B, source, tags) : make(style, $A, source, tags) + end + return nothing + end +end + +@generated function _unionmake(style::StructStyle, ::Type{T}, source) where {T} + types = Base.uniontypes(T) + length(types) == 2 || return :(return nothing) + A, B = types + return quote + a_arr = arraylike(style, $A) + b_arr = arraylike(style, $B) + if a_arr && !b_arr + return arraylike(style, source) ? make(style, $A, source) : make(style, $B, source) + elseif b_arr && !a_arr + return arraylike(style, source) ? make(style, $B, source) : make(style, $A, source) + end + return nothing + end +end + +# see the 4-arg method: must specialize per target type (`::Type{T}`, not `T::Type`) if VERSION < v"1.11" mem(n) = Vector{Any}(undef, n) @@ -995,6 +1035,41 @@ struct TupleClosure{T,A,S} i::Ptr{Int} end +# generated for the same reason as `findfield`: literal field indices/types keep the +# `make` calls concrete (type-stable, and resolvable under `--trim`) +@static if TRIM_SPECIALIZE + +@generated function _tupleclosure(f::TupleClosure{T}, k, v) where {T} + ex = Expr(:block) + push!(ex.args, :(Base.@_inline_meta)) + for i = 1:fieldcount(T) + ft = fieldtype(T, i) + push!(ex.args, quote + if typeof(k) == Int + if k == $i + intval, intst = make(f.style, $ft, v) + @inbounds f.vals[$i] = intval + return intst + end + else + if unsafe_load(f.i) == $i + unsafe_store!(f.i, $i + 1) + elseval, elsest = make(f.style, $ft, v) + @inbounds f.vals[$i] = elseval + return elsest + end + end + end) + end + push!(ex.args, :(return unknownfield(f.style, T, k, v))) + return ex +end + + +(f::TupleClosure{T,A,S})(k, v) where {T,A,S} = _tupleclosure(f, k, v) + +else + function (f::TupleClosure{T,A,S})(k, v) where {T,A,S} st = _foreach(T) do i if typeof(k) == Int @@ -1016,6 +1091,8 @@ function (f::TupleClosure{T,A,S})(k, v) where {T,A,S} return st isa _MatchedState ? st.value : unknownfield(f.style, T, k, v) end +end # @static if TRIM_SPECIALIZE + function maketuple(style, ::Type{T}, source) where {T} vals = mem(fieldcount(T)) ref = Ref(1) @@ -1032,7 +1109,10 @@ struct DictClosure{T,S} end function (f::DictClosure{T,S})(k, v) where {T,S} - val, st = make(f.style, _valtype(f.dict), v) + # NTuple{2,Any} assert: an Any-valued target (e.g. Dict{String,Any}) makes + # the make() return type opaque, and destructuring an opaque value is + # dynamic dispatch under `juliac --trim` + val, st = make(f.style, _valtype(f.dict), v)::NTuple{2,Any} addkeyval!(f.dict, liftkey(f.style, _keytype(f.dict), k), val) return st end @@ -1050,7 +1130,7 @@ struct ArrayClosure{T,S} end function (f::ArrayClosure{T,S})(_, v) where {T,S} - val, st = make(f.style, eltype(f.arr), v) + val, st = make(f.style, eltype(f.arr), v)::NTuple{2,Any} push!(f.arr, val) return st end @@ -1062,7 +1142,7 @@ struct FixedArrayClosure{A,S} end function (f::FixedArrayClosure{A,S})(_, v) where {A,S} - val, st = make(f.style, eltype(f.arr), v) + val, st = make(f.style, eltype(f.arr), v)::NTuple{2,Any} i = f.idx[] @inbounds f.arr[i] = val f.idx[] = i + 1 @@ -1134,6 +1214,55 @@ end setval!(vals::T, x, i) where {T} = _setfield!(vals, i, x) +# generated with literal field indices/types: a runtime `i` makes `fieldtype(T, i)` +# abstract, which funnels every `make` call into the unspecialized `T::Type` method — +# type-unstable and a pile of unresolvable dynamic calls under `--trim` +@static if TRIM_SPECIALIZE + +# literal field indices and field types: the runtime-`i` closure form below +# makes `fieldtype(T, i)` abstract, funneling every make call through one +# generic instance — unresolvable under --trim +@generated function findfield(::Type{T}, k, v, f) where {T} + ex = Expr(:block) + push!(ex.args, :(Base.@_inline_meta)) + for i = 1:fieldcount(T) + ft = fieldtype(T, i) + push!(ex.args, quote + if typeof(k) == Symbol + let fn = f.fsyms[$i], ftags = f.ftags[$i] + field = get(ftags, :name, fn) + if keyeq(k, field) || keyeq(k, fn) + symval, symst = make(f.style, $ft, v, ftags) + setval!(f.vals, symval, $i) + return symst + end + end + elseif typeof(k) == Int + if k == $i + let ftags = f.ftags[$i] + intval, intst = make(f.style, $ft, v, ftags) + setval!(f.vals, intval, $i) + return intst + end + end + else + let fstr = f.fstrs[$i], ftags = f.ftags[$i] + field = get(ftags, :name, fstr) + if keyeq(k, field) + strval, strst = make(f.style, $ft, v, ftags) + setval!(f.vals, strval, $i) + return strst + end + end + end + end) + end + push!(ex.args, :(return unknownfield(f.style, T, k, v))) + return ex +end + +else + function findfield(::Type{T}, k, v, f) where {T} st = _foreach(T) do i if typeof(k) == Symbol @@ -1167,6 +1296,8 @@ function findfield(::Type{T}, k, v, f) where {T} return st isa _MatchedState ? st.value : unknownfield(f.style, T, k, v) end +end # @static if TRIM_SPECIALIZE + (f::StructClosure{T,A,S,FS,FSS,FT})(k, v) where {T,A,S,FS,FSS,FT} = findfield(T, k, v, f) @inline makenoarg(style, ::Type{T}, source) where {T} = makenoarg(style, initialize(style, T, source), source) diff --git a/test/trim_compile_tests.jl b/test/trim_compile_tests.jl index da649d8..d259438 100644 --- a/test/trim_compile_tests.jl +++ b/test/trim_compile_tests.jl @@ -41,6 +41,9 @@ function _setup_trim_env() println(output) error("failed to set up trim test environment") end + # the trim workload must build against the trim-specialized machinery — + # the default (fast-TTFX) forms are intentionally not verifier-resolvable + write(joinpath(env_path, "LocalPreferences.toml"), "[StructUtils]\ntrim_specialize = true\n") println("[trim] temp environment ready") return env_path end