Test and restore behavior around to_json changing depth

ext/json/ext/generator/generator.c

@@ -1295,6 +1295,8 @@ static void generate_json_fallback(FBuffer *buffer, struct generate_json_data *d
     VALUE tmp;
     if (rb_respond_to(obj, i_to_json)) {
         tmp = rb_funcall(obj, i_to_json, 1, vstate_get(data));
+        GET_STATE(data->vstate);
+        data->depth = state->depth;
         Check_Type(tmp, T_STRING);
         fbuffer_append_str(buffer, tmp);
     } else {

java/src/json/ext/Generator.java

@@ -462,10 +462,10 @@ static void generateArray(ThreadContext context, Session session, RubyArray<IRub
             generateFor(context, session, element, buffer);
         }
 
-        int oldDepth = state.decreaseDepth();
+        state.depth = --depth;
         if (!arrayNLEmpty) {
             buffer.write(arrayNLBytes, arrayNLBegin, arrayNLSize);
-            Utils.repeatWrite(buffer, indentUnit, oldDepth);
+            Utils.repeatWrite(buffer, indentUnit, depth);
         }
 
         buffer.write((byte) ']');

java/src/json/ext/GeneratorState.java

@@ -107,7 +107,7 @@ public class GeneratorState extends RubyObject {
     /**
      * The current depth (inside a #to_json call)
      */
-    private int depth = 0;
+    protected int depth = 0;
 
     static final ObjectAllocator ALLOCATOR = GeneratorState::new;
 

lib/json/truffle_ruby/generator.rb

@@ -626,7 +626,7 @@ def json_transform(state)
               end
               first = false
             }
-            depth = state.depth -= 1
+            state.depth = depth -= 1
             result << state.array_nl
             result << state.indent * depth if indent
             result << ']'

test/json/json_generator_test.rb

@@ -321,6 +321,32 @@ def test_allow_nan
     end
   end
 
+  def test_depth_bad_to_json
+    obj = Object.new
+    def obj.to_json(state)
+      state.depth += 1
+      "{#{state.object_nl}"\
+        "#{state.indent * state.depth}\"foo\":#{state.space}1#{state.object_nl}"\
+        "#{state.indent * (state.depth - 1)}}"
+    end
+    assert_equal <<~JSON.chomp, JSON.pretty_generate([obj] * 2)
+      [
+        {
+          "foo": 1
+        },
+        {
+            "foo": 1
+          }
+      ]
+    JSON
+    state = JSON::State.new(object_nl: "\n", array_nl: "\n", space: " ", indent: "  ")
+    state.generate(obj)
+    assert_equal 1, state.depth # FIXME
+    state.depth = 0
+    state.generate([obj])
+    assert_equal 0, state.depth
+  end
+
   def test_depth
     pretty = { object_nl: "\n", array_nl: "\n", space: " ", indent: "  " }
     state = JSON.state.new(**pretty)