From 982fb2e5d7e9f1515a18e4c2edb1cab4995ed5ab Mon Sep 17 00:00:00 2001
From: Charles Oliver Nutter <headius@headius.com>
Date: Tue, 28 May 2024 13:07:09 -0500
Subject: [PATCH] Capture shared variable iter state immediately

If not captured immediately, contained objects will be dumped and
clear this state, breaking the variable count just acquired.
---
 .../org/jruby/runtime/marshal/NewMarshal.java | 72 +++++++++----------
 1 file changed, 33 insertions(+), 39 deletions(-)

diff --git a/core/src/main/java/org/jruby/runtime/marshal/NewMarshal.java b/core/src/main/java/org/jruby/runtime/marshal/NewMarshal.java
index 28b12647b54..f92e7df32ea 100644
--- a/core/src/main/java/org/jruby/runtime/marshal/NewMarshal.java
+++ b/core/src/main/java/org/jruby/runtime/marshal/NewMarshal.java
@@ -212,31 +212,29 @@ public void writeDirectly(ThreadContext context, RubyOutputStream out, IRubyObje
 
         } else {
 
-            try {
-                var ivarAccessors = checkVariables(value, shouldMarshalEncoding);
-
-                if (doVariables) {
-                    // object has instance vars and isn't a class, get a snapshot to be marshalled
-                    // and output the ivar header here
+            var ivarAccessors = checkVariables(value, shouldMarshalEncoding);
+            int size = variableCount;
+            clearVariableState();
 
-                    // write `I' instance var signet if class is NOT a direct subclass of Object
-                    out.write(TYPE_IVAR);
-                    dumpBaseObject(context, out, value, nativeClassIndex);
+            if (doVariables) {
+                // object has instance vars and isn't a class, get a snapshot to be marshalled
+                // and output the ivar header here
 
-                    if (shouldMarshalEncoding) {
-                        writeInt(out, variableCount + 1); // vars preceded by encoding
-                        writeEncoding(context, out, ((MarshalEncoding) value).getMarshalEncoding());
-                    } else {
-                        writeInt(out, variableCount);
-                    }
+                // write `I' instance var signet if class is NOT a direct subclass of Object
+                out.write(TYPE_IVAR);
+                dumpBaseObject(context, out, value, nativeClassIndex);
 
-                    ivarAccessors.forEach(new VariableDumper(context, out, value));
+                if (shouldMarshalEncoding) {
+                    writeInt(out, size + 1); // vars preceded by encoding
+                    writeEncoding(context, out, ((MarshalEncoding) value).getMarshalEncoding());
                 } else {
-                    // no variables, no encoding
-                    dumpBaseObject(context, out, value, nativeClassIndex);
+                    writeInt(out, size);
                 }
-            } finally {
-                clearVariableState();
+
+                ivarAccessors.forEach(new VariableDumper(context, out, value));
+            } else {
+                // no variables, no encoding
+                dumpBaseObject(context, out, value, nativeClassIndex);
             }
 
         }
@@ -458,21 +456,19 @@ private void userCommon(ThreadContext context, RubyOutputStream out, IRubyObject
         RubyString marshaled = castToString(context, dumpResult);
 
         if (marshaled.hasVariables()) {
-            try {
-                var ivarAccessors = countVariables(marshaled);
+            var ivarAccessors = countVariables(marshaled);
+            int size = variableCount;
+            clearVariableState();
 
-                if (variableCount > 0) {
-                    out.write(TYPE_IVAR);
-                    dumpUserdefBase(out, runtime, klass, marshaled);
+            if (size > 0) {
+                out.write(TYPE_IVAR);
+                dumpUserdefBase(out, runtime, klass, marshaled);
 
-                    writeInt(out, variableCount);
+                writeInt(out, size);
 
-                    ivarAccessors.forEach(new VariableDumper(context, out, marshaled));
-                } else {
-                    dumpUserdefBase(out, runtime, klass, marshaled);
-                }
-            } finally {
-                clearVariableState();
+                ivarAccessors.forEach(new VariableDumper(context, out, marshaled));
+            } else {
+                dumpUserdefBase(out, runtime, klass, marshaled);
             }
         } else {
             dumpUserdefBase(out, runtime, klass, marshaled);
@@ -506,15 +502,13 @@ public void dumpVariables(ThreadContext context, RubyOutputStream out, IRubyObje
     }
 
     public void dumpVariables(ThreadContext context, RubyOutputStream out, IRubyObject value, int extraSize) {
-        try {
-            Map<String, VariableAccessor> ivarAccessors = countVariables(value, extraSize);
+        Map<String, VariableAccessor> ivarAccessors = countVariables(value, extraSize);
+        int size = variableCount;
+        clearVariableState();
 
-            writeInt(out, variableCount);
+        writeInt(out, size);
 
-            ivarAccessors.forEach(new VariableDumper(context, out, value));
-        } finally {
-            clearVariableState();
-        }
+        ivarAccessors.forEach(new VariableDumper(context, out, value));
     }
 
     private Map<String, VariableAccessor> countVariables(IRubyObject value) {