8343188: Investigate ways to simplify MemorySegment::ofBuffer

Reviewed-by: mcimadamore

src/java.base/share/classes/java/nio/Buffer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2000, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2000, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -37,11 +37,9 @@ import jdk.internal.misc.VM.BufferPool;
 import jdk.internal.util.Preconditions;
 import jdk.internal.vm.annotation.ForceInline;
 
-import java.io.FileDescriptor;
 import java.lang.foreign.MemorySegment;
 import java.lang.ref.Reference;
 import java.util.List;
-import java.util.Objects;
 import java.util.Spliterator;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -780,6 +778,23 @@ public abstract sealed class Buffer
         return Preconditions.checkIndex(i, limit - nb + 1, IOOBE_FORMATTER);
     }
 
+    /**
+     * {@return the scale shifts for this Buffer}
+     * <p>
+     * The scale shifts are:
+     *   ByteBuffer:               0
+     *   ShortBuffer, CharBuffer:  1
+     *   IntBuffer, FloatBuffer:   2
+     *   LongBuffer, DoubleBuffer: 3
+     */
+    abstract int scaleShifts();
+
+    abstract AbstractMemorySegmentImpl heapSegment(Object base,
+                                                   long offset,
+                                                   long length,
+                                                   boolean readOnly,
+                                                   MemorySessionImpl bufferScope);
+
     final int markValue() {                             // package-private
         return mark;
     }
@@ -832,6 +847,7 @@ public abstract sealed class Buffer
                     return new HeapByteBuffer(hb, -1, 0, capacity, capacity, offset, segment);
                 }
 
+                @ForceInline
                 @Override
                 public Object getBufferBase(Buffer buffer) {
                     return buffer.base();
@@ -906,6 +922,23 @@ public abstract sealed class Buffer
                 public int pageSize() {
                     return Bits.pageSize();
                 }
+
+                @ForceInline
+                @Override
+                public int scaleShifts(Buffer buffer) {
+                    return buffer.scaleShifts();
+                }
+
+                @ForceInline
+                @Override
+                public AbstractMemorySegmentImpl heapSegment(Buffer buffer,
+                                                             Object base,
+                                                             long offset,
+                                                             long length,
+                                                             boolean readOnly,
+                                                             MemorySessionImpl bufferScope) {
+                    return buffer.heapSegment(base, offset, length, readOnly, bufferScope);
+                }
             });
     }
 

src/java.base/share/classes/java/nio/ByteBufferAs-X-Buffer.java.template

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2000, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2000, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -28,6 +28,10 @@
 package java.nio;
 
 import java.lang.foreign.MemorySegment;
+import jdk.internal.foreign.AbstractMemorySegmentImpl;
+import jdk.internal.foreign.MemorySessionImpl;
+import jdk.internal.foreign.SegmentFactories;
+import jdk.internal.vm.annotation.ForceInline;
 import java.util.Objects;
 import jdk.internal.misc.Unsafe;
 
@@ -246,6 +250,21 @@ class ByteBufferAs$Type$Buffer$RW$$BO$                  // package-private
 
 #end[char]
 
+    @ForceInline
+    @Override
+    int scaleShifts() {
+        return Integer.numberOfTrailingZeros($Fulltype$.BYTES);
+    }
+
+    @ForceInline
+    @Override
+    AbstractMemorySegmentImpl heapSegment(Object base,
+                                          long offset,
+                                          long length,
+                                          boolean readOnly,
+                                          MemorySessionImpl bufferScope) {
+        return SegmentFactories.arrayOfByteSegment(base, offset, length, readOnly, bufferScope);
+    }
 
     public ByteOrder order() {
 #if[boB]

src/java.base/share/classes/java/nio/Direct-X-Buffer.java.template

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2000, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2000, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -31,7 +31,10 @@ import java.io.FileDescriptor;
 import java.lang.foreign.MemorySegment;
 import java.lang.ref.Reference;
 import java.util.Objects;
+import jdk.internal.foreign.AbstractMemorySegmentImpl;
 import jdk.internal.foreign.MemorySessionImpl;
+import jdk.internal.foreign.SegmentFactories;
+import jdk.internal.vm.annotation.ForceInline;
 import jdk.internal.misc.ScopedMemoryAccess.ScopedAccessError;
 import jdk.internal.misc.VM;
 import jdk.internal.ref.Cleaner;
@@ -528,6 +531,24 @@ class Direct$Type$Buffer$RW$$BO$
     }
 #end[char]
 
