From 755114c1a35ad93d83def5e69dd6ce370e8070b8 Mon Sep 17 00:00:00 2001
From: David Renshaw <dwrenshaw@gmail.com>
Date: Sat, 4 May 2019 10:33:55 -0400
Subject: [PATCH] implement scratch space reuse

---
 .../main/java/org/capnproto/BuilderArena.java | 12 +++++++++
 .../java/org/capnproto/MessageBuilder.java    | 27 +++++++++++++++++++
 .../java/org/capnproto/SegmentBuilder.java    |  7 +++++
 3 files changed, 46 insertions(+)

diff --git a/runtime/src/main/java/org/capnproto/BuilderArena.java b/runtime/src/main/java/org/capnproto/BuilderArena.java
index 53f6851..37dc774 100644
--- a/runtime/src/main/java/org/capnproto/BuilderArena.java
+++ b/runtime/src/main/java/org/capnproto/BuilderArena.java
@@ -52,6 +52,18 @@ public final class BuilderArena implements Arena {
         this.allocator = allocator;
     }
 
+    public BuilderArena(Allocator allocator, ByteBuffer firstSegment) {
+        this.segments = new ArrayList<SegmentBuilder>();
+        SegmentBuilder newSegment = new SegmentBuilder(
+            firstSegment,
+            this);
+        newSegment.buffer.order(ByteOrder.LITTLE_ENDIAN);
+        newSegment.id = 0;
+        this.segments.add(newSegment);
+
+        this.allocator = allocator;
+    }
+
     @Override
     public final SegmentReader tryGetSegment(int id) {
         return this.segments.get(id);
diff --git a/runtime/src/main/java/org/capnproto/MessageBuilder.java b/runtime/src/main/java/org/capnproto/MessageBuilder.java
index b3b46d9..0dd2282 100644
--- a/runtime/src/main/java/org/capnproto/MessageBuilder.java
+++ b/runtime/src/main/java/org/capnproto/MessageBuilder.java
@@ -40,10 +40,26 @@ public final class MessageBuilder {
                                       allocationStrategy);
     }
 
+    /**
+     * Constructs a new MessageBuilder from an Allocator.
+     */
     public MessageBuilder(Allocator allocator) {
         this.arena = new BuilderArena(allocator);
     }
 
+    /**
+     * Constructs a new MessageBuilder from an Allocator and a given first segment buffer.
+     * This is useful for reusing the first segment buffer between messages, to avoid
+     * repeated allocations.
+     *
+     * You MUST ensure that firstSegment contains only zeroes before calling this method.
+     * If you are reusing firstSegment from another message, then it suffices to call
+     * clearFirstSegment() on that message.
+     */
+    public MessageBuilder(Allocator allocator, java.nio.ByteBuffer firstSegment) {
+        this.arena = new BuilderArena(allocator, firstSegment);
+    }
+
     private AnyPointer.Builder getRootInternal() {
         if (this.arena.segments.isEmpty()) {
             this.arena.allocate(1);
@@ -78,4 +94,15 @@ public final class MessageBuilder {
     public final java.nio.ByteBuffer[] getSegmentsForOutput() {
         return this.arena.getSegmentsForOutput();
     }
+
+    /**
+     * Sets the first segment buffer to contain all zeros so that it can be reused in
+     * another message. (See the MessageBuilder(Allocator, ByteBuffer) constructor above.)
+     *
+     * After calling this method, the message will be corrupted. Therefore, you need to make
+     * sure to write the message (via getSegmentsForOutput()) before calling this.
+     */
+    public final void clearFirstSegment() {
+        this.arena.segments.get(0).clear();
+    }
 }
diff --git a/runtime/src/main/java/org/capnproto/SegmentBuilder.java b/runtime/src/main/java/org/capnproto/SegmentBuilder.java
index 4aec968..e9ac57a 100644
--- a/runtime/src/main/java/org/capnproto/SegmentBuilder.java
+++ b/runtime/src/main/java/org/capnproto/SegmentBuilder.java
@@ -72,4 +72,11 @@ public final class SegmentBuilder extends SegmentReader {
     public final void put(int index, long value) {
         buffer.putLong(index * Constants.BYTES_PER_WORD, value);
     }
+
+    public final void clear() {
+        for (int ii = 0; ii <  this.pos; ++ii) {
+            this.put(ii, 0);
+        }
+        this.pos = 0;
+    }
 }