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

*** 76,136 ****
      private static final byte[] MAGIC_BYTES = {'B', 'I', 'G', 'V'};
  
      private final ConstantPool constantPool;
      private final ByteBuffer buffer;
      private final WritableByteChannel channel;
      final int versionMajor;
      final int versionMinor;
  
!     GraphProtocol(WritableByteChannel channel, int major, int minor) 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();
      }
  
      GraphProtocol(GraphProtocol<?, ?, ?, ?, ?, ?, ?, ?, ?, ?> parent) {
          this.versionMajor = parent.versionMajor;
          this.versionMinor = parent.versionMinor;
          this.constantPool = parent.constantPool;
          this.buffer = parent.buffer;
          this.channel = parent.channel;
      }
  
      @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);
              }
!         } else {
!             writePoolObject(formatTitle(graph, id, format, args));
          }
-         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);
      }
  
      public final void endGroup() throws IOException {
!         writeByte(CLOSE_GROUP);
      }
  
      @Override
      public final void close() {
          try {
--- 76,173 ----
      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, 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;
!         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 {
!         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));
              }
!             writeGraph(graph, properties);
!             flushEmbedded();
!             flush();
!         } finally {
!             printing = false;
          }
      }
  
      public final void beginGroup(Graph noGraph, String name, String shortName, ResolvedJavaMethod method, int bci, Map<? extends Object, ? extends Object> properties) throws IOException {
!         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 {
!         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,287 ****
--- 315,331 ----
          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,365 ****
--- 400,426 ----
              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,823 ****
--- 877,890 ----
          public char add(Object obj) {
              Character id = nextAvailableId();
              put(obj, id);
              return id;
          }
+ 
+         void reset() {
+             clear();
+             availableIds.clear();
+             nextId = 0;
+         }
      }
  
  }