Streamline typed JSON parsing for trim compilation

Project.toml

@@ -22,7 +22,7 @@ Arrow = "2.8.0"
 ArrowTypes = "2.2"
 Parsers = "1, 2"
 PrecompileTools = "1"
-StructUtils = "2.8"
+StructUtils = "2.8.3"
 julia = "1.9"
 
 [extras]

src/JSON.jl

@@ -40,7 +40,14 @@ end
 
 @enum Error InvalidJSON UnexpectedEOF ExpectedOpeningObjectChar ExpectedOpeningQuoteChar ExpectedOpeningArrayChar ExpectedClosingArrayChar ExpectedComma ExpectedColon ExpectedNewline InvalidChar InvalidNumber InvalidUTF16
 
-@noinline function invalid(error, buf, pos::Int, T)
+@generated _typename(::Type{T}) where {T} = QuoteNode(string(T))
+
+@noinline invalid(error, buf, pos::Int, ::Type{T}) where {T} =
+    _invalid(error, buf, pos, _typename(T))
+@noinline invalid(error, buf, pos::Int, typename::String) =
+    _invalid(error, buf, pos, typename)
+
+@noinline function _invalid(error, buf, pos::Int, typename::String)
     # compute which line the error falls on by counting “\n” bytes up to pos
     cus = buf isa AbstractString ? codeunits(buf) : buf
     line_no = count(b -> b == UInt8('\n'), view(cus, 1:pos)) + 1
@@ -64,7 +71,7 @@ end
     # we call @invoke here to avoid --trim verify errors
     caret = @invoke(repeat(" "::String, (erri + 2)::Integer)) * "^"
     msg = """
-    invalid JSON at byte position $(pos) (line $line_no) parsing type $T: $error
+    invalid JSON at byte position $(pos) (line $line_no) parsing type $(typename): $error
     $snippet$(error == UnexpectedEOF ? " <EOF>" : "...")
     $caret
     """
@@ -140,10 +147,29 @@ print(a, indent=nothing) = print(stdout, a, indent)
 "See [`json`](@ref)."
 print
 
+# typed-parse workload struct: exercising one struct with the common field
+# shapes caches the shared make/lift/array-chain inference in this package's
+# image, so downstream typed parses hit those caches instead of re-inferring
+# (which also keeps `juliac --trim` edge inference precise on nested families)
+struct _WorkloadInner
+    x::Int
+    y::Float64
+end
+struct _Workload
+    item::Union{Nothing,_WorkloadInner}
+    items::Vector{_WorkloadInner}
+    tags::Vector{String}
+    note::Union{Nothing,String}
+end
+
 @compile_workload begin
     x = JSON.parse("{\"a\": 1, \"b\": null, \"c\": true, \"d\": false, \"e\": \"\", \"f\": [1,null,true], \"g\": {\"key\": \"value\"}}")
     json = JSON.json(x)
     isvalidjson(json)
+    JSON.parse(
+        "{\"item\":{\"x\":1,\"y\":2.0},\"items\":[{\"x\":3,\"y\":4.0}],\"tags\":[\"p\"],\"note\":null}",
+        _Workload,
+    )
 end
 
 

src/lazy.jl

@@ -182,14 +182,21 @@ Selectors.@selectors LazyValues
 Base.lastindex(x::LazyValues) = length(x)
 
 # this ensures LazyValues can be "sources" in StructUtils.make
-@inline function StructUtils.applyeach(::StructUtils.StructStyle, f, x::LazyValues)
+function StructUtils.applyeach(::StructUtils.StructStyle, f, x::LazyValues)
     type = gettype(x)
     if type == JSONTypes.OBJECT
         return applyobject(f, x)
     elseif type == JSONTypes.ARRAY
         return applyarray(f, x)
     end
-    throw(ArgumentError("applyeach not applicable for `$(typeof(x))` with JSON type = `$type`"))
+    typename = get(JSONTypes.names, type, "UNKNOWN")
+    throw(ArgumentError(string(
+        "applyeach not applicable for `",
+        _typename(typeof(x)),
+        "` with JSON type = `JSONTypes.",
+        typename,
+        '`',
+    )))
 end
 
 @inline function Base.foreach(f, x::LazyValues)
@@ -254,7 +261,7 @@ end
 
 # core JSON object parsing function
 # takes a `keyvalfunc` that is applied to each key/value pair
