Move Slice to use MemorySegment

airlift · Jul 1, 2024 · 203338e · 203338e
1 parent 86fe760
commit 203338e
Show file tree

Hide file tree

Showing 9 changed files with 244 additions and 421 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,24 @@ public final class SizeOf
 
     private static final int SIMPLE_ENTRY_INSTANCE_SIZE = instanceSize(AbstractMap.SimpleEntry.class);
 
+    public static long sizeOf(MemorySegment segment)
+    {
+        // final long length;
+        // final boolean readOnly;
+        // final MemorySessionImpl scope (int state);
+        long size = LONG_INSTANCE_SIZE + BOOLEAN_INSTANCE_SIZE + INTEGER_INSTANCE_SIZE;
+        if (segment.isNative()) {
+            return size
+                    + LONG_INSTANCE_SIZE; // addr;
+        }
+
+        return size
+                + LONG_INSTANCE_SIZE // offset
+                + segment.heapBase() // base
+                    .map(value -> sizeOfByteArray(((byte[]) value).length))
+                    .orElse(0L);
+    }
+
     public static long sizeOf(boolean[] array)
     {
         return (array == null) ? 0 : sizeOfBooleanArray(array.length);