+#if[byte]
+    @ForceInline
+    @Override
+    int scaleShifts() {
+        return 0;
+    }
+
+    @ForceInline
+    @Override
+    AbstractMemorySegmentImpl heapSegment(Object base,
+                                          long offset,
+                                          long length,
+                                          boolean readOnly,
+                                          MemorySessionImpl bufferScope) {
+        // Direct buffers are not backed by an array.
+        throw new UnsupportedOperationException();
+    }
+#end[byte]
 
 #if[byte]
     // #BIN

src/java.base/share/classes/java/nio/Heap-X-Buffer.java.template

@@ -28,6 +28,10 @@
 package java.nio;
 
 import java.lang.foreign.MemorySegment;
+import jdk.internal.foreign.AbstractMemorySegmentImpl;
+import jdk.internal.foreign.MemorySessionImpl;
+import jdk.internal.foreign.SegmentFactories;
+import jdk.internal.vm.annotation.ForceInline;
 import java.util.Objects;
 
 /**
@@ -735,6 +739,23 @@ class Heap$Type$Buffer$RW$
 
 #end[char]
 
+#if[byte]
+    @ForceInline
+    @Override
+    int scaleShifts() {
+        return 0;
+    }
+
+    @ForceInline
+    @Override
+    AbstractMemorySegmentImpl heapSegment(Object base,
+                                          long offset,
+                                          long length,
+                                          boolean readOnly,
+                                          MemorySessionImpl bufferScope) {
+        return SegmentFactories.arrayOf$Type$Segment(base, offset, length, readOnly, bufferScope);
+    }
+#end[byte]
 
 #if[!byte]
 

src/java.base/share/classes/java/nio/X-Buffer.java.template

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2000, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2000, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -37,6 +37,10 @@ import java.util.stream.StreamSupport;
 import java.util.stream.$Streamtype$Stream;
 #end[streamableType]
 
+import jdk.internal.foreign.AbstractMemorySegmentImpl;
+import jdk.internal.foreign.MemorySessionImpl;
+import jdk.internal.foreign.SegmentFactories;
+import jdk.internal.vm.annotation.ForceInline;
 import java.lang.foreign.MemorySegment;
 import java.util.Objects;
 import jdk.internal.util.ArraysSupport;
@@ -2321,6 +2325,22 @@ public abstract sealed class $Type$Buffer
 
 #end[byte]
 
+    @ForceInline
+    @Override
+    int scaleShifts() {
+        return Integer.numberOfTrailingZeros($Fulltype$.BYTES);
+    }
+
+    @ForceInline
+    @Override
+    AbstractMemorySegmentImpl heapSegment(Object base,
+                                          long offset,
+                                          long length,
+                                          boolean readOnly,
+                                          MemorySessionImpl bufferScope) {
+        return SegmentFactories.arrayOf$Type$Segment(base, offset, length, readOnly, bufferScope);
+    }
+
 #if[streamableType]
 
 #if[char]

src/java.base/share/classes/jdk/internal/access/JavaNioAccess.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2007, 2022, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2007, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -27,10 +27,11 @@ package jdk.internal.access;
 
 import jdk.internal.access.foreign.MappedMemoryUtilsProxy;
 import jdk.internal.access.foreign.UnmapperProxy;
+import jdk.internal.foreign.AbstractMemorySegmentImpl;
+import jdk.internal.foreign.MemorySessionImpl;
 import jdk.internal.misc.VM.BufferPool;
 
 import java.lang.foreign.MemorySegment;
-import java.io.FileDescriptor;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 
@@ -127,4 +128,14 @@ public interface JavaNioAccess {
      * Used by {@code jdk.internal.foreign.NativeMemorySegmentImpl}.
      */
     int pageSize();
