Skip to content

fix: make the struct machinery resolvable under juliac --trim #62

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
quinnj wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix: make the struct machinery resolvable under juliac --trim #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a2090c1
fix: make the struct machinery resolvable under juliac --trim
quinnj Jul 11, 2026
0212e6c
fix: typeassert the union-split closure calls so the optimizer can't …
quinnj Jul 11, 2026
c9cc25a
fix: hoist the applyeach field value before the union-split call
quinnj Jul 11, 2026
1ba17ae
fix: assert make() pair returns in the container closures
quinnj Jul 14, 2026
32329db
perf: gate the trim specialization behind a compile-time preference
quinnj Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -17,6 +18,7 @@ StructUtilsStaticArraysCoreExt = ["StaticArraysCore"]
StructUtilsTablesExt = ["Tables"]

[compat]
Preferences = "1"
Measurements = "2"
StaticArraysCore = "1"
Tables = "1"
Expand Down
265 changes: 198 additions & 67 deletions src/StructUtils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
3 changes: 3 additions & 0 deletions test/trim_compile_tests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading