Move Slice to use MemorySegment

- Add out-of-order read methods to BasicSliceInput - Add Slice.mismatch method - Small tweaks and new benchmarks
airlift · Jul 3, 2024 · 383c10b · 383c10b
1 parent b04abb3
commit 383c10b
Show file tree

Hide file tree

Showing 13 changed files with 545 additions and 479 deletions.
diff --git a/src/main/java/io/airlift/slice/BasicSliceInput.java b/src/main/java/io/airlift/slice/BasicSliceInput.java
@@ -150,6 +150,49 @@ public double readDouble()
         return v;
     }
 
+    public double getDouble(int offset)
+    {
+        return slice.getDouble(position + offset);
+    }
+
+    public float getFloat(int offset)
+    {
+        return slice.getFloat(position + offset);
+    }
+
+    public float getShort(int offset)
+    {
+        return slice.getShort(position + offset);
+    }
+
+    public int getInt(int offset)
+    {
+        return slice.getInt(position + offset);
+    }
+
+    public byte getByte(int offset)
+    {
+        return slice.getByte(position + offset);
+    }
+
+    public long getLong(int offset)
+    {
+        return slice.getLong(position + offset);
+    }
+
+    public byte[] getByteArray()
+    {
+        return slice.byteArray();
+    }
+
+    /**
+     * Returns the start index the content of this slice within the byte array wrapped by this slice.
+     */
+    public int getByteArrayOffset()
+    {
+        return position + slice.byteArrayOffset();
+    }
+
     /**
      * Returned slice is a view over {@code slice}
      */
@@ -164,6 +207,13 @@ public Slice readSlice(int length)
         return newSlice;
     }
 
+    public byte[] readBytes()
+    {
+        byte[] bytes = slice.getBytes();
+        position = slice.length();
+        return bytes;
+    }
+
     @Override
     public int read(byte[] destination, int destinationIndex, int length)
     {

diff --git a/src/main/java/io/airlift/slice/InputStreamSliceInput.java b/src/main/java/io/airlift/slice/InputStreamSliceInput.java
@@ -110,7 +110,7 @@ public boolean readBoolean()
     public byte readByte()
     {
         ensureAvailable(SIZE_OF_BYTE);
-        byte v = slice.getByteUnchecked(bufferPosition);
+        byte v = slice.getByte(bufferPosition);
         bufferPosition += SIZE_OF_BYTE;
         return v;
     }
@@ -125,7 +125,7 @@ public int readUnsignedByte()
     public short readShort()
     {
         ensureAvailable(SIZE_OF_SHORT);
-        short v = slice.getShortUnchecked(bufferPosition);
+        short v = slice.getShort(bufferPosition);
         bufferPosition += SIZE_OF_SHORT;
         return v;
     }
@@ -140,7 +140,7 @@ public int readUnsignedShort()
     public int readInt()
     {
         ensureAvailable(SIZE_OF_INT);
-        int v = slice.getIntUnchecked(bufferPosition);
+        int v = slice.getInt(bufferPosition);
         bufferPosition += SIZE_OF_INT;
         return v;
     }
@@ -149,7 +149,7 @@ public int readInt()
     public long readLong()
     {
         ensureAvailable(SIZE_OF_LONG);
-        long v = slice.getLongUnchecked(bufferPosition);
+        long v = slice.getLong(bufferPosition);
         bufferPosition += SIZE_OF_LONG;
         return v;
     }
@@ -174,7 +174,7 @@ public int read()
         }
 
         verify(availableBytes() > 0);
-        int v = slice.getByteUnchecked(bufferPosition) & 0xFF;
+        int v = slice.getByte(bufferPosition) & 0xFF;
         bufferPosition += SIZE_OF_BYTE;
         return v;
     }

diff --git a/src/main/java/io/airlift/slice/JvmUtils.java b/src/main/java/io/airlift/slice/JvmUtils.java
diff --git a/src/main/java/io/airlift/slice/OutputStreamSliceOutput.java b/src/main/java/io/airlift/slice/OutputStreamSliceOutput.java
@@ -133,31 +133,31 @@ public boolean isWritable()
     public void writeByte(int value)
     {
         ensureWritableBytes(SIZE_OF_BYTE);
-        slice.setByteUnchecked(bufferPosition, value);
+        slice.setByte(bufferPosition, value);
         bufferPosition += SIZE_OF_BYTE;
     }
 
     @Override
     public void writeShort(int value)
     {
         ensureWritableBytes(SIZE_OF_SHORT);
-        slice.setShortUnchecked(bufferPosition, value);
+        slice.setShort(bufferPosition, value);
         bufferPosition += SIZE_OF_SHORT;
     }
 
     @Override
     public void writeInt(int value)
     {
         ensureWritableBytes(SIZE_OF_INT);
-        slice.setIntUnchecked(bufferPosition, value);
+        slice.setInt(bufferPosition, value);
         bufferPosition += SIZE_OF_INT;
     }
 
     @Override
     public void writeLong(long value)
     {
         ensureWritableBytes(SIZE_OF_LONG);
-        slice.setLongUnchecked(bufferPosition, value);
+        slice.setLong(bufferPosition, value);
         bufferPosition += SIZE_OF_LONG;
     }
 

diff --git a/src/main/java/io/airlift/slice/SizeOf.java b/src/main/java/io/airlift/slice/SizeOf.java
@@ -20,6 +20,7 @@
 import org.openjdk.jol.vm.VM;
 import org.openjdk.jol.vm.VirtualMachine;
 
+import java.lang.foreign.MemorySegment;
 import java.util.AbstractMap;
 import java.util.List;
 import java.util.Map;
@@ -76,9 +77,29 @@ public final class SizeOf
     public static final int OPTIONAL_DOUBLE_INSTANCE_SIZE = instanceSize(OptionalDouble.class);
 
     public static final int STRING_INSTANCE_SIZE = instanceSize(String.class);
+    public static final int MEMORY_SEGMENT_INSTANCE_SIZE = instanceSize(MemorySegment.ofArray(new byte[0]).getClass());
 
     private static final int SIMPLE_ENTRY_INSTANCE_SIZE = instanceSize(AbstractMap.SimpleEntry.class);
 
+    public static long sizeOf(MemorySegment segment)
+    {
+        if (segment.isNative()) {
+            return MEMORY_SEGMENT_INSTANCE_SIZE;
+        }
+
+        return MEMORY_SEGMENT_INSTANCE_SIZE + segment.heapBase() // base
+                .map(value -> switch (value) {
+                    case byte[] byteArray -> sizeOf(byteArray);
+                    case short[] shortArray -> sizeOf(shortArray);
+                    case int[] intArray -> sizeOf(intArray);
+                    case long[] longArray -> sizeOf(longArray);
+                    case float[] floatArray -> sizeOf(floatArray);
+                    case double[] doubleArray -> sizeOf(doubleArray);
+                    default -> throw new UnsupportedOperationException("Unsupported heap type: " + value.getClass());
+                })
+                .orElseThrow();
+    }
+
     public static long sizeOf(boolean[] array)
     {
         return (array == null) ? 0 : sizeOfBooleanArray(array.length);