Allow for segmented output streams and a SWAR-based basic StringEncoder implementation.

samyron · byroot · commit 5274d5d19f8a · 2025-08-28T20:03:00.000+02:00
diff --git a/Rakefile b/Rakefile
@@ -116,7 +116,7 @@ if defined?(RUBY_ENGINE) and RUBY_ENGINE == 'jruby'
   file JRUBY_GENERATOR_JAR => :compile do
     cd 'java/src' do
       generator_classes = FileList[
-        "json/ext/ByteList*.class",
+        "json/ext/*ByteList*.class",
         "json/ext/OptionsReader*.class",
         "json/ext/Generator*.class",
         "json/ext/RuntimeInfo*.class",
diff --git a/java/src/json/ext/AbstractByteListDirectOutputStream.java b/java/src/json/ext/AbstractByteListDirectOutputStream.java
@@ -0,0 +1,31 @@
+package json.ext;
+
+import java.io.OutputStream;
+
+import org.jcodings.Encoding;
+import org.jruby.util.ByteList;
+
+abstract class AbstractByteListDirectOutputStream extends OutputStream {
+
+    private static final String PROP_SEGMENTED_BUFFER = "json.useSegmentedOutputStream";
+    private static final String PROP_SEGMENTED_BUFFER_DEFAULT = "true";
+
+    private static final boolean USE_SEGMENTED_BUFFER;
+
+    static {
+        String useSegmentedOutputStream = System.getProperty(PROP_SEGMENTED_BUFFER, PROP_SEGMENTED_BUFFER_DEFAULT);
+        USE_SEGMENTED_BUFFER = Boolean.parseBoolean(useSegmentedOutputStream);
+        // XXX Is there a logger we can use here?
+        // System.out.println("Using segmented output stream: " + USE_SEGMENTED_BUFFER);
+    }
+
+    public static AbstractByteListDirectOutputStream create(int estimatedSize) {
+        if (USE_SEGMENTED_BUFFER) {
+            return new SegmentedByteListDirectOutputStream(estimatedSize);
+        } else {
+            return new ByteListDirectOutputStream(estimatedSize);
+        }
+    }
+
+    public abstract ByteList toByteListDirect(Encoding encoding);
+}
diff --git a/java/src/json/ext/ByteListDirectOutputStream.java b/java/src/json/ext/ByteListDirectOutputStream.java
@@ -4,10 +4,9 @@
 import org.jruby.util.ByteList;
 
 import java.io.IOException;
-import java.io.OutputStream;
 import java.util.Arrays;
 
-public class ByteListDirectOutputStream extends OutputStream {
+public class ByteListDirectOutputStream extends AbstractByteListDirectOutputStream {
     private byte[] buffer;
     private int length;
 
diff --git a/java/src/json/ext/Generator.java b/java/src/json/ext/Generator.java
@@ -26,6 +26,8 @@
 import org.jruby.util.IOOutputStream;
 import org.jruby.util.TypeConverter;
 
+import json.ext.ByteListDirectOutputStream;
+
 import java.io.BufferedOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
@@ -252,7 +254,7 @@ int guessSize(ThreadContext context, Session session, T object) {
         }
 
         RubyString generateNew(ThreadContext context, Session session, T object) {
-            ByteListDirectOutputStream buffer = new ByteListDirectOutputStream(guessSize(context, session, object));
+            AbstractByteListDirectOutputStream buffer = AbstractByteListDirectOutputStream.create(guessSize(context, session, object));
             generateToBuffer(context, session, object, buffer);
             return RubyString.newString(context.runtime, buffer.toByteListDirect(UTF8Encoding.INSTANCE));
         }
diff --git a/java/src/json/ext/LinkedSegmentedByteListDirectOutputStream.java b/java/src/json/ext/LinkedSegmentedByteListDirectOutputStream.java
@@ -0,0 +1,106 @@
+package json.ext;
+
+import org.jcodings.Encoding;
+import org.jruby.util.ByteList;
+
+import java.io.IOException;
+
+public class LinkedSegmentedByteListDirectOutputStream extends AbstractByteListDirectOutputStream {
+    private Segment head;
+    private int length;
+    private Segment current;
+    private int numSegments;
+
+    private static class Segment {
+        static final int DEFAULT_SEGMENT_SIZE = 1024;
+        byte[] buffer;
+        int length;
+        Segment next;
+
+        Segment() {
+            this(DEFAULT_SEGMENT_SIZE);
+        }
+
+        Segment(int size) {
+            if (size <= 0) {
+                size = DEFAULT_SEGMENT_SIZE;
+            }
+            buffer = new byte[Math.max(size, DEFAULT_SEGMENT_SIZE)];
+        }
+    }
+
+    LinkedSegmentedByteListDirectOutputStream() {
+        this(Segment.DEFAULT_SEGMENT_SIZE);
+    }
+
+    LinkedSegmentedByteListDirectOutputStream(int size) {
+        current = head = new Segment(size);
+    }
+
+    public ByteList toByteListDirect(Encoding encoding) {
+        byte[] buffer = new byte[length];
+        Segment segment = head;
+        int pos = 0;
+        while (segment != null) {
+            System.arraycopy(segment.buffer, 0, buffer, pos, segment.length);
+            pos += segment.length;
+            segment = segment.next;
+        }
+        return new ByteList(buffer, 0, length, encoding, false);
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        Segment c = current;
+        if (c.length == c.buffer.length) {
+            // This check is deliberately in the case the current segment is full. We want to 
+            // avoid this check in the common case where we have space in the current segment.
+            if (this.length + 1 < 0) {
+                throw new IOException("Total length exceeds maximum length of an array.");
+            }
+            if (c.next == null) {
+                numSegments++;
+                c.next = new Segment(c.buffer.length * 2);
+            }
+            c = c.next;
+            current = c;
+        }
+        c.buffer[c.length++] = (byte)b;
+        length++;
+    }
+
+    @Override
+    public void write(byte[] bytes, int start, int length) throws IOException {
+        Segment c = current;
+        int remaining = length;
+
+        while (remaining > 0) {
+            if (c.length == c.buffer.length) {
+                // This check is deliberately in the case the current segment is full. We want to 
+                // avoid this check in the common case where we have space in the current segment.
+                if (this.length + remaining < 0) {
+                    throw new IOException("Total length exceeds maximum length of an array.");
+                }
+                if (c.next == null) {
+                    numSegments++;
+                    c.next = new Segment(c.buffer.length * 2);
+                }
+                c = c.next;
+                current = c;
+            }
+            int currentLength = c.length;
+            int currentCapacity = c.buffer.length;
+            int copyLength = Math.min(remaining, currentCapacity - currentLength);
+            System.arraycopy(bytes, start, c.buffer, currentLength, copyLength);
+            c.length += copyLength;
+            this.length += copyLength;
+            start += copyLength;
+            remaining -= copyLength;
+        }
+    }
+
+    @Override
+    public void write(byte[] bytes) throws IOException {
+        write(bytes, 0, bytes.length);
+    }
+}
diff --git a/java/src/json/ext/SegmentedByteListDirectOutputStream.java b/java/src/json/ext/SegmentedByteListDirectOutputStream.java
@@ -0,0 +1,79 @@
+package json.ext;
+
+import org.jcodings.Encoding;
+import org.jruby.util.ByteList;
+
+import java.io.IOException;
+
+public class SegmentedByteListDirectOutputStream extends AbstractByteListDirectOutputStream {
+    private static final int DEFAULT_CAPACITY = 1024;
+
+    private int totalLength;
+    private byte[][] segments = new byte[21][];
+    private int currentSegmentIndex;
+    private int currentSegmentLength;
+    private byte[] currentSegment;
+
+    SegmentedByteListDirectOutputStream(int size) {
+        currentSegment = new byte[Math.max(size, DEFAULT_CAPACITY)];
+        segments[0] = currentSegment;
+    }
+
+    public ByteList toByteListDirect(Encoding encoding) {
+        byte[] buffer = new byte[totalLength];
+        int pos = 0;
+        // We handle the current segment separately.
+        for (int i = 0; i < currentSegmentIndex; i++) {
+            byte[] segment = segments[i];
+            System.arraycopy(segment, 0, buffer, pos, segment.length);
+            pos += segment.length;
+        }
+        System.arraycopy(currentSegment, 0, buffer, pos, currentSegmentLength);
+        return new ByteList(buffer, 0, totalLength, encoding, false);
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        if (currentSegmentLength == currentSegment.length) {
+            if (totalLength + 1 < 0) {
+                throw new IOException("Total length exceeds maximum length of an array.");
+            }
+            currentSegmentIndex++;
+            int capacity = currentSegment.length * 2;
+            currentSegment = new byte[capacity];
+            currentSegmentLength = 0;
+            segments[currentSegmentIndex] = currentSegment;
+        }
+        currentSegment[currentSegmentLength++] = (byte) b;
+        totalLength++;
+    }
+
+    @Override
+    public void write(byte[] bytes, int start, int length) throws IOException {
+        int remaining = length;
+
+        while (remaining > 0) {
+            if (currentSegmentLength == currentSegment.length) {
+                if (totalLength + remaining < 0) {
+                    throw new IOException("Total length exceeds maximum length of an array.");
+                }
+                currentSegmentIndex++;
+                int capacity = currentSegment.length * 2;
+                currentSegment = new byte[capacity];
+                currentSegmentLength = 0;
+                segments[currentSegmentIndex] = currentSegment;
+            }
+            int toWrite = Math.min(remaining, currentSegment.length - currentSegmentLength);
+            System.arraycopy(bytes, start, currentSegment, currentSegmentLength, toWrite);
+            currentSegmentLength += toWrite;
+            start += toWrite;
+            remaining -= toWrite;
+        }
+        totalLength += length;
+    }
+
+    @Override
+    public void write(byte[] bytes) throws IOException {
+        write(bytes, 0, bytes.length);
+    }
+}
diff --git a/java/src/json/ext/StringEncoder.java b/java/src/json/ext/StringEncoder.java