-# `keyvalfunc` is provided a PtrString => LazyValue pair
+# `keyvalfunc` receives a PtrString => LazyValue pair
 # `keyvalfunc` can return `StructUtils.EarlyReturn` to short-circuit parsing
 # otherwise, it should return a `pos::Int` value that notes the next position to continue parsing
 # to materialize the key, call `convert(String, key)`
@@ -445,7 +452,9 @@ function Base.convert(::Type{String}, x::PtrString)
     return unsafe_string(x.ptr, x.len)
 end
 
-Base.convert(::Type{Symbol}, x::PtrString) = ccall(:jl_symbol_n, Ref{Symbol}, (Ptr{UInt8}, Int), x.ptr, x.len)
+Base.convert(::Type{Symbol}, x::PtrString) = x.escaped ?
+    Symbol(convert(String, x)) :
+    ccall(:jl_symbol_n, Ref{Symbol}, (Ptr{UInt8}, Int), x.ptr, x.len)
 
 function Base.convert(::Type{T}, x::PtrString) where {T <: Enum}
     sym = convert(Symbol, x)
@@ -455,17 +464,27 @@ function Base.convert(::Type{T}, x::PtrString) where {T <: Enum}
     throw(ArgumentError("invalid `$T` string value: \"$sym\""))
 end
 
-Base.:(==)(x::PtrString, y::AbstractString) = x.len == sizeof(y) && ccall(:memcmp, Cint, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), x.ptr, pointer(y), x.len) == 0
+Base.:(==)(x::PtrString, y::AbstractString) = x.escaped ?
+    convert(String, x) == y :
+    x.len == sizeof(y) && ccall(:memcmp, Cint, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), x.ptr, pointer(y), x.len) == 0
 Base.:(==)(x::AbstractString, y::PtrString) = y == x
-Base.:(==)(x::PtrString, y::PtrString) = x.len == y.len && ccall(:memcmp, Cint, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), x.ptr, y.ptr, x.len) == 0
+Base.:(==)(x::PtrString, y::PtrString) = x.escaped || y.escaped ?
+    convert(String, x) == convert(String, y) :
+    x.len == y.len && ccall(:memcmp, Cint, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), x.ptr, y.ptr, x.len) == 0
 Base.isequal(x::PtrString, y::AbstractString) = x == y
 Base.isequal(x::AbstractString, y::PtrString) = y == x
 Base.isequal(x::PtrString, y::PtrString) = x == y
-Base.hash(x::PtrString, h::UInt) = hash(unsafe_string(x.ptr, x.len), h)
+Base.hash(x::PtrString, h::UInt) = x.escaped ?
+    hash(convert(String, x), h) :
+    hash(unsafe_string(x.ptr, x.len), h)
 StructUtils.keyeq(x::PtrString, y::AbstractString) = x == y
 StructUtils.keyeq(x::PtrString, y::String) = x == y
 StructUtils.keyeq(x::PtrString, y::Symbol) = convert(Symbol, x) == y
 
+# JSON owns PtrString and its comparison semantics, so it can use the ordered
+# field cursor without changing StructUtils' behavior for arbitrary key types.
+@inline StructUtils.orderedfieldmatch(x::PtrString, field::String) = x == field
+
 # core JSON string parsing function
 # returns a PtrString and the next position to parse
 # a PtrString is a semi-lazy, internal-only representation

src/parse.jl

@@ -172,6 +172,7 @@ StructUtils.initialize(::JSONReadStyle, ::Type{Object}, source) = DEFAULT_OBJECT
 # this allows struct fields to specify tags under the json key specifically to override JSON behavior
 StructUtils.fieldtagkey(::JSONStyle) = :json
 StructUtils.defaultstate(st::JSONReadStyle) = StructUtils.defaultstate(st.style)
+StructUtils.orderedfields(::JSONReadStyle) = true
 
 # forward StructUtils API to the inner style so user-provided JSONStyle dispatches are honored
 StructUtils.dictlike(st::JSONReadStyle, ::Type{T}) where {T} = StructUtils.dictlike(st.style, T)
