support custom allocators for builder segments

runtime/src/main/java/org/capnproto/Allocator.java

@@ -0,0 +1,13 @@
+package org.capnproto;
+
+/**
+ * An object that allocates memory for a Cap'n Proto message as it is being built.
+ */
+interface Allocator {
+    /**
+     * Allocates a ByteBuffer to be used as a segment in a message. The returned
+     * buffer must contain at least `minimumSize` bytes, all of which MUST be
+     * set to zero.
+     */
+   public java.nio.ByteBuffer allocateSegment(int minimumSize);
+}

runtime/src/main/java/org/capnproto/BuilderArena.java

@@ -26,7 +26,6 @@ import java.nio.ByteOrder;
 import java.util.ArrayList;
 
 public final class BuilderArena implements Arena {
-
     public enum AllocationStrategy {
         FIXED_SIZE,
         GROW_HEURISTICALLY
@@ -37,19 +36,20 @@ public final class BuilderArena implements Arena {
         AllocationStrategy.GROW_HEURISTICALLY;
 
     public final ArrayList<SegmentBuilder> segments;
-
-    public int nextSize;
-    public final AllocationStrategy allocationStrategy;
-
+    private Allocator allocator;
 
     public BuilderArena(int firstSegmentSizeWords, AllocationStrategy allocationStrategy) {
         this.segments = new ArrayList<SegmentBuilder>();
-        this.nextSize = firstSegmentSizeWords;
-        this.allocationStrategy = allocationStrategy;
-        SegmentBuilder segment0 = new SegmentBuilder(
-            ByteBuffer.allocate(firstSegmentSizeWords * Constants.BYTES_PER_WORD), this);
-        segment0.buffer.order(ByteOrder.LITTLE_ENDIAN);
-        this.segments.add(segment0);
+        {
+            DefaultAllocator allocator = new DefaultAllocator(allocationStrategy);
+            allocator.setNextAllocationSizeBytes(firstSegmentSizeWords * Constants.BYTES_PER_WORD);
+            this.allocator = allocator;
+        }
+    }
+
+    public BuilderArena(Allocator allocator) {
+        this.segments = new ArrayList<SegmentBuilder>();
+        this.allocator = allocator;
     }
 
     @Override
@@ -76,32 +76,19 @@ public final class BuilderArena implements Arena {
     }
 
     public AllocateResult allocate(int amount) {
-
         int len = this.segments.size();
+
         // we allocate the first segment in the constructor.
-
-        int result = this.segments.get(len - 1).allocate(amount);
-        if (result != SegmentBuilder.FAILED_ALLOCATION) {
-            return new AllocateResult(this.segments.get(len - 1), result);
+        if (len > 0) {
+            int result = this.segments.get(len - 1).allocate(amount);
+            if (result != SegmentBuilder.FAILED_ALLOCATION) {
+                return new AllocateResult(this.segments.get(len - 1), result);
+            }
         }
-
-        // allocate_owned_memory
-
-        int size = Math.max(amount, this.nextSize);
         SegmentBuilder newSegment = new SegmentBuilder(
-            ByteBuffer.allocate(size * Constants.BYTES_PER_WORD),
+            this.allocator.allocateSegment(amount * Constants.BYTES_PER_WORD),
             this);
 
-        switch (this.allocationStrategy) {
-        case GROW_HEURISTICALLY:
-            this.nextSize += size;
-            break;
-        default:
-            break;
-        }
-
-        // --------
-
         newSegment.buffer.order(ByteOrder.LITTLE_ENDIAN);
         newSegment.id = len;
         this.segments.add(newSegment);

runtime/src/main/java/org/capnproto/DefaultAllocator.java

@@ -0,0 +1,65 @@
+package org.capnproto;
+
+import java.nio.ByteBuffer;
+
+import org.capnproto.BuilderArena.AllocationStrategy;
+
+class DefaultAllocator implements Allocator {
+
+    // (minimum) number of bytes in the next allocation
+    private int nextSize = 0;
+
+    public enum ByteBufferAllocationStyle {
+        REGULAR,
+        DIRECT
+    }
+    public ByteBufferAllocationStyle allocationStyle = ByteBufferAllocationStyle.REGULAR;
+
+    public AllocationStrategy allocationStrategy =
+        AllocationStrategy.GROW_HEURISTICALLY;
+
+    public DefaultAllocator(AllocationStrategy allocationStrategy) {
+        this.allocationStrategy = allocationStrategy;
+    }
+
+    public DefaultAllocator(ByteBufferAllocationStyle style) {
+        this.allocationStyle = style;
+    }
+
+    public DefaultAllocator(AllocationStrategy allocationStrategy,
+                            ByteBufferAllocationStyle style) {
+        this.allocationStrategy = allocationStrategy;
+        this.allocationStyle = style;
+    }
+
+    public void setNextAllocationSizeBytes(int nextSize) {
+        this.nextSize = nextSize;
+    }
+
+    @Override
+    public java.nio.ByteBuffer allocateSegment(int minimumSize) {
+        int size = Math.max(minimumSize, this.nextSize);
+        ByteBuffer result = null;
+        switch (allocationStyle) {
+          case REGULAR:
+            result = ByteBuffer.allocate(size);
+            break;
+          case DIRECT:
+            result = ByteBuffer.allocateDirect(size);
+        }
+
+        switch (this.allocationStrategy) {
+            case GROW_HEURISTICALLY:
+                this.nextSize += size;
+                break;
+            case FIXED_SIZE:
+                if (nextSize == 0) {
+                    nextSize = size;
+                }
+                break;
+        }
+
+        this.nextSize += size;
+        return result;
+    }
+}

runtime/src/main/java/org/capnproto/MessageBuilder.java

@@ -40,7 +40,14 @@ public final class MessageBuilder {
                                       allocationStrategy);
     }
 
+    public MessageBuilder(Allocator allocator) {
+        this.arena = new BuilderArena(allocator);
+    }
+
     private AnyPointer.Builder getRootInternal() {
+        if (this.arena.segments.isEmpty()) {
+            this.arena.allocate(1);
+        }
         SegmentBuilder rootSegment = this.arena.segments.get(0);
         if (rootSegment.currentSize() == 0) {
             int location = rootSegment.allocate(1);