Hack to avoid inference recursion limit (#15)

* add a hack to avoid inference recursion limit to allow more optimization

For this target:
```julia
using JSONBase, BenchmarkTools
struct A
    a::Int
    b::Int
    c::Int
    d::Int
end
@benchmark JSONBase.materialize("""{ "a": 1, "b": 2, "c": 3, "d": 4}""", A)
```

Before:
```julia
BenchmarkTools.Trial: 10000 samples with 331 evaluations.
 Range (min … max):  258.937 ns …   4.373 μs  ┊ GC (min … max): 0.00% … 91.32%
 Time  (median):     279.202 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   291.054 ns ± 185.340 ns  ┊ GC (mean ± σ):  3.24% ±  4.74%

           ▃█▃    ▂
  ▁▂▂▂▁▁▁▁▃███▃▂▂▄█▇▅▅▅▄▂▂▂▃▃▄▃▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  259 ns           Histogram: frequency by time          334 ns <

 Memory estimate: 384 bytes, allocs estimate: 10.
```

After:
```julia
BenchmarkTools.Trial: 10000 samples with 750 evaluations.
 Range (min … max):  173.333 ns …  1.618 μs  ┊ GC (min … max): 0.00% … 86.43%
 Time  (median):     179.389 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   184.518 ns ± 58.847 ns  ┊ GC (mean ± σ):  1.40% ±  3.86%

          ▂▄▆██▇▅▄▁▁ ▂▃▄▅▄▄▄▄▂▁▁▁▂▂▂▂▂▂▂▂▂▂▁                   ▂
  ▂▅▄▅▅▃▃█████████████████████████████████████████████▇▇▅▄▅▆▅▅ █
  173 ns        Histogram: log(frequency) by time       200 ns <

 Memory estimate: 128 bytes, allocs estimate: 2.
```

closes #2

* Rebase/update PR

* fix 32-bit

* fix 32-bit more

---------

Co-authored-by: Shuhei Kadowaki <aviatesk@gmail.com>

LICENSE.md

@@ -1,4 +1,4 @@
-The Example.jl package is licensed under the MIT "Expat" License:
+The JSON.jl package is licensed under the MIT "Expat" License:
 
 > Copyright (c) 2022: Jacob Quinn
 >

Project.toml

@@ -14,8 +14,9 @@ Parsers = "2.5.1"
 julia = "1.6"
 
 [extras]
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "OrderedCollections"]
+test = ["Test", "JET", "OrderedCollections"]

src/JSONBase.jl

@@ -17,7 +17,13 @@ include("utils.jl")
 include("interfaces.jl")
 using .API
 
-# TODO: why is this needed?
+# in `materialize`, given an initial LazyValue/BinaryValue
+# and a possible Union or abstract type `T`, we want to
+# concretize to a more specific type based on runtime values in `x`
+# `choosetype` can be overloaded for custom scenarios, but by default
+# we can at least cover the case where `T` is a Union
+# and `x` is an object, array, or string and strip away any
+# `Nothing` or `Missing` types (very common Union types)
 function API.choosetype(::Type{T}, x) where {T}
   if T isa Union
       type = gettype(x)
@@ -30,10 +36,8 @@ function API.choosetype(::Type{T}, x) where {T}
   return T
 end
 
-# TODO: remove?
 pass(args...) = Continue(0)
 
-# TODO: document this, why, where is it used, move to utils.jl?
 struct LengthClosure
   len::Ptr{Int}
 end
@@ -88,6 +92,70 @@ print(io::IO, obj, indent=nothing) = json(io, obj; pretty=something(indent, 0))
 print(a, indent=nothing) = print(stdout, a, indent)
 json(a, indent::Integer) = json(a; pretty=indent)
 
