Introduce JSON::Coder

etiennebarrie · byroot · byroot · commit a086f4bcda0c · 2025-01-20T12:51:16.000+01:00
Co-authored-by: Jean Boussier &lt;jean.boussier@gmail.com&gt;
diff --git a/benchmark/encoder.rb b/benchmark/encoder.rb
@@ -17,8 +17,10 @@
 
 def implementations(ruby_obj)
   state = JSON::State.new(JSON.dump_default_options)
+  coder = JSON::Coder.new
   {
     json: ["json", proc { JSON.generate(ruby_obj) }],
+    json_coder: ["json_coder", proc { coder.dump(ruby_obj) }],
     oj: ["oj", proc { Oj.dump(ruby_obj) }],
   }
 end
diff --git a/benchmark/parser.rb b/benchmark/parser.rb
@@ -15,9 +15,11 @@
 
 def benchmark_parsing(name, json_output)
   puts "== Parsing #{name} (#{json_output.size} bytes)"
+  coder = JSON::Coder.new
 
   Benchmark.ips do |x|
     x.report("json")      { JSON.parse(json_output) } if RUN[:json]
+    x.report("json_coder") { coder.load(json_output) } if RUN[:json_coder]
     x.report("oj")        { Oj.load(json_output) } if RUN[:oj]
     x.report("Oj::Parser") { Oj::Parser.new(:usual).parse(json_output) } if RUN[:oj]
     x.report("rapidjson") { RapidJSON.parse(json_output) } if RUN[:rapidjson]
diff --git a/ext/json/ext/generator/generator.c b/ext/json/ext/generator/generator.c
@@ -12,6 +12,7 @@ typedef struct JSON_Generator_StateStruct {
     VALUE space_before;
     VALUE object_nl;
     VALUE array_nl;
+    VALUE as_json;
 
     long max_nesting;
     long depth;
@@ -30,8 +31,8 @@ typedef struct JSON_Generator_StateStruct {
 static VALUE mJSON, cState, cFragment, mString_Extend, eGeneratorError, eNestingError, Encoding_UTF_8;
 
 static ID i_to_s, i_to_json, i_new, i_pack, i_unpack, i_create_id, i_extend, i_encode;
-static ID sym_indent, sym_space, sym_space_before, sym_object_nl, sym_array_nl, sym_max_nesting, sym_allow_nan,
-          sym_ascii_only, sym_depth, sym_buffer_initial_length, sym_script_safe, sym_escape_slash, sym_strict;
+static VALUE sym_indent, sym_space, sym_space_before, sym_object_nl, sym_array_nl, sym_max_nesting, sym_allow_nan,
+             sym_ascii_only, sym_depth, sym_buffer_initial_length, sym_script_safe, sym_escape_slash, sym_strict, sym_as_json;
 
 
 #define GET_STATE_TO(self, state) \
@@ -648,6 +649,7 @@ static void State_mark(void *ptr)
     rb_gc_mark_movable(state->space_before);
     rb_gc_mark_movable(state->object_nl);
     rb_gc_mark_movable(state->array_nl);
+    rb_gc_mark_movable(state->as_json);
 }
 
 static void State_compact(void *ptr)
@@ -658,6 +660,7 @@ static void State_compact(void *ptr)
     state->space_before = rb_gc_location(state->space_before);
     state->object_nl = rb_gc_location(state->object_nl);
     state->array_nl = rb_gc_location(state->array_nl);
+    state->as_json = rb_gc_location(state->as_json);
 }
 
 static void State_free(void *ptr)
@@ -714,6 +717,7 @@ static void vstate_spill(struct generate_json_data *data)
     RB_OBJ_WRITTEN(vstate, Qundef, state->space_before);
     RB_OBJ_WRITTEN(vstate, Qundef, state->object_nl);
     RB_OBJ_WRITTEN(vstate, Qundef, state->array_nl);
+    RB_OBJ_WRITTEN(vstate, Qundef, state->as_json);
 }
 
 static inline VALUE vstate_get(struct generate_json_data *data)
