src/jdk.internal.vm.compiler/share/classes/org.graalvm.graphio/src/org/graalvm/graphio/GraphProtocol.java

@@ -76,61 +76,98 @@
     private static final byte[] MAGIC_BYTES = {'B', 'I', 'G', 'V'};
 
     private final ConstantPool constantPool;
     private final ByteBuffer buffer;
     private final WritableByteChannel channel;
+    private final boolean embedded;
     final int versionMajor;
     final int versionMinor;
+    private boolean printing;
 
-    GraphProtocol(WritableByteChannel channel, int major, int minor) throws IOException {
+    GraphProtocol(WritableByteChannel channel, int major, int minor, boolean embedded) throws IOException {
         if (major > 6 || (major == 6 && minor > 0)) {
             throw new IllegalArgumentException("Unrecognized version " + major + "." + minor);
         }
         this.versionMajor = major;
         this.versionMinor = minor;
         this.constantPool = new ConstantPool();
         this.buffer = ByteBuffer.allocateDirect(256 * 1024);
         this.channel = channel;
-        writeVersion();
+        this.embedded = embedded;
+        if (!embedded) {
+            writeVersion();
+            flushEmbedded();
+        }
     }
 
     GraphProtocol(GraphProtocol<?, ?, ?, ?, ?, ?, ?, ?, ?, ?> parent) {
         this.versionMajor = parent.versionMajor;
         this.versionMinor = parent.versionMinor;
         this.constantPool = parent.constantPool;
         this.buffer = parent.buffer;
         this.channel = parent.channel;
+        this.embedded = parent.embedded;
     }
 
     @SuppressWarnings("all")
     public final void print(Graph graph, Map<? extends Object, ? extends Object> properties, int id, String format, Object... args) throws IOException {
-        writeByte(BEGIN_GRAPH);
-        if (versionMajor >= 3) {
-            writeInt(id);
-            writeString(format);
-            writeInt(args.length);
-            for (Object a : args) {
-                writePropertyObject(graph, a);
+        printing = true;
+        try {
+            writeByte(BEGIN_GRAPH);
+            if (versionMajor >= 3) {
+                writeInt(id);
+                writeString(format);
+                writeInt(args.length);
+                for (Object a : args) {
+                    writePropertyObject(graph, a);
+                }
+            } else {
+                writePoolObject(formatTitle(graph, id, format, args));
             }
-        } else {
-            writePoolObject(formatTitle(graph, id, format, args));
+            writeGraph(graph, properties);
+            flushEmbedded();
+            flush();
+        } finally {
+            printing = false;
         }
-        writeGraph(graph, properties);
-        flush();
     }
 
     public final void beginGroup(Graph noGraph, String name, String shortName, ResolvedJavaMethod method, int bci, Map<? extends Object, ? extends Object> properties) throws IOException {
-        writeByte(BEGIN_GROUP);
-        writePoolObject(name);
-        writePoolObject(shortName);
-        writePoolObject(method);
-        writeInt(bci);
-        writeProperties(noGraph, properties);
+        printing = true;
+        try {
+            writeByte(BEGIN_GROUP);
+            writePoolObject(name);
+            writePoolObject(shortName);
+            writePoolObject(method);
+            writeInt(bci);
+            writeProperties(noGraph, properties);
+            flushEmbedded();
+        } finally {
+            printing = false;
+        }
     }
 
     public final void endGroup() throws IOException {
-        writeByte(CLOSE_GROUP);
+        printing = true;
+        try {
+            writeByte(CLOSE_GROUP);
+            flushEmbedded();
+        } finally {
+            printing = false;
+        }
+    }
+
+    final int write(ByteBuffer src) throws IOException {
+        if (printing) {
+            throw new IllegalStateException("Trying to write during graph print.");
+        }
+        constantPool.reset();
+        return writeBytesRaw(src);
+    }
+
+    final boolean isOpen() {
+        return channel.isOpen();
     }
 
     @Override
     public final void close() {
         try {

@@ -278,10 +315,17 @@
         writeBytesRaw(MAGIC_BYTES);
         writeByte(versionMajor);
         writeByte(versionMinor);
     }
 
+    private void flushEmbedded() throws IOException {
+        if (embedded) {
+            flush();
+            constantPool.reset();
+        }
+    }
+
     private void flush() throws IOException {
         buffer.flip();
         /*
          * Try not to let interrupted threads abort the write. There's still a race here but an
          * interrupt that's been pending for a long time shouldn't stop this writing.

@@ -356,10 +400,27 @@
             buffer.put(b, bytesWritten, toWrite);
             bytesWritten += toWrite;
         }
     }
 
+    private int writeBytesRaw(ByteBuffer b) throws IOException {
+        int limit = b.limit();
+        int written = 0;
+        while (b.position() < limit) {
+            int toWrite = Math.min(limit - b.position(), buffer.capacity());
+            ensureAvailable(toWrite);
+            b.limit(b.position() + toWrite);
+            try {
+                buffer.put(b);
+                written += toWrite;
+            } finally {
+                b.limit(limit);
+            }
+        }
+        return written;
+    }
+
     private void writeInts(int[] b) throws IOException {
         if (b == null) {
             writeInt(-1);
         } else {
             writeInt(b.length);

@@ -816,8 +877,14 @@
         public char add(Object obj) {
             Character id = nextAvailableId();
             put(obj, id);
             return id;
         }
+
+        void reset() {
+            clear();
+            availableIds.clear();
+            nextId = 0;
+        }
     }
 
 }