+
+    int scaleShifts(Buffer buffer);
+
+    AbstractMemorySegmentImpl heapSegment(Buffer buffer,
+                                          Object base,
+                                          long offset,
+                                          long length,
+                                          boolean readOnly,
+                                          MemorySessionImpl bufferScope);
+
 }

src/java.base/share/classes/jdk/internal/foreign/AbstractMemorySegmentImpl.java

@@ -30,12 +30,6 @@ import java.lang.reflect.Array;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
-import java.nio.CharBuffer;
-import java.nio.DoubleBuffer;
-import java.nio.FloatBuffer;
-import java.nio.IntBuffer;
-import java.nio.LongBuffer;
-import java.nio.ShortBuffer;
 import java.nio.charset.Charset;
 import java.util.*;
 import java.util.function.BiFunction;
@@ -51,14 +45,11 @@ import jdk.internal.access.foreign.UnmapperProxy;
 import jdk.internal.misc.ScopedMemoryAccess;
 import jdk.internal.reflect.CallerSensitive;
 import jdk.internal.reflect.Reflection;
-import jdk.internal.util.Architecture;
 import jdk.internal.util.ArraysSupport;
 import jdk.internal.util.Preconditions;
 import jdk.internal.vm.annotation.ForceInline;
 import sun.nio.ch.DirectBuffer;
 
