Eliminate SORT_KEYS_PROC const and make default proc shareable

Author: byroot

ext/json/ext/generator/generator.c

@@ -37,7 +37,7 @@ typedef struct JSON_Generator_StateStruct {
     VALUE sort_keys;
 } JSON_Generator_State;
 
-static VALUE mJSON, cState, cFragment, eGeneratorError, eNestingError, Encoding_UTF_8;
+static VALUE mJSON, cState, cFragment, eGeneratorError, eNestingError, Encoding_UTF_8, default_sort_keys_proc;
 
 static ID i_to_s, i_to_json, i_new, i_encode;
 static VALUE sym_indent, sym_space, sym_space_before, sym_object_nl, sym_array_nl, sym_max_nesting, sym_allow_nan, sym_allow_duplicate_key,
@@ -1732,12 +1732,25 @@ static VALUE cState_ascii_only_set(VALUE self, VALUE enable)
     return Qnil;
 }
 
+static VALUE cState_set_default_sort_keys_proc(VALUE self, VALUE proc)
+{
+    if (!rb_obj_is_proc(proc)) {
+        rb_raise(rb_eTypeError, "sort_key_proc must be a Proc");
+    }
+    return default_sort_keys_proc = proc;
+}
+
 static VALUE normalize_sort_keys(VALUE value)
 {
     if (rb_obj_is_proc(value)) {
         return value;
+    } else if (value == Qtrue) {
+        return default_sort_keys_proc;
+    } else if (RTEST(value)) {
+        rb_raise(rb_eTypeError, "The `sort_keys` argument must be a boolean or a Proc");
+    } else {
+        return Qfalse;
     }
-    return RTEST(value) ? rb_const_get_at(mJSON, rb_intern("SORT_KEYS_PROC")) : Qfalse;
 }
 
 /*
@@ -1958,6 +1971,8 @@ void Init_generator(void)
     VALUE mExt = rb_define_module_under(mJSON, "Ext");
     VALUE mGenerator = rb_define_module_under(mExt, "Generator");
 
+    rb_global_variable(&default_sort_keys_proc);
+
     rb_global_variable(&eGeneratorError);
     eGeneratorError = rb_path2class("JSON::GeneratorError");
 
@@ -1967,6 +1982,8 @@ void Init_generator(void)
     cState = rb_define_class_under(mGenerator, "State", rb_cObject);
     rb_define_alloc_func(cState, cState_s_allocate);
     rb_define_singleton_method(cState, "from_state", cState_from_state_s, 1);
+    rb_define_singleton_method(cState, "default_sort_keys_proc=", cState_set_default_sort_keys_proc, 1);
+
     rb_define_method(cState, "initialize", cState_initialize, -1);
     rb_define_alias(cState, "initialize", "initialize"); // avoid method redefinition warnings
     rb_define_private_method(cState, "_configure", cState_configure, 1);

java/src/json/ext/Generator.java

@@ -20,6 +20,7 @@
 import org.jruby.RubyProc;
 import org.jruby.RubyString;
 import org.jruby.RubySymbol;
+import org.jruby.anno.JRubyMethod;
 import org.jruby.runtime.Helpers;
 import org.jruby.runtime.ThreadContext;
 import org.jruby.runtime.builtin.IRubyObject;

java/src/json/ext/GeneratorState.java

@@ -37,6 +37,8 @@ public class GeneratorState extends RubyObject {
     private boolean allowDuplicateKey = false;
     private boolean deprecateDuplicateKey = true;
 
+    private static IRubyObject defaultSortKeyProc;
+
     /**
      * The indenting unit string. Will be repeated several times for larger
      * indenting levels.
@@ -165,6 +167,12 @@ static GeneratorState fromState(ThreadContext context, RuntimeInfo info,
         return (GeneratorState)klass.newInstance(context, context.nil);
     }
 
+    @JRubyMethod(meta=true, name="default_sort_keys_proc=")
+    public static IRubyObject setDefaultSortKeyProc(IRubyObject klass, IRubyObject proc) {
+        defaultSortKeyProc = proc;
+        return proc;
+    }
+
     /**
      * <code>State#initialize(opts = {})</code>
      * <p>
@@ -451,7 +459,7 @@ public RubyProc getSortKeysProc() {
     private static IRubyObject normalizeSortKeys(ThreadContext context, IRubyObject value) {
         if (value instanceof RubyProc) return value;
         if (value != null && value.isTrue()) {
-            return context.runtime.getModule("JSON").getConstant("SORT_KEYS_PROC");
+            return defaultSortKeyProc;
         }
         return null;
     }

lib/json/common.rb

@@ -155,6 +155,15 @@ def parser=(parser) # :nodoc:
     # Set the module _generator_ to be used by JSON.
     def generator=(generator) # :nodoc:
       old, $VERBOSE = $VERBOSE, nil
+
+      # The default proc used when the +sort_keys+ generation option is +true+.
+      # It returns a new hash with the entries sorted by their keys.
+      sort_keys_proc = ->(hash) { hash.sort.to_h }
+      if defined?(::Ractor) && Ractor.respond_to?(:shareable_lambda)
+        sort_keys_proc = Ractor.shareable_lambda(&sort_keys_proc)
+      end
+      generator::State.default_sort_keys_proc = sort_keys_proc
+
       @generator = generator
       if generator.const_defined?(:GeneratorMethods)
         generator_methods = generator::GeneratorMethods
@@ -243,10 +252,6 @@ def self.create_id
 
   MinusInfinity = -Infinity
 
-  # The default proc used when the +sort_keys+ generation option is +true+.
-  # It returns a new hash with the entries sorted by their keys.
-  SORT_KEYS_PROC = ->(hash) { hash.sort.to_h } # :nodoc:
-
   # The base exception for JSON errors.
   class JSONError < StandardError; end
 

lib/json/truffle_ruby/generator.rb

@@ -111,6 +111,8 @@ def self.valid_utf8?(string)
       # This class is used to create State instances, that are use to hold data
       # while generating a JSON text from a Ruby data structure.
       class State
+        singleton_class.attr_accessor :default_sort_keys_proc # :nodoc:
+
         def self.generate(obj, opts = nil, io = nil)
           new(opts).generate(obj, io)
         end
@@ -207,7 +209,16 @@ def initialize(opts = nil)
         attr_reader :sort_keys
 
         def sort_keys=(value) # :nodoc:
-          @sort_keys = Proc === value ? value : (value ? JSON::SORT_KEYS_PROC : false)
+          @sort_keys = case value
+          when Proc
+            value
+          when true
+            State.default_sort_keys_proc
+          when nil, false
+            false
+          else
+            raise TypeError, "The `sort_keys` argument must be a boolean or a Proc"
+          end
         end
 
         # :stopdoc:

test/json/json_generator_test.rb

@@ -240,7 +240,7 @@ def test_generate_sort_keys_with_proc
 
     # A truthy sort_keys is normalized to the default sorting proc.
     state = State.new(sort_keys: true)
-    assert_same JSON::SORT_KEYS_PROC, state.sort_keys
+    assert_instance_of Proc, state.sort_keys
   end
 
   def test_generate_custom