Use Arrays.copyOfRange in Slices.ensureSize

airlift · Aug 16, 2023 · 962bc1b · 962bc1b
1 parent 9f9d495
commit 962bc1b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git a/src/main/java/io/airlift/slice/Slices.java b/src/main/java/io/airlift/slice/Slices.java
@@ -15,6 +15,7 @@
 
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
+import java.util.Arrays;
 import java.util.Random;
 import java.util.concurrent.ThreadLocalRandom;
 
@@ -39,6 +40,10 @@ private Slices() {}
 
     public static Slice ensureSize(Slice existingSlice, int minWritableBytes)
     {
+        if (minWritableBytes < 0) {
+            throw new IllegalArgumentException("minWritableBytes is negative");
+        }
+
         if (existingSlice == null) {
             return allocate(minWritableBytes);
         }
@@ -66,9 +71,13 @@ public static Slice ensureSize(Slice existingSlice, int minWritableBytes)
             }
         }
 
-        Slice newSlice = allocate(newCapacity);
-        newSlice.setBytes(0, existingSlice, 0, existingSlice.length());
-        return newSlice;
+        byte[] bytes = existingSlice.byteArray();
+        int offset = existingSlice.byteArrayOffset();
+        // copy of range can be more efficient because it avoids zeroing memory
+        byte[] copy = Arrays.copyOfRange(bytes, offset, offset + newCapacity);
+        // but if the slice is a view, we need to zero the end of the array
+        Arrays.fill(copy, existingSlice.length(), bytes.length - offset, (byte) 0);
+        return new Slice(copy);
     }
 
     public static Slice allocate(int capacity)

diff --git a/src/test/java/io/airlift/slice/TestSlices.java b/src/test/java/io/airlift/slice/TestSlices.java
@@ -122,8 +122,25 @@ public void testRandom()
     @Test
     public void testEnsureSize()
     {
-        Slice fourBytes = wrappedBuffer(new byte[] {1, 2, 3, 4});
+        Slice slice = Slices.utf8Slice("testValueAbc").slice(4, 5);
+        assertThat(slice.byteArray()).hasSize(12);
+        assertThat(slice.byteArrayOffset()).isEqualTo(4);
+        assertThat(slice.length()).isEqualTo(5);
+        assertThat(slice.toStringUtf8()).isEqualTo("Value");
+
+        // grow to a size within the existing outer slice
+        Slice newSlice = ensureSize(slice, 6);
+        // new byte array is always allocated if size changes
+        assertThat(newSlice.byteArray()).isNotSameAs(slice.byteArray());
+        // size is doubled when small
+        assertThat(newSlice.byteArray()).hasSize(10);
+        // new Slice covers the entire byte array range
+        assertThat(newSlice.byteArrayOffset()).isEqualTo(0);
+        assertThat(newSlice.length()).isEqualTo(10);
+        // the existing data outside the slice view is not copied
+        assertThat(newSlice.toStringUtf8()).isEqualTo("Value\0\0\0\0\0");
 
+        Slice fourBytes = wrappedBuffer(new byte[] {1, 2, 3, 4});
         assertThat(ensureSize(null, 42).length()).isEqualTo(42);
         assertThat(ensureSize(fourBytes, 3)).isSameAs(fourBytes);
         assertThat(ensureSize(fourBytes, 8)).isEqualTo(wrappedBuffer(new byte[] {1, 2, 3, 4, 0, 0, 0, 0}));