@@ -982,6 +986,8 @@ static void generate_json_fragment(FBuffer *buffer, struct generate_json_data *d
 static void generate_json(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj)
 {
     VALUE tmp;
+    bool as_json_called = false;
+start:
     if (obj == Qnil) {
         generate_json_null(buffer, data, state, obj);
     } else if (obj == Qfalse) {
@@ -1025,7 +1031,13 @@ static void generate_json(FBuffer *buffer, struct generate_json_data *data, JSON
             default:
             general:
                 if (state->strict) {
-                    raise_generator_error(obj, "%"PRIsVALUE" not allowed in JSON", CLASS_OF(obj));
+                    if (RTEST(state->as_json) && !as_json_called) {
+                        obj = rb_proc_call_with_block(state->as_json, 1, &obj, Qnil);
+                        as_json_called = true;
+                        goto start;
+                    } else {
+                        raise_generator_error(obj, "%"PRIsVALUE" not allowed in JSON", CLASS_OF(obj));
+                    }
                 } else if (rb_respond_to(obj, i_to_json)) {
                     tmp = rb_funcall(obj, i_to_json, 1, vstate_get(data));
                     Check_Type(tmp, T_STRING);
@@ -1126,6 +1138,7 @@ static VALUE cState_init_copy(VALUE obj, VALUE orig)
     objState->space_before = origState->space_before;
     objState->object_nl = origState->object_nl;
     objState->array_nl = origState->array_nl;
+    objState->as_json = origState->as_json;
     return obj;
 }
 
@@ -1277,6 +1290,28 @@ static VALUE cState_array_nl_set(VALUE self, VALUE array_nl)
     return Qnil;
 }
 
+/*
+ * call-seq: as_json()
+ *
+ * This string is put at the end of a line that holds a JSON array.
+ */
+static VALUE cState_as_json(VALUE self)
+{
+    GET_STATE(self);
+    return state->as_json;
+}
+
+/*
+ * call-seq: as_json=(as_json)
+ *
+ * This string is put at the end of a line that holds a JSON array.
+ */
+static VALUE cState_as_json_set(VALUE self, VALUE as_json)
+{
+    GET_STATE(self);
+    RB_OBJ_WRITE(self, &state->as_json, rb_convert_type(as_json, T_DATA, "Proc", "to_proc"));
+    return Qnil;
+}
 
 /*
 * call-seq: check_circular?
@@ -1498,6 +1533,7 @@ static int configure_state_i(VALUE key, VALUE val, VALUE _arg)
     else if (key == sym_script_safe)           { state->script_safe = RTEST(val); }
     else if (key == sym_escape_slash)          { state->script_safe = RTEST(val); }
     else if (key == sym_strict)                { state->strict = RTEST(val); }
+    else if (key == sym_as_json)               { state->as_json = rb_convert_type(val, T_DATA, "Proc", "to_proc"); }
     return ST_CONTINUE;
 }
 
@@ -1589,6 +1625,8 @@ void Init_generator(void)
     rb_define_method(cState, "object_nl=", cState_object_nl_set, 1);
     rb_define_method(cState, "array_nl", cState_array_nl, 0);
     rb_define_method(cState, "array_nl=", cState_array_nl_set, 1);
+    rb_define_method(cState, "as_json", cState_as_json, 0);
+    rb_define_method(cState, "as_json=", cState_as_json_set, 1);
     rb_define_method(cState, "max_nesting", cState_max_nesting, 0);
     rb_define_method(cState, "max_nesting=", cState_max_nesting_set, 1);
     rb_define_method(cState, "script_safe", cState_script_safe, 0);
@@ -1680,6 +1718,7 @@ void Init_generator(void)
     sym_script_safe = ID2SYM(rb_intern("script_safe"));
     sym_escape_slash = ID2SYM(rb_intern("escape_slash"));
     sym_strict = ID2SYM(rb_intern("strict"));
+    sym_as_json = ID2SYM(rb_intern("as_json"));
 
     usascii_encindex = rb_usascii_encindex();
     utf8_encindex = rb_utf8_encindex();
diff --git a/java/src/json/ext/Generator.java b/java/src/json/ext/Generator.java
@@ -510,6 +510,14 @@ void generate(ThreadContext context, Session session, IRubyObject object, Output
             RubyString generateNew(ThreadContext context, Session session, IRubyObject object) {
                 GeneratorState state = session.getState(context);
                 if (state.strict()) {
+                    if (state.getAsJSON() != null ) {
+                        IRubyObject value = state.getAsJSON().call(context, object);
+                        Handler handler = getHandlerFor(context.runtime, value);
+                        if (handler == GENERIC_HANDLER) {
+                            throw Utils.buildGeneratorError(context, object, value + " returned by as_json not allowed in JSON").toThrowable();
+                        }
+                        return handler.generateNew(context, session, value);
+                    }
                     throw Utils.buildGeneratorError(context, object, object + " not allowed in JSON").toThrowable();
                 } else if (object.respondsTo("to_json")) {
                     IRubyObject result = object.callMethod(context, "to_json", state);
diff --git a/java/src/json/ext/GeneratorState.java b/java/src/json/ext/GeneratorState.java
@@ -14,6 +14,7 @@
 import org.jruby.RubyInteger;
 import org.jruby.RubyNumeric;
 import org.jruby.RubyObject;
+import org.jruby.RubyProc;
 import org.jruby.RubyString;
 import org.jruby.anno.JRubyMethod;
 import org.jruby.runtime.Block;
@@ -22,6 +23,7 @@
 import org.jruby.runtime.Visibility;
 import org.jruby.runtime.builtin.IRubyObject;
 import org.jruby.util.ByteList;
+import org.jruby.util.TypeConverter;
 
 /**
  * The <code>JSON::Ext::Generator::State</code> class.
@@ -58,6 +60,8 @@ public class GeneratorState extends RubyObject {
      */
     private ByteList arrayNl = ByteList.EMPTY_BYTELIST;
 
+    private RubyProc asJSON;
+
     /**
      * The maximum level of nesting of structures allowed.
      * <code>0</code> means disabled.
@@ -211,6 +215,7 @@ public IRubyObject initialize_copy(ThreadContext context, IRubyObject vOrig) {
         this.spaceBefore = orig.spaceBefore;
         this.objectNl = orig.objectNl;
         this.arrayNl = orig.arrayNl;
+        this.asJSON = orig.asJSON;
         this.maxNesting = orig.maxNesting;
         this.allowNaN = orig.allowNaN;
         this.asciiOnly = orig.asciiOnly;
@@ -353,6 +358,22 @@ public IRubyObject array_nl_set(ThreadContext context,
         return arrayNl;
     }
 
+    public RubyProc getAsJSON() {
+        return asJSON;
+    }
+
+    @JRubyMethod(name="as_json")
+    public IRubyObject as_json_get(ThreadContext context) {
+        return asJSON == null ? context.getRuntime().getFalse() : asJSON;
+    }
+
+    @JRubyMethod(name="as_json=")
+    public IRubyObject as_json_set(ThreadContext context,
+                                   IRubyObject asJSON) {
+        this.asJSON = (RubyProc)TypeConverter.convertToType(asJSON, context.getRuntime().getProc(), "to_proc");
+        return asJSON;
+    }
+
     @JRubyMethod(name="check_circular?")
     public RubyBoolean check_circular_p(ThreadContext context) {
         return RubyBoolean.newBoolean(context, maxNesting != 0);
@@ -487,6 +508,8 @@ public IRubyObject _configure(ThreadContext context, IRubyObject vOpts) {
         ByteList arrayNl = opts.getString("array_nl");
         if (arrayNl != null) this.arrayNl = arrayNl;
 
+        this.asJSON = opts.getProc("as_json");
+
         ByteList objectNl = opts.getString("object_nl");
         if (objectNl != null) this.objectNl = objectNl;
 
@@ -522,6 +545,7 @@ public RubyHash to_h(ThreadContext context) {
         result.op_aset(context, runtime.newSymbol("space_before"), space_before_get(context));
         result.op_aset(context, runtime.newSymbol("object_nl"), object_nl_get(context));
         result.op_aset(context, runtime.newSymbol("array_nl"), array_nl_get(context));
+        result.op_aset(context, runtime.newSymbol("as_json"), as_json_get(context));
         result.op_aset(context, runtime.newSymbol("allow_nan"), allow_nan_p(context));
         result.op_aset(context, runtime.newSymbol("ascii_only"), ascii_only_p(context));
         result.op_aset(context, runtime.newSymbol("max_nesting"), max_nesting_get(context));
diff --git a/java/src/json/ext/OptionsReader.java b/java/src/json/ext/OptionsReader.java
@@ -10,10 +10,12 @@
 import org.jruby.RubyClass;
 import org.jruby.RubyHash;
 import org.jruby.RubyNumeric;
+import org.jruby.RubyProc;
 import org.jruby.RubyString;
 import org.jruby.runtime.ThreadContext;
 import org.jruby.runtime.builtin.IRubyObject;
 import org.jruby.util.ByteList;
+import org.jruby.util.TypeConverter;
 
 final class OptionsReader {
     private final ThreadContext context;
@@ -110,4 +112,10 @@ public RubyHash getHash(String key) {
         if (value == null || value.isNil()) return new RubyHash(runtime);
         return (RubyHash) value;
     }
+
+    RubyProc getProc(String key) {
+        IRubyObject value = get(key);
+        if (value == null) return null;
+        return (RubyProc)TypeConverter.convertToType(value, runtime.getProc(), "to_proc");
+    }
 }
diff --git a/lib/json/common.rb b/lib/json/common.rb
@@ -844,6 +844,27 @@ def merge_dump_options(opts, strict: NOT_SET)
   class << self
     private :merge_dump_options
   end
+
+  class Coder
+    def initialize(options = nil, &as_json)
+      if options.nil?
+        options = { strict: true }
+      else
+        options[:strict] = true
+      end
+      options[:as_json] = as_json if as_json
+      @state = State.new(options)
+      @parser_config = Ext::Parser::Config.new(options)
+    end
+
+    def dump(...)
+      @state.generate(...)
+    end
+
+    def load(source)
+      @parser_config.parse(source)
+    end
+  end
 end
 
 module ::Kernel
diff --git a/lib/json/ext/generator/state.rb b/lib/json/ext/generator/state.rb
@@ -58,6 +58,7 @@ def to_h
             space_before: space_before,
             object_nl: object_nl,
             array_nl: array_nl,
+            as_json: as_json,
             allow_nan: allow_nan?,
             ascii_only: ascii_only?,
             max_nesting: max_nesting,
diff --git a/lib/json/truffle_ruby/generator.rb b/lib/json/truffle_ruby/generator.rb
@@ -142,6 +142,7 @@ def initialize(opts = nil)
           @array_nl              = ''
           @allow_nan             = false
           @ascii_only            = false
+          @as_json               = false
           @depth                 = 0
           @buffer_initial_length = 1024
           @script_safe           = false
@@ -167,6 +168,9 @@ def initialize(opts = nil)
         # This string is put at the end of a line that holds a JSON array.
         attr_accessor :array_nl
 
+        # This proc converts unsupported types into native JSON types.
+        attr_accessor :as_json
+
         # This integer returns the maximum level of data structure nesting in
         # the generated JSON, max_nesting = 0 if no maximum is checked.
         attr_accessor :max_nesting
@@ -251,6 +255,7 @@ def configure(opts)
           @object_nl             = opts[:object_nl]     || '' if opts.key?(:object_nl)
           @array_nl              = opts[:array_nl]      || '' if opts.key?(:array_nl)
           @allow_nan             = !!opts[:allow_nan]         if opts.key?(:allow_nan)
+          @as_json               = opts[:as_json].to_proc     if opts.key?(:as_json)
           @ascii_only            = opts[:ascii_only]          if opts.key?(:ascii_only)
           @depth                 = opts[:depth] || 0
           @buffer_initial_length ||= opts[:buffer_initial_length]
@@ -403,8 +408,20 @@ module Object
           # it to a JSON string, and returns the result. This is a fallback, if no
           # special method #to_json was defined for some object.
           def to_json(state = nil, *)
-            if state && State.from_state(state).strict?
-              raise GeneratorError.new("#{self.class} not allowed in JSON", self)
+            state = State.from_state(state) if state
+            if state&.strict?
+              value = self
+              if state.strict? && !(false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value)
+                if state.as_json
+                  value = state.as_json.call(value)
+                  unless false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value
+                    raise GeneratorError.new("#{value.class} returned by #{state.as_json} not allowed in JSON", value)
+                  end
+                  value.to_json(state)
+                else
+                  raise GeneratorError.new("#{value.class} not allowed in JSON", value)
+                end
+              end
             else
               to_s.to_json
             end
@@ -455,7 +472,15 @@ def json_transform(state)
 
               result = +"#{result}#{key_json}#{state.space_before}:#{state.space}"
               if state.strict? && !(false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value)
-                raise GeneratorError.new("#{value.class} not allowed in JSON", value)
+                if state.as_json
+                  value = state.as_json.call(value)
+                  unless false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value
+                    raise GeneratorError.new("#{value.class} returned by #{state.as_json} not allowed in JSON", value)
+                  end
+                  result << value.to_json(state)
+                else
+                  raise GeneratorError.new("#{value.class} not allowed in JSON", value)
+                end
               elsif value.respond_to?(:to_json)
                 result << value.to_json(state)
               else
@@ -508,7 +533,15 @@ def json_transform(state)
               result << delim unless first
               result << state.indent * depth if indent
               if state.strict? && !(false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value)
-                raise GeneratorError.new("#{value.class} not allowed in JSON", value)
+                if state.as_json
+                  value = state.as_json.call(value)
+                  unless false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value
+                    raise GeneratorError.new("#{value.class} returned by #{state.as_json} not allowed in JSON", value)
+                  end
+                  result << value.to_json(state)
+                else
+                  raise GeneratorError.new("#{value.class} not allowed in JSON", value)
+                end
               elsif value.respond_to?(:to_json)
                 result << value.to_json(state)
               else
diff --git a/test/json/json_coder_test.rb b/test/json/json_coder_test.rb
diff --git a/test/json/json_generator_test.rb b/test/json/json_generator_test.rb

Original file line number	Diff line number	Diff line change
`@@ -17,8 +17,10 @@`
`17`	`17`
`18`	`18`	`def implementations(ruby_obj)`
`19`	`19`	`state = JSON::State.new(JSON.dump_default_options)`
	`20`	`+ coder = JSON::Coder.new`
`20`	`21`	`{`
`21`	`22`	`json: ["json", proc { JSON.generate(ruby_obj) }],`
	`23`	`+ json_coder: ["json_coder", proc { coder.dump(ruby_obj) }],`
`22`	`24`	`oj: ["oj", proc { Oj.dump(ruby_obj) }],`
`23`	`25`	`}`
`24`	`26`	`end`