+# HACK to avoid inference recursion limit and the de-optimization:
+# This works since know the inference recursion will terminate due to the fact that this
+# method is only called when materializing a struct with definite number of fields, i.e.
+# that is not self-referencing, so it is guaranteed that there are no cycles in a recursive
+# `materialize` call. Especially, the `fieldcount` call in the struct fallback case within
+# the `materialize` should have errored for this case.
+# TODO we should revisit this hack when we start to support https://github.com/quinnj/JSONBase.jl/issues/3
+function validate_recursion_relation_sig(f, nargs::Int, sig)
+  @nospecialize f sig
+  sig = Base.unwrap_unionall(sig)
+  @assert sig isa DataType "unexpected `recursion_relation` call"
+  @assert sig.name === Tuple.name "unexpected `recursion_relation` call"
+  @assert length(sig.parameters) == nargs "unexpected `recursion_relation` call"
+  @assert sig.parameters[1] == typeof(f) "unexpected `recursion_relation` call"
+  return sig
+end
+@static if hasfield(Method, :recursion_relation)
+  let applyobject_recursion_relation = function (
+          method::Method, topmost::Union{Nothing,Method},
+          @nospecialize(sig), @nospecialize(topmostsig))
+          # Core.println("applyobject")
+          # Core.println("  method = ", method)
+          # Core.println("  topmost = ", topmost)
+          # Core.println("  sig = ", sig)
+          # Core.println("  topmostsig = ", topmostsig)
+          sig = validate_recursion_relation_sig(applyobject, 3, sig)
+          topmostsig = validate_recursion_relation_sig(applyobject, 3, topmostsig)
+          return sig.parameters[2] ≠ topmostsig.parameters[2]
+      end
+      method = only(methods(applyobject, (Any,LazyValue,)))
+      method.recursion_relation = applyobject_recursion_relation
+  end
+  let applyfield_recursion_relation = function (
+          method::Method, topmost::Union{Nothing,Method},
+          @nospecialize(sig), @nospecialize(topmostsig))
+          # Core.println("applyfield")
+          # Core.println("  method = ", method)
+          # Core.println("  topmost = ", topmost)
+          # Core.println("  sig = ", sig)
+          # Core.println("  topmostsig = ", topmostsig)
+          sig = validate_recursion_relation_sig(applyfield, 7, sig)
+          topmostsig = validate_recursion_relation_sig(applyfield, 7, topmostsig)
+          return sig.parameters[2] ≠ topmostsig.parameters[2]
+      end
+      method = only(methods(applyfield, (Type,Any,Type,Any,Any,Any)))
+      method.recursion_relation = applyfield_recursion_relation
+  end
+  let _materialize_recursion_relation = function (
+          method::Method, topmost::Union{Nothing,Method},
+          @nospecialize(sig), @nospecialize(topmostsig))
+          # Core.println("_materialize")
+          # Core.println("  method = ", method)
+          # Core.println("  topmost = ", topmost)
+          # Core.println("  sig = ", sig)
+          # Core.println("  topmostsig = ", topmostsig)
+          sig = validate_recursion_relation_sig(_materialize, 6, sig)
+          topmostsig = validate_recursion_relation_sig(_materialize, 6, topmostsig)
+          return sig.parameters[4] ≠ topmostsig.parameters[4]
+      end
+      method = only(methods(_materialize, (Any,LazyValue,Type,Any,Type)))
+      method.recursion_relation = _materialize_recursion_relation
+  end
+end
+
 end # module
 
 #TODO
@@ -114,4 +182,4 @@ end # module
      # how to form json
        # create Dict/NamedTuple/Array and call tojson
        # use struct and call tojson
-       # support jsonlines output
+       # support jsonlines output

src/lazy.jl

@@ -242,7 +242,7 @@ _applyobject(f::F, x) where {F} = applyobject(f, x)
         # we're now positioned at the start of the value
         val = lazy(buf, pos, len, b, opts)
         ret = keyvalfunc(key, val)
-        # if ret is not an Continue, then we're 
+        # if ret is not an Continue, then we're
         # short-circuiting parsing via e.g. selection syntax
         # so return immediately
         ret isa Continue || return ret

src/utils.jl

