Add a sort_keys option to the generator.

Fix: #976

Co-Authored-By: Jean Boussier <byroot@ruby-lang.org>

Author: samyron

ext/json/ext/generator/generator.c

@@ -34,13 +34,14 @@ typedef struct JSON_Generator_StateStruct {
     bool ascii_only;
     bool script_safe;
     bool strict;
+    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,
-             sym_ascii_only, sym_depth, sym_buffer_initial_length, sym_script_safe, sym_escape_slash, sym_strict, sym_as_json;
+             sym_ascii_only, sym_depth, sym_buffer_initial_length, sym_script_safe, sym_escape_slash, sym_strict, sym_as_json, sym_sort_keys;
 
 
 #define GET_STATE_TO(self, state) \
@@ -709,6 +710,7 @@ static void State_mark(void *ptr)
     rb_gc_mark_movable(state->object_nl);
     rb_gc_mark_movable(state->array_nl);
     rb_gc_mark_movable(state->as_json);
+    rb_gc_mark_movable(state->sort_keys);
 }
 
 static void State_compact(void *ptr)
@@ -720,6 +722,7 @@ static void State_compact(void *ptr)
     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);
+    state->sort_keys = rb_gc_location(state->sort_keys);
 }
 
 static size_t State_memsize(const void *ptr)
@@ -769,6 +772,7 @@ static void vstate_spill(struct generate_json_data *data)
     RB_OBJ_WRITTEN(vstate, Qundef, state->object_nl);
     RB_OBJ_WRITTEN(vstate, Qundef, state->array_nl);
     RB_OBJ_WRITTEN(vstate, Qundef, state->as_json);
+    RB_OBJ_WRITTEN(vstate, Qundef, state->sort_keys);
 }
 
 static inline VALUE json_call_to_json(struct generate_json_data *data, VALUE obj)
@@ -1050,6 +1054,11 @@ static inline long increase_depth(struct generate_json_data *data)
 
 static void generate_json_object(FBuffer *buffer, struct generate_json_data *data, VALUE obj)
 {
+    if (RB_UNLIKELY(data->state->sort_keys)) {
+        obj = rb_proc_call_with_block(data->state->sort_keys, 1, &obj, Qnil);
+        Check_Type(obj, T_HASH);
+    }
+
     long depth = increase_depth(data);
 
     if (RHASH_SIZE(obj) == 0) {
@@ -1376,6 +1385,7 @@ static VALUE cState_init_copy(VALUE obj, VALUE orig)
     RB_OBJ_WRITTEN(obj, Qundef, objState->object_nl);
     RB_OBJ_WRITTEN(obj, Qundef, objState->array_nl);
     RB_OBJ_WRITTEN(obj, Qundef, objState->as_json);
+    RB_OBJ_WRITTEN(obj, Qundef, objState->sort_keys);
 
     return obj;
 }
@@ -1722,6 +1732,55 @@ 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;
+    }
+}
+
+/*
+ * call-seq: sort_keys
+ *
+ * Get the value of sort_keys.
+ */
+static VALUE cState_sort_keys_p(VALUE self)
+{
+    GET_STATE(self);
+    return state->sort_keys;
+}
+
+/*
+ * call-seq: sort_keys=(value)
+ *
+ * value is a boolean or a proc. If the value is the boolean true, object keys
+ * will be sorted lexicographically in ascending order.
+ *
+ * If the value is a proc, it receives the entire Hash and must return a Hash
+ * with its pairs in the desired order, allowing for arbitrary sorting.
+ */
+static VALUE cState_sort_keys_set(VALUE self, VALUE value)
+{
+    rb_check_frozen(self);
+    GET_STATE(self);
+    RB_OBJ_WRITE(self, &state->sort_keys, normalize_sort_keys(value));
+    return Qnil;
+}
+
 static VALUE cState_allow_duplicate_key_p(VALUE self)
 {
     GET_STATE(self);
@@ -1832,6 +1891,9 @@ static int configure_state_i(VALUE key, VALUE val, VALUE _arg)
         state->as_json_single_arg = proc && rb_proc_arity(proc) == 1;
         state_write_value(data, &state->as_json, proc);
     }
+    else if (key == sym_sort_keys)             {
+        state_write_value(data, &state->sort_keys, normalize_sort_keys(val));
+    }
     return ST_CONTINUE;
 }
 
@@ -1909,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");
 
@@ -1918,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);
@@ -1957,6 +2023,8 @@ void Init_generator(void)
     rb_define_method(cState, "buffer_initial_length=", cState_buffer_initial_length_set, 1);
     rb_define_method(cState, "generate", cState_generate, -1);
     rb_define_method(cState, "_generate_no_fallback", cState_generate_no_fallback, -1);
+    rb_define_method(cState, "sort_keys", cState_sort_keys_p, 0);
+    rb_define_method(cState, "sort_keys=", cState_sort_keys_set, 1);
 
     rb_define_private_method(cState, "allow_duplicate_key?", cState_allow_duplicate_key_p, 0);
 
@@ -1986,6 +2054,7 @@ void Init_generator(void)
     sym_strict = ID2SYM(rb_intern("strict"));
     sym_as_json = ID2SYM(rb_intern("as_json"));
     sym_allow_duplicate_key = ID2SYM(rb_intern("allow_duplicate_key"));
+    sym_sort_keys = ID2SYM(rb_intern("sort_keys"));
 
     usascii_encindex = rb_usascii_encindex();
     utf8_encindex = rb_utf8_encindex();

java/src/json/ext/Generator.java

@@ -17,6 +17,7 @@
 import org.jruby.RubyFloat;
 import org.jruby.RubyHash;
 import org.jruby.RubyIO;