-import static java.lang.foreign.ValueLayout.JAVA_BYTE;
-
 /**
  * This abstract class provides an immutable implementation for the {@code MemorySegment} interface. This class contains information
  * about the segment's spatial and temporal bounds; each memory segment implementation is associated with an owner thread which is set at creation time.
@@ -521,48 +512,53 @@ public abstract sealed class AbstractMemorySegmentImpl
                 unsafeGetBase());
     }
 
-    public static AbstractMemorySegmentImpl ofBuffer(Buffer bb) {
+    @ForceInline
-        Objects.requireNonNull(bb);
+    public static AbstractMemorySegmentImpl ofBuffer(Buffer b) {
-        Object base = NIO_ACCESS.getBufferBase(bb);
+        // Implicit null check via NIO_ACCESS.scaleShifts(b)
-        if (!bb.isDirect() && base == null) {
+        final int scaleShifts = NIO_ACCESS.scaleShifts(b);
+        return ofBuffer(b, offset(b, scaleShifts), length(b, scaleShifts));
+    }
+
+    @ForceInline
+    private static AbstractMemorySegmentImpl ofBuffer(Buffer b, long offset, long length) {
+        final Object base = NIO_ACCESS.getBufferBase(b);
+        return (base == null)
+                ? nativeSegment(b, offset, length)
+                : NIO_ACCESS.heapSegment(b, base, offset, length, b.isReadOnly(), bufferScope(b));
+    }
+
+    @ForceInline
+    private static long offset(Buffer b, int scaleShifts) {
+        final long bbAddress = NIO_ACCESS.getBufferAddress(b);
+        return bbAddress + (((long) b.position()) << scaleShifts);
+    }
+
+    @ForceInline
+    private static long length(Buffer b, int scaleShifts) {
+        return ((long) b.limit() - b.position()) << scaleShifts;
+    }
+
+    @ForceInline
+    private static AbstractMemorySegmentImpl nativeSegment(Buffer b, long offset, long length) {
+        if (!b.isDirect()) {
             throw new IllegalArgumentException("The provided heap buffer is not backed by an array.");
         }
-        long bbAddress = NIO_ACCESS.getBufferAddress(bb);
+        final UnmapperProxy unmapper = NIO_ACCESS.unmapper(b);
-        UnmapperProxy unmapper = NIO_ACCESS.unmapper(bb);
+        return unmapper == null
-
+                ? new NativeMemorySegmentImpl(offset, length, b.isReadOnly(), bufferScope(b))
-        int pos = bb.position();
+                : new MappedMemorySegmentImpl(offset, unmapper, length, b.isReadOnly(), bufferScope(b));
-        int limit = bb.limit();
-        int size = limit - pos;
-
-        AbstractMemorySegmentImpl bufferSegment = (AbstractMemorySegmentImpl) NIO_ACCESS.bufferSegment(bb);
-        boolean readOnly = bb.isReadOnly();
-        int scaleFactor = getScaleFactor(bb);
-        final MemorySessionImpl bufferScope;
-        if (bufferSegment != null) {
-            bufferScope = bufferSegment.scope;
-        } else {
-            bufferScope = MemorySessionImpl.createHeap(bufferRef(bb));
-        }
-        long off = bbAddress + ((long)pos << scaleFactor);
-        long len = (long)size << scaleFactor;
-        if (base != null) {
-            return switch (base) {
-                case byte[]   _ -> new HeapMemorySegmentImpl.OfByte(off, base, len, readOnly, bufferScope);
-                case short[]  _ -> new HeapMemorySegmentImpl.OfShort(off, base, len, readOnly, bufferScope);
-                case char[]   _ -> new HeapMemorySegmentImpl.OfChar(off, base, len, readOnly, bufferScope);
-                case int[]    _ -> new HeapMemorySegmentImpl.OfInt(off, base, len, readOnly, bufferScope);
-                case float[]  _ -> new HeapMemorySegmentImpl.OfFloat(off, base, len, readOnly, bufferScope);
-                case long[]   _ -> new HeapMemorySegmentImpl.OfLong(off, base, len, readOnly, bufferScope);
-                case double[] _ -> new HeapMemorySegmentImpl.OfDouble(off, base, len, readOnly, bufferScope);
-                default         -> throw new AssertionError("Cannot get here");
-            };
-        } else if (unmapper == null) {
-            return new NativeMemorySegmentImpl(off, len, readOnly, bufferScope);
-        } else {
-            return new MappedMemorySegmentImpl(off, unmapper, len, readOnly, bufferScope);
-        }
     }
 
+    @ForceInline
+    private static MemorySessionImpl bufferScope(Buffer b) {
+        final AbstractMemorySegmentImpl bufferSegment =
+                (AbstractMemorySegmentImpl) NIO_ACCESS.bufferSegment(b);
+        return bufferSegment == null
+                ? MemorySessionImpl.createHeap(bufferRef(b))
+                : bufferSegment.scope;
+    }
+
+    @ForceInline
     private static Object bufferRef(Buffer buffer) {
         if (buffer instanceof DirectBuffer directBuffer) {
             // direct buffer, return either the buffer attachment (for slices and views), or the buffer itself
@@ -660,15 +656,6 @@ public abstract sealed class AbstractMemorySegmentImpl
         }
     }
 
-    private static int getScaleFactor(Buffer buffer) {
-        return switch (buffer) {
-            case ByteBuffer   _                 -> 0;
-            case CharBuffer   _, ShortBuffer  _ -> 1;
-            case IntBuffer    _, FloatBuffer  _ -> 2;
-            case LongBuffer   _, DoubleBuffer _ -> 3;
-        };
-    }
-
     // accessors
 
     @ForceInline

src/java.base/share/classes/jdk/internal/foreign/SegmentFactories.java

@@ -140,6 +140,64 @@ public class SegmentFactories {
                 MemorySessionImpl.createHeap(arr));
     }
 
+    // Buffer conversion factories
+
+    public static AbstractMemorySegmentImpl arrayOfByteSegment(Object base,
+                                                               long offset,
+                                                               long length,
+                                                               boolean readOnly,
+                                                               MemorySessionImpl bufferScope) {
+        return new HeapMemorySegmentImpl.OfByte(offset, base, length, readOnly, bufferScope);
+    }
+
+    public static AbstractMemorySegmentImpl arrayOfShortSegment(Object base,
+                                                                long offset,
+                                                                long length,
+                                                                boolean readOnly,
+                                                                MemorySessionImpl bufferScope) {
+        return new HeapMemorySegmentImpl.OfShort(offset, base, length, readOnly, bufferScope);
+    }
+
+    public static AbstractMemorySegmentImpl arrayOfCharSegment(Object base,
+                                                               long offset,
+                                                               long length,
+                                                               boolean readOnly,
+                                                               MemorySessionImpl bufferScope) {
+        return new HeapMemorySegmentImpl.OfChar(offset, base, length, readOnly, bufferScope);
+    }
+
+    public static AbstractMemorySegmentImpl arrayOfIntSegment(Object base,
+                                                              long offset,
+                                                              long length,
+                                                              boolean readOnly,
+                                                              MemorySessionImpl bufferScope) {
+        return new HeapMemorySegmentImpl.OfInt(offset, base, length, readOnly, bufferScope);
+    }
+
+    public static AbstractMemorySegmentImpl arrayOfFloatSegment(Object base,
+                                                                long offset,
+                                                                long length,
+                                                                boolean readOnly,
+                                                                MemorySessionImpl bufferScope) {
+        return new HeapMemorySegmentImpl.OfFloat(offset, base, length, readOnly, bufferScope);
+    }
+
+    public static AbstractMemorySegmentImpl arrayOfLongSegment(Object base,
+                                                               long offset,
+                                                               long length,
+                                                               boolean readOnly,
+                                                               MemorySessionImpl bufferScope) {
+        return new HeapMemorySegmentImpl.OfLong(offset, base, length, readOnly, bufferScope);
+    }
+
+    public static AbstractMemorySegmentImpl arrayOfDoubleSegment(Object base,
+                                                                 long offset,
+                                                                 long length,
+                                                                 boolean readOnly,
+                                                                 MemorySessionImpl bufferScope) {
+        return new HeapMemorySegmentImpl.OfDouble(offset, base, length, readOnly, bufferScope);
+    }
+
     public static MemorySegment allocateSegment(long byteSize, long byteAlignment, MemorySessionImpl sessionImpl,
                                                   boolean shouldReserve) {
         ensureInitialized();

test/micro/org/openjdk/bench/java/lang/foreign/SegmentOfBuffer.java

@@ -0,0 +1,71 @@
+/*
+ *  Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ *  DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ *  This code is free software; you can redistribute it and/or modify it
+ *  under the terms of the GNU General Public License version 2 only, as
+ *  published by the Free Software Foundation.
+ *
+ *  This code is distributed in the hope that it will be useful, but WITHOUT
+ *  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ *  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ *  version 2 for more details (a copy is included in the LICENSE file that
+ *  accompanied this code).
+ *
+ *  You should have received a copy of the GNU General Public License version
+ *  2 along with this work; if not, write to the Free Software Foundation,
+ *  Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ *   Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ *  or visit www.oracle.com if you need additional information or have any
+ *  questions.
+ *
+ */
+
+package org.openjdk.bench.java.lang.foreign;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.lang.foreign.MemorySegment;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.concurrent.TimeUnit;
+
+@BenchmarkMode(Mode.AverageTime)
+@Warmup(iterations = 5, time = 500, timeUnit = TimeUnit.MILLISECONDS)
+@Measurement(iterations = 10, time = 500, timeUnit = TimeUnit.MILLISECONDS)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@State(Scope.Benchmark)
+public class SegmentOfBuffer {
+
+    private final ByteBuffer buffer = ByteBuffer
+            .allocateDirect(0x1000)
+            .order(ByteOrder.nativeOrder());
+
+    @Benchmark
+    @Fork(value = 3)
+    public long ofBuffer() {
+        return MemorySegment.ofBuffer(buffer).address();
+    }
+
+    @Benchmark
+    @Fork(value = 3, jvmArgsAppend = "-XX:CompileCommand=inline,jdk.internal.foreign.AbstractMemorySegmentImpl::ofBuffer,false")
+    public long ofBufferInlineFalse() {
+        return MemorySegment.ofBuffer(buffer).address();
+    }
+
+    @Benchmark
+    @Fork(value = 3, jvmArgsAppend = "-XX:CompileCommand=inline,jdk.internal.foreign.AbstractMemorySegmentImpl::ofBuffer,true")
+    public long ofBufferInlineTrue() {
+        return MemorySegment.ofBuffer(buffer).address();
+    }
+
+}