@@ -8,7 +8,7 @@ $(Base.String(buf[max(1, pos-25):min(end, pos+25)]))
 """))
 
 # like nonnothingtype + nonmissingtype together
-non_nothing_missing_type(::Type{T}) where {T} = Base.typesplit(Base.typesplit(T, Nothing), Missing)
+non_nothing_missing_type(@nospecialize T) = Base.typesplit(Base.typesplit(T, Nothing), Missing)
 
 # helper struct we pack lazy-specific keyword args into
 # held by LazyValue for access

test/json.jl

@@ -126,7 +126,7 @@ end
     # SimpleVector writing
     @test JSONBase.json(Core.svec(1, 2, 3)) == "[1,2,3]"
     # Ptr writing
-    @test JSONBase.json(C_NULL) == "\"Ptr{Nothing} @0x0000000000000000\""
+    sizeof(Int) == 8 && @test JSONBase.json(C_NULL) == "\"Ptr{Nothing} @0x0000000000000000\""
     # DataType writing
     @test JSONBase.json(Float64) == "\"Float64\""
     @test JSONBase.json(Union{Missing, Float64}) == "\"Union{Missing, Float64}\""

test/optimization.jl

@@ -0,0 +1,43 @@
+module test_optimization
+
+using JSONBase, JET
+
+struct OptimizationFailureChecker end
+function JET.configured_reports(::OptimizationFailureChecker, reports::Vector{JET.InferenceErrorReport})
+    return filter(reports) do @nospecialize report::JET.InferenceErrorReport
+        isa(report, JET.OptimizationFailureReport)
+    end
+end
+
+# https://github.com/quinnj/JSONBase.jl/issues/2
+struct Simple
+    a::Int
+    b::Int
+end
+@test_opt annotate_types=true report_config=OptimizationFailureChecker() JSONBase.materialize("""{ "a": 1, "b": 2 }""", Simple)
+
+struct Inner
+    b::Int
+end
+struct Outer
+    a::Int
+    b::Inner
+end
+@test_opt annotate_types=true report_config=OptimizationFailureChecker() JSONBase.materialize("""{ "a": 1, "b": { "b": 2 } }""", Outer)
+
+struct SelfRecur
+    a1::Int
+    a2::Union{Nothing,SelfRecur}
+end
+@test_opt annotate_types=true report_config=OptimizationFailureChecker() JSONBase.materialize("""{ "a1": 1, "a2": { "a1": 2 } }""", SelfRecur)
+
+struct RecurInner{T}
+    a::T
+end
+struct RecurOuter
+    a1::Int
+    a2::Union{Nothing,RecurInner{RecurOuter}}
+end
+@test_opt annotate_types=true report_config=OptimizationFailureChecker() JSONBase.materialize("""{ "a1": 1, "a2": { "a": { "a1": 2 } } }""", RecurOuter)
+
+end # module test_optimization

test/runtests.jl

@@ -122,7 +122,7 @@ end
     ["Gilbert", "2013", 24, true]
     ["Alexa", "2013", 29, true]
     ["May", "2012B", 14, false]
-    ["Deloise", "2012A", 19, true] 
+    ["Deloise", "2012A", 19, true]
     """; jsonlines=true, float64=true) ==
     [["Name", "Session", "Score", "Completed"],
      ["Gilbert", "2013", 24.0, true],
@@ -368,3 +368,6 @@ end
 include("struct.jl")
 include("json.jl")
 include("numbers.jl")
+# @static if VERSION ≥ v"1.8"
+#     @testset "Optimization test with JET" include("optimization.jl")
+# end

test/struct.jl

@@ -245,7 +245,7 @@ end
     @test JSONBase.materialize("""{"type": "car","make": "Mercedes-Benz","model": "S500","seatingCapacity": 5,"topSpeed": 250.1}""", Vehicle) == Car("car", "Mercedes-Benz", "S500", 5, 250.1)
     @test JSONBase.materialize("""{"type": "truck","make": "Isuzu","model": "NQR","payloadCapacity": 7500.5}""", Vehicle) == Truck("truck", "Isuzu", "NQR", 7500.5)
     # union
-    @test JSONBase.materialize("""{"id": 1, "name": "2", "rate": 3}""", J) == J(1, "2", 3)
+    @test JSONBase.materialize("""{"id": 1, "name": "2", "rate": 3}""", J) == J(1, "2", Int64(3))
     @test JSONBase.materialize("""{"id": null, "name": null, "rate": 3.14}""", J) == J(nothing, nothing, 3.14)
     # test K
     @test JSONBase.materialize("""{"id": 1, "value": null}""", K) == K(1, missing)