@@ -188,15 +189,43 @@ function jsonreadstyle(::Type{T}, ::Type{O}, null, style::StructStyle, unknown_f
     ignore_unknown_fields =
         unknown_fields === :ignore ? true :
         unknown_fields === :error ? false :
-        throw(ArgumentError("`unknown_fields` must be `:ignore` or `:error`, got `$(repr(unknown_fields))`"))
+        throw(ArgumentError(string(
+            "`unknown_fields` must be `:ignore` or `:error`, got ",
+            _unknownoption(unknown_fields),
+        )))
     if T === Any && !ignore_unknown_fields
         throw(ArgumentError("`unknown_fields` is only supported when parsing into a target type or existing object"))
     end
     return JSONReadStyle{O}(null, style, ignore_unknown_fields)
 end
 
-@noinline unknownfielderror(::Type{T}, key) where {T} =
-    ArgumentError("encountered unknown JSON member $(repr(key)) while parsing `$T`")
+# Use the writer's existing byte escape table. This keeps error text accurate
+# without pulling generic `repr` and type-display machinery into trim images.
+function _quotedstring(key::AbstractString)
+    n = escapelength(key) + 2
+    buf = Vector{UInt8}(undef, n)
+    _string(buf, 1, key, nothing, n)
+    return String(buf)
+end
+
+function _unknownoption(option::Symbol)
+    value = String(option)
+    return Base.isidentifier(value) ? string(':', value) :
+        string("Symbol(", _quotedstring(value), ')')
+end
+_unknownkey(key::PtrString) = _quotedstring(convert(String, key))
+_unknownkey(key::AbstractString) = _quotedstring(key)
+_unknownkey(key::Symbol) = _unknownoption(key)
+_unknownkey(key::Integer) = string(key)
+_unknownkey(key) = "<key>"
+
+@noinline unknownfielderror(::Type{T}, key) where {T} = ArgumentError(string(
+    "encountered unknown JSON member ",
+    _unknownkey(key),
+    " while parsing `",
+    _typename(T),
+    "`",
+))
 
 function StructUtils.unknownfield(st::JSONReadStyle, ::Type{T}, key, value) where {T}
     st.ignore_unknown_fields || throw(unknownfielderror(T, key))
@@ -217,20 +246,28 @@ parse(io::Union{IO,Base.AbstractCmd}, ::Type{T}=Any; kw...) where {T} = parse(Ba
 
 parse!(io::Union{IO,Base.AbstractCmd}, x::T; kw...) where {T} = parse!(Base.read(io), x; kw...)
 
+# No forced @inline through the entry chain: inlining the typed descent into
+# these forwarding bodies makes each entry's compilation unit re-optimize the
+# entire per-type parse tower instead of calling the already-compiled
+# instances (measured at +12s on the first parse of a 35-field Union-typed
+# struct, and 30-50% of the first parse of a 13-type struct family). The buffer
+# entry therefore calls the positional core directly. Typed customization for
+# buffer inputs stays at the StructUtils style/make/lift boundary. A parse
+# method specialized on LazyValue applies when the caller passes one directly.
 parse(buf::Union{AbstractVector{UInt8},AbstractString}, ::Type{T}=Any;
     dicttype::Type{O}=DEFAULT_OBJECT_TYPE, null=nothing, style::StructStyle=StructUtils.DefaultStyle(),
     unknown_fields::Symbol=:ignore, kw...) where {T,O} =
-    @inline parse(lazy(buf; kw...), T; dicttype, null, style, unknown_fields)
+    _parse(lazy(buf; kw...), T, dicttype, null, jsonreadstyle(T, O, null, style, unknown_fields))
 
 parse!(buf::Union{AbstractVector{UInt8},AbstractString}, x::T;
     dicttype::Type{O}=DEFAULT_OBJECT_TYPE, null=nothing, style::StructStyle=StructUtils.DefaultStyle(),
     unknown_fields::Symbol=:ignore, kw...) where {T,O} =
-    @inline parse!(lazy(buf; kw...), x; dicttype, null, style, unknown_fields)
+    StructUtils.make!(jsonreadstyle(typeof(x), O, null, style, unknown_fields), x, lazy(buf; kw...))
 
 parse(x::LazyValue, ::Type{T}=Any;
     dicttype::Type{O}=DEFAULT_OBJECT_TYPE, null=nothing, style::StructStyle=StructUtils.DefaultStyle(),
     unknown_fields::Symbol=:ignore) where {T,O} =
-    @inline _parse(x, T, dicttype, null, jsonreadstyle(T, O, null, style, unknown_fields))
+    _parse(x, T, dicttype, null, jsonreadstyle(T, O, null, style, unknown_fields))
 
 function _parse(x::LazyValue, ::Type{T}, dicttype::Type{O}, null, style::StructStyle) where {T,O}
     y, pos = StructUtils.make(style, T, x)