Repository: amzn/ion-java (java)

Issue:
Annotations should only preallocate if they wrap a container
When writing containers in binary Ion, writers must wait until the complete contents of the container have been encoded before the container's header (which includes its length) can be written. 

As an optimization, some writers go ahead and buffer some placeholder bytes for the header. If the header fits in those bytes, they can be updated in-place and the entire buffer can be flushed. If they don't fit, then the writer needs to encode the data somewhere else (e.g. in another buffer) and then flush the appropriate buffer slices to the output stream in the correct order.

The `IonRawBinaryWriter` has [a configuration option called `preallocationMode`](https://github.com/amzn/ion-java/blob/2518f3107555f4f3e4c2b9a4274658257afea308/src/com/amazon/ion/impl/bin/IonRawBinaryWriter.java#L504). The `PreallocationMode` enum allows users to specify how much space should be reserved for writing the `Length` component of any container-like values:

https://github.com/amzn/ion-java/blob/2518f3107555f4f3e4c2b9a4274658257afea308/src/com/amazon/ion/impl/bin/IonRawBinaryWriter.java#L154-L187

Somewhat surprisingly, this setting impacts not only how container lengths are written, but also how annotation wrappers are written regardless of their contents. For example, using `preallocateMode=PREALLOCATE_2`, the binary writer would write this Ion value:

```js
'LoneliestNumber'::1
```

as the following bytes:

```
// Assumes a symbol ID of $186 for 'LoneliestNumber'

ee 00 84 82 01 ba 21
 ^  ^  ^  ^  ^     ^--- Ion integer 1
 |  |  |  |  |
 |  |  |  |  +--------- VarUInt-encoded `186`
 |  |  |  +------------ VarUInt length of annotations sequence: 2 bytes
 |  |  +--------------- VarUInt length of annotations sequence + value: 4 bytes
 |  +------------------ Over-allocated placeholder byte
 +--------------------- Beginning of annotations wrapper with VarUInt length to follow
```

We should consider only preallocating space when the wrapped value is a container.


Pull request description:
Compacts preallocated headers when possible
*Issue #, if available:* 

Fixes #399.

*Description of changes:*

`IonRawBinaryWriter` can be configured to preallocate buffer bytes to
hold the length prefix of a new container or annotations wrapper. This
option trades processing speed for data size. Preallocating more bytes
makes the encoded data larger, but increases the chances that the
container's length (once calculated) can be written directly into the
existing buffer.

This patch modifies the writer such that, when finalizing the encoding
of a container or annotations wrapper, the writer will check to see if
the encoded length is small enough to fit in the type descriptor byte.
If it is, the encoded bytes are shifted backwards in the buffer to
eliminate the unnecessary preallocation. This means that the encoding of
small containers isn't penalized for having preallocation enabled.

Benchmarks showed no substantial difference in processing speed from the
previous release, so this data size improvement does not come at the
cost of throughput.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.


Patch:
diff --git a/src/com/amazon/ion/impl/bin/IonRawBinaryWriter.java b/src/com/amazon/ion/impl/bin/IonRawBinaryWriter.java
index bc7d63b5..de6a719b 100644
--- a/src/com/amazon/ion/impl/bin/IonRawBinaryWriter.java
+++ b/src/com/amazon/ion/impl/bin/IonRawBinaryWriter.java
@@ -202,6 +202,15 @@ import java.util.NoSuchElementException;
 
         /*package*/ abstract void patchLength(final WriteBuffer buffer, final long position, final long length);
 
+        /**
+         * Returns the number of header bytes that this mode would preallocate to hold the VarUInt-encoded length of
+         * the current value. This number is equal to the total header length (i.e. `typedLength`) minus one, as it does
+         * not include the type descriptor byte. (Examples: PREALLOCATE_0 returns `0`, PREALLOCATE_1 returns `1`, etc.)
+         */
+        int numberOfLengthBytes() {
+            return typedLength - 1;
+        }
+
         /*package*/ static PreallocationMode withPadSize(final int pad)
         {
             switch (pad)
@@ -723,48 +732,70 @@ import java.util.NoSuchElementException;
 
     private ContainerInfo popContainer()
     {
-        final ContainerInfo current = containers.pop();
-        if (current == null)
+        final ContainerInfo currentContainer = containers.pop();
+        if (currentContainer == null)
         {
             throw new IllegalStateException("Tried to pop container state without said container");
         }
 
         // only patch for real containers and annotations -- we use VALUE for tracking only
-        final long length = current.length;
-        if (current.type != ContainerType.VALUE)
+        long length = currentContainer.length;
+        if (currentContainer.type != ContainerType.VALUE)
         {
             // patch in the length
-            final long position = current.position;
-            if (current.length <= preallocationMode.contentMaxLength && preallocationMode != PreallocationMode.PREALLOCATE_0)
+            final long positionOfFirstLengthByte = currentContainer.position;
+            if (length <= 0xD) {
+                // The body of this container/wrapper is small enough that its length can fit in the lower nibble of
+                // the type descriptor byte; we don't need the extra length bytes that were preallocated (if any).
+                // We'll shift the encoded body of the container/wrapper backwards in the buffer to overwrite them.
+
+                // The number of bytes we need to shift by is determined by the writer's preallocation mode.
+                final int numberOfBytesToShiftBy = preallocationMode.numberOfLengthBytes();
+
+                // `length` is the encoded length of the container/wrapper we're stepping out of. It does not
+                // include any header bytes. In this `if` branch, we've confirmed that `length` is <= 0xD,
+                // so this downcast from `long` to `int` is safe.
+                final int lengthOfSliceToShift = (int) length;
+
+                // Shift the container/wrapper body backwards in the buffer. Because this only happens when
+                // `lengthOfSliceToShift` is 13 or fewer bytes, this will usually be a very fast memcpy.
+                // It's slightly more work if the slice we're shifting happens to straddle two memory blocks
+                // inside the buffer.
+                buffer.shiftBytesLeft(lengthOfSliceToShift, numberOfBytesToShiftBy);
+
+                // Overwrite the lower nibble of the original type descriptor byte with the body's encoded length.
+                final long typeDescriptorPosition = positionOfFirstLengthByte - 1;
+                final long type = (buffer.getUInt8At(typeDescriptorPosition) & 0xF0) | length;
+                buffer.writeUInt8At(typeDescriptorPosition, type);
+
+                // We've reclaimed some number of bytes; adjust the container length as appropriate.
+                length -= numberOfBytesToShiftBy;
+            }
+            else if (currentContainer.length <= preallocationMode.contentMaxLength)
             {
-                preallocationMode.patchLength(buffer, position, length);
+                // The container's encoded body is too long to fit the length in the type descriptor byte, but it will
+                // fit in the preallocated length bytes that were added to the buffer when the container was started.
+                // Update those bytes with the VarUInt encoding of the length value.
+                preallocationMode.patchLength(buffer, positionOfFirstLengthByte, length);
             }
             else
             {
-                // side patch
-                if (current.length <= 0xD && preallocationMode == PreallocationMode.PREALLOCATE_0)
-                {
-                    // XXX if we're not using padding we can get here and optimize the length a little without side patching!
-                    final long typePosition = position - 1;
-                    final long type = (buffer.getUInt8At(typePosition) & 0xF0) | current.length;
-                    buffer.writeUInt8At(typePosition, type);
-                }
-                else
-                {
-                    addPatchPoint(position, preallocationMode.typedLength - 1, length);
-                }
+                // The container's encoded body is too long to fit in the length bytes that were preallocated.
+                // Write the VarUInt encoding of the length in a secondary buffer and make a note to include that
+                // when we go to flush the primary buffer to the output stream.
+                addPatchPoint(positionOfFirstLengthByte, preallocationMode.numberOfLengthBytes(), length);
             }
         }
-        if (current.patches != null)
+        if (currentContainer.patches != null)
         {
             // at this point, we've appended our patch points upward, lets make sure we get
             // our child patch points in
-            extendPatchPoints(current.patches);
+            extendPatchPoints(currentContainer.patches);
         }
 
         // make sure to record length upward
         updateLength(length);
-        return current;
+        return currentContainer;
     }
 
     private void writeVarUInt(final long value)
diff --git a/src/com/amazon/ion/impl/bin/WriteBuffer.java b/src/com/amazon/ion/impl/bin/WriteBuffer.java
index e3973c56..0df5ba68 100644
--- a/src/com/amazon/ion/impl/bin/WriteBuffer.java
+++ b/src/com/amazon/ion/impl/bin/WriteBuffer.java
@@ -170,6 +170,122 @@ import java.util.List;
         block.limit += len;
     }
 
+    /**
+     * Shifts the last `length` bytes in the buffer to the left. This can be used when a value's header was
+     * preallocated but the value's encoded size proved to be much smaller than anticipated.
+     *
+     * The caller must guarantee that the buffer contains enough bytes to perform the requested shift.
+     *
+     * @param length    The number of bytes at the end of the buffer that we'll be shifting to the left.
+     * @param shiftBy   The number of bytes to the left that we'll be shifting.
+     */
+    public void shiftBytesLeft(int length, int shiftBy) {
+        if (shiftBy == 0) {
+            // Nothing to do.
+            return;
+        }
+
+        // If all of the bytes that we need to shift are in the current block, do a simple memcpy.
+        if (current.limit >= length + shiftBy) {
+            shiftBytesLeftWithinASingleBlock(length, shiftBy);
+            return;
+        }
+
+        // Otherwise, the slice we're shifting straddles multiple blocks. We'll need to iterate across those blocks
+        // applying shifting logic to each one.
+        shiftBytesLeftAcrossBlocks(length, shiftBy);
+    }
+
+    /**
+     * Shifts the last `length` bytes in the buffer to the left. The caller must guarantee that the `current` Block
+     * contains at least `length + shiftBy` bytes. This ensures that we're shifting a contiguous slice of bytes within
+     * a single block.
+     *
+     * @param length    The number of bytes at the end of the buffer that we'll be shifting to the left.
+     * @param shiftBy   The number of bytes to the left that we'll be shifting.
+     */
+    private void shiftBytesLeftWithinASingleBlock(int length, int shiftBy) {
+        int startOfSliceToShift = current.limit - length;
+        System.arraycopy(
+                current.data,
+                startOfSliceToShift,
+                current.data,
+                startOfSliceToShift - shiftBy,
+                length
+        );
+        // Update the `limit` (cursor position) within the current block to reflect that
+        // we have reclaimed `length` bytes of space in the buffer.
+        current.limit -= shiftBy;
+    }
+
+    /**
+     * Shifts the last `length` bytes in the buffer to the left. Unlike
+     * {@link #shiftBytesLeftWithinASingleBlock(int, int)}, this method supports shifting bytes across multiple blocks
+     * in the buffer.
+     *
+     * @param length    The number of bytes at the end of the buffer that we'll be shifting to the left.
+     * @param shiftBy   The number of bytes to the left that we'll be shifting.
+     */
+    private void shiftBytesLeftAcrossBlocks(int length, int shiftBy) {
+        // Our starting position is the first byte that we plan to shift backwards. The `bufferOffset` is the cursor's
+        // current position in the stream; we can use this to derive which block it's in as well as its position within
+        // that block.
+        long bufferOffset = position() - length;
+
+        Block block = null;
+        while (length > 0)
+        {
+            // Using the `bufferOffset`, determine which block we're in...
+            int blockIndex = index(bufferOffset);
+            block = blocks.get(blockIndex);
+            // ...and our offset within that block.
+            int blockOffset = offset(bufferOffset);
+
+            // If the block offset is within `shiftBy` bytes of beginning of the block, some bytes we're shifting
+            // will end up in the previous block. Here's an illustrated example:
+            //
+            // shiftBy = 2, blockIndex = 1, blockOffset = 1
+            //                        v---- Cursor is here
+            // Before: [A B C D E] [F G H I J]
+            // After : [A B C D G] [F G H I J]
+            //  Now this is G --^       ^-- And the cursor is here, ready to continue copy the rest.
+            if (blockOffset < shiftBy) {
+                Block previousBlock = blocks.get(blockIndex - 1);
+                int numberOfBytesToShift = Math.min(length, shiftBy) - blockOffset;
+                System.arraycopy(block.data, blockOffset, previousBlock.data, previousBlock.data.length - numberOfBytesToShift, numberOfBytesToShift);
+
+                // Now that we've shifted some bytes, update our position within the buffer.
+                bufferOffset += numberOfBytesToShift;
+                length -= numberOfBytesToShift;
+
+                // If there are no more bytes to shift...
+                if (length == 0) {
+                    // ...lower the `limit` because we've reclaimed some bytes in this block...
+                    block.limit -= numberOfBytesToShift;
+                    // ...and early return.
+                    return;
+                }
+                // Otherwise, use our new buffer offset to recalculate our block-specific position.
+                blockIndex = index(bufferOffset);
+                block = blocks.get(blockIndex);
+                blockOffset = offset(bufferOffset);
+            }
+
+            // At this point, the block offset is at least `shiftBy` bytes away from the beginning of the block, so
+            // we can do a memcpy to shift the rest of the bytes in the block.
+            int numberOfBytesToShift = Math.min(length, block.data.length - blockOffset);
+            System.arraycopy(block.data, blockOffset, block.data, blockOffset - shiftBy, numberOfBytesToShift);
+
+            // Update our counters and see if there are any more bytes to shift.
+            bufferOffset += numberOfBytesToShift;
+            length -= numberOfBytesToShift;
+        }
+        if (block != null) {
+            // We've reclaimed some space in this block; lower the `limit` accordingly.
+            block.limit -= shiftBy;
+        }
+    }
+
     /** Writes an array of bytes to the buffer expanding if necessary, defaulting to the entire array. */
     public void writeBytes(byte[] bytes)
     {


Test patch:
diff --git a/test/com/amazon/ion/impl/bin/WriteBufferTest.java b/test/com/amazon/ion/impl/bin/WriteBufferTest.java
index 5a5ddb5a..24822820 100644
--- a/test/com/amazon/ion/impl/bin/WriteBufferTest.java
+++ b/test/com/amazon/ion/impl/bin/WriteBufferTest.java
@@ -921,4 +921,42 @@ public class WriteBufferTest
         buf.truncate(3);
         assertBuffer("ARG".getBytes("UTF-8"));
     }
+
+    @Test
+    public void shiftBytesLeftWithinFirstBufferBlock() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // All bytes are being shifted within the first buffer block.
+        buf.writeBytes("01234567".getBytes());
+        buf.shiftBytesLeft(4, 1);
+        assertBuffer("0124567".getBytes());
+    }
+
+    @Test
+    public void shiftBytesLeftWithinLastBufferBlock() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // All bytes are being shifted within the last buffer block.
+        buf.writeBytes("0123456789ABCDEF".getBytes());
+        buf.shiftBytesLeft(3, 2);
+        assertBuffer("0123456789ADEF".getBytes());
+    }
+
+    @Test
+    public void shiftBytesLeftAcrossBufferBlocks() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // Some bytes are shifted to the previous block, some are shifted within
+        // the last block.
+        buf.writeBytes("0123456789ABCDEF".getBytes());
+        buf.shiftBytesLeft(8, 3);
+        assertBuffer("0123489ABCDEF".getBytes());
+    }
+
+    @Test
+    public void shiftBytesLeftAcrossBufferBlocksExclusively() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // Unlike `shiftBytesLeftAcrossBufferBlocks`, EVERY byte that is shifted
+        // in the buffer ends up in the previous block.
+        buf.writeBytes("0123456789ABCDEF".getBytes());
+        buf.shiftBytesLeft(5, 5);
+        assertBuffer("012345BCDEF".getBytes());
+    }
 }