Move Slice to use MemorySegment

airlift · Jul 2, 2024 · 6ec178e · 6ec178e
1 parent b04abb3
commit 6ec178e
Show file tree

Hide file tree

Showing 10 changed files with 265 additions and 507 deletions.
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;
@@ -79,6 +80,27 @@ public final class SizeOf
 
     private static final int SIMPLE_ENTRY_INSTANCE_SIZE = instanceSize(AbstractMap.SimpleEntry.class);
 
+    private static final int MEMORY_SEGMENT_INSTANCE_SIZE = instanceSize(MemorySegment.ofArray(new byte[0]).getClass());
+
+    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);