+import org.jruby.RubyProc;
 import org.jruby.RubyString;
 import org.jruby.RubySymbol;
 import org.jruby.runtime.Helpers;
@@ -573,6 +574,11 @@ static void generateHash(ThreadContext context, Session session, RubyHash object
             return;
         }
 
+        RubyProc sortKeysProc = state.getSortKeysProc();
+        if (sortKeysProc != null) {
+            object = (RubyHash) Helpers.invoke(context, sortKeysProc, "call", object);
+        }
+
         final ByteList objectNl = state.getObjectNl();
         byte[] objectNLBytes = objectNl.unsafeBytes();
         final byte[] indent = Utils.repeat(state.getIndent(), depth);

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.
@@ -104,6 +106,13 @@ public class GeneratorState extends RubyObject {
     private int bufferInitialLength = DEFAULT_BUFFER_INITIAL_LENGTH;
     static final int DEFAULT_BUFFER_INITIAL_LENGTH = 1024;
 
+    /**
+     * Controls key sorting when generating JSON. <code>null</code> means keys
+     * are emitted in insertion order; a true value sorts keys lexicographically;
+     * a {@link RubyProc} is used as a comparator receiving two [key, value] pairs.
+     */
+    private IRubyObject sortKeys;
+
     /**
      * The current depth (inside a #to_json call)
      */
@@ -158,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>
@@ -222,6 +237,7 @@ public IRubyObject initialize_copy(ThreadContext context, IRubyObject vOrig) {
 
         this.allowDuplicateKey = orig.allowDuplicateKey;
         this.deprecateDuplicateKey = orig.deprecateDuplicateKey;
+        this.sortKeys = orig.sortKeys;
 
         return this;
     }
@@ -431,6 +447,23 @@ public boolean strict() {
         return strict;
     }
 
+    /**
+     * Returns the proc used to sort the keys of an object, or
+     * <code>null</code> if keys should not be sorted. The proc receives the
+     * entire Hash and returns a Hash with its pairs in the desired order.
+     */
+    public RubyProc getSortKeysProc() {
+        return sortKeys instanceof RubyProc ? (RubyProc) sortKeys : null;
+    }
+
+    private static IRubyObject normalizeSortKeys(ThreadContext context, IRubyObject value) {
+        if (value instanceof RubyProc) return value;
+        if (value != null && value.isTrue()) {
+            return defaultSortKeyProc;
+        }
+        return null;
+    }
+
     @JRubyMethod(name={"strict","strict?"})
     public RubyBoolean strict_get(ThreadContext context) {
         return RubyBoolean.newBoolean(context, strict);
@@ -474,6 +507,18 @@ public IRubyObject buffer_initial_length_set(IRubyObject buffer_initial_length)
         return buffer_initial_length;
     }
 
+    @JRubyMethod(name="sort_keys")
+    public IRubyObject sort_keys_get(ThreadContext context) {
+        return sortKeys == null ? context.getRuntime().getFalse() : sortKeys;
+    }
+
+    @JRubyMethod(name="sort_keys=")
+    public IRubyObject sort_keys_set(ThreadContext context, IRubyObject sortKeys) {
+        checkFrozen();
+        this.sortKeys = normalizeSortKeys(context, sortKeys);
+        return sortKeys;
+    }
+
     public int getDepth() {
         return depth;
     }
@@ -568,6 +613,9 @@ public IRubyObject _configure(ThreadContext context, IRubyObject vOpts) {
             this.allowDuplicateKey = opts.getBool("allow_duplicate_key", false);
             this.deprecateDuplicateKey = false;
         }
+
+        sortKeys = normalizeSortKeys(context, opts.get("sort_keys"));
+
         return this;
     }
 
@@ -596,6 +644,7 @@ public RubyHash to_h(ThreadContext context) {
         result.op_aset(context, runtime.newSymbol("strict"), strict_get(context));
         result.op_aset(context, runtime.newSymbol("depth"), depth_get(context));
         result.op_aset(context, runtime.newSymbol("buffer_initial_length"), buffer_initial_length_get(context));
+        result.op_aset(context, runtime.newSymbol("sort_keys"), sort_keys_get(context));
 
         if (this.allowDuplicateKey) {
             if (!this.deprecateDuplicateKey) {

lib/json.rb

@@ -408,14 +408,18 @@
 #   to be inserted after each \JSON object; defaults to the empty \String, <tt>''</tt>.
 # - Option +indent+ (\String) specifies the string (usually spaces) to be
 #   used for indentation; defaults to the empty \String, <tt>''</tt>;
-#   defaults to the empty \String, <tt>''</tt>;
 #   has no effect unless options +array_nl+ or +object_nl+ specify newlines.
 # - Option +space+ (\String) specifies a string (usually a space) to be
 #   inserted after the colon in each \JSON object's pair;
 #   defaults to the empty \String, <tt>''</tt>.
 # - Option +space_before+ (\String) specifies a string (usually a space) to be
 #   inserted before the colon in each \JSON object's pair;
 #   defaults to the empty \String, <tt>''</tt>.
+# - Option +sort_keys+ (boolean or \Proc) controls whether and how the keys of a
+#   hash are sorted when generating the output; defaults to <tt>false</tt>.
+#   When +true+, keys are sorted lexicographically. When a \Proc, it receives
+#   the entire \Hash and must return a \Hash with its pairs in the desired
+#   order, allowing for arbitrary sort orders.
 #
 # In this example, +obj+ is used first to generate the shortest
 # \JSON data (no whitespace), then again with all formatting options

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

lib/json/ext/generator/state.rb

@@ -54,6 +54,7 @@ def to_h
             strict: strict?,
             depth: depth,
             buffer_initial_length: buffer_initial_length,
+            sort_keys: sort_keys
           }
 
           allow_duplicate_key = allow_duplicate_key?