8303757: MemorySegment::reinterpret should accept an Arena

Reviewed-by: jvernee

src/java.base/share/classes/java/lang/foreign/MemorySegment.java

@@ -380,14 +380,14 @@
  *}
  * <p>
  * In some cases, a client might additionally want to assign new temporal bounds to a zero-length memory segment.
- * This can be done using the {@link #reinterpret(long, Scope, Consumer)} method, which returns a
+ * This can be done using the {@link #reinterpret(long, Arena, Consumer)} method, which returns a
  * new native segment with the desired size and the same temporal bounds as those in the provided arena:
  *
  * {@snippet lang = java:
  * MemorySegment foreign = null;
  * try (Arena arena = Arena.ofConfined()) {
  *       foreign = someSegment.get(ValueLayout.ADDRESS, 0)           // size = 0, scope = always alive
- *                            .reinterpret(4, arena.scope(), null);  // size = 4, scope = arena.scope()
+ *                            .reinterpret(4, arena, null);          // size = 4, scope = arena.scope()
  *       int x = foreign.get(ValueLayout.JAVA_INT, 0);               // ok
  * }
  * int x = foreign.get(ValueLayout.JAVA_INT, 0); // throws IllegalStateException
@@ -412,7 +412,7 @@
  * be associated with a given native pointer. In these cases using an unbounded address layout might be preferable.
  * <p>
  * All the methods which can be used to manipulate zero-length memory segments
- * ({@link #reinterpret(long)}, {@link #reinterpret(Scope, Consumer)}, {@link #reinterpret(long, Scope, Consumer)} and
+ * ({@link #reinterpret(long)}, {@link #reinterpret(Arena, Consumer)}, {@link #reinterpret(long, Arena, Consumer)} and
  * {@link AddressLayout#withTargetLayout(MemoryLayout)}) are
  * <a href="package-summary.html#restricted"><em>restricted</em></a> methods, and should be used with caution:
  * assigning a segment incorrect spatial and/or temporal bounds could result in a VM crash when attempting to access
@@ -570,10 +570,6 @@ default MemorySegment asSlice(long offset, MemoryLayout layout) {
 
     /**
      * Returns a new memory segment that has the same address and scope as this segment, but with the provided size.
-     * Equivalent to the following code:
-     * {@snippet lang=java :
-     * reinterpret(newSize, scope(), null);
-     * }
      * <p>
      * This method is <a href="package-summary.html#restricted"><em>restricted</em></a>.
      * Restricted methods are unsafe, and, if used incorrectly, their use might crash
@@ -592,17 +588,38 @@ default MemorySegment asSlice(long offset, MemoryLayout layout) {
 
     /**
      * Returns a new memory segment with the same address and size as this segment, but with the provided scope.
-     * Equivalent to the following code:
+     * As such, the returned segment cannot be accessed after the provided arena has been closed.
+     * Moreover, the returned segment can be accessed compatibly with the confinement restrictions associated with the
+     * provided arena: that is, if the provided arena is a {@linkplain Arena#ofConfined() confined arena},
+     * the returned segment can only be accessed by the arena's owner thread, regardless of the confinement restrictions
+     * associated with this segment. In other words, this method returns a segment that behaves as if it had been allocated
+     * using the provided arena.
+     * <p>
+     * Clients can specify an optional cleanup action that should be executed when the provided scope becomes
+     * invalid. This cleanup action receives a fresh memory segment that is obtained from this segment as follows:
      * {@snippet lang=java :
-     * reinterpret(byteSize(), scope, cleanup);
+     * MemorySegment cleanupSegment = MemorySegment.ofAddress(this.address());
      * }
+     * That is, the cleanup action receives a segment that is associated with a fresh scope that is always alive,
+     * and is accessible from any thread. The size of the segment accepted by the cleanup action is {@link #byteSize()}.
+     * <p>
+     * This method is <a href="package-summary.html#restricted"><em>restricted</em></a>.
+     * Restricted methods are unsafe, and, if used incorrectly, their use might crash
+     * the JVM or, worse, silently result in memory corruption. Thus, clients should refrain from depending on
+     * restricted methods, and use safe and supported functionalities, where possible.
+     *
+     * @apiNote The cleanup action (if present) should take care not to leak the received segment to external
+     * clients which might access the segment after its backing region of memory is no longer available. Furthermore,
+     * if the provided scope is the scope of an {@linkplain Arena#ofAuto() automatic arena}, the cleanup action
+     * must not prevent the scope from becoming <a href="../../../java/lang/ref/package.html#reachability">unreachable</a>.
+     * A failure to do so will permanently prevent the regions of memory allocated by the automatic arena from being deallocated.
      * <p>
      * This method is <a href="package-summary.html#restricted"><em>restricted</em></a>.
      * Restricted methods are unsafe, and, if used incorrectly, their use might crash
      * the JVM or, worse, silently result in memory corruption. Thus, clients should refrain from depending on
      * restricted methods, and use safe and supported functionalities, where possible.
      *
-     * @param newScope the scope of the returned segment.
+     * @param arena the arena to be associated with the returned segment.
      * @param cleanup the cleanup action that should be executed when the provided arena is closed (can be {@code null}).
      * @return a new memory segment with unbounded size.
      * @throws IllegalArgumentException if {@code newSize < 0}.
@@ -611,17 +628,16 @@ default MemorySegment asSlice(long offset, MemoryLayout layout) {
      * @throws IllegalCallerException If the caller is in a module that does not have native access enabled.
      */
     @CallerSensitive
-    MemorySegment reinterpret(Scope newScope, Consumer<MemorySegment> cleanup);
+    MemorySegment reinterpret(Arena arena, Consumer<MemorySegment> cleanup);
 
     /**
      * Returns a new segment with the same address as this segment, but with the provided size and scope.
-     * As such, the returned segment cannot be accessed after the provided
-     * scope has been invalidated. Moreover, if the provided scope is an arena scope,
-     * the returned segment can be accessed compatibly with the confinement restrictions associated with the
-     * corresponding arena: that is, if the provided scope is the scope of a {@linkplain Arena#ofConfined() confined arena},
+     * As such, the returned segment cannot be accessed after the provided arena has been closed.
+     * Moreover, if the returned segment can be accessed compatibly with the confinement restrictions associated with the
+     * provided arena: that is, if the provided arena is a {@linkplain Arena#ofConfined() confined arena},
      * the returned segment can only be accessed by the arena's owner thread, regardless of the confinement restrictions
-     * associated with this segment. In other words, if the provided scope is an arena scope, this method returns a segment
-     * that behaves as if it had been allocated using the arena associated with the provided scope.
+     * associated with this segment. In other words, this method returns a segment that behaves as if it had been allocated
+     * using the provided arena.
      * <p>
      * Clients can specify an optional cleanup action that should be executed when the provided scope becomes
      * invalid. This cleanup action receives a fresh memory segment that is obtained from this segment as follows:
@@ -643,7 +659,7 @@ default MemorySegment asSlice(long offset, MemoryLayout layout) {
      * A failure to do so will permanently prevent the regions of memory allocated by the automatic arena from being deallocated.
      *
      * @param newSize the size of the returned segment.
-     * @param newScope the scope of the returned segment.
+     * @param arena the arena to be associated with the returned segment.
      * @param cleanup the cleanup action that should be executed when the provided arena is closed (can be {@code null}).
      * @return a new segment that has the same address as this segment, but with new size and its scope set to
      * that of the provided arena.
@@ -653,7 +669,7 @@ default MemorySegment asSlice(long offset, MemoryLayout layout) {
      * @throws IllegalCallerException If the caller is in a module that does not have native access enabled.
      */
     @CallerSensitive
-    MemorySegment reinterpret(long newSize, Scope newScope, Consumer<MemorySegment> cleanup);
+    MemorySegment reinterpret(long newSize, Arena arena, Consumer<MemorySegment> cleanup);
 
     /**
      * {@return {@code true}, if this segment is read-only}

src/java.base/share/classes/java/lang/foreign/SymbolLookup.java

@@ -176,7 +176,7 @@ static SymbolLookup loaderLookup() {
             return addr == 0L ?
                     Optional.empty() :
                     Optional.of(MemorySegment.ofAddress(addr)
-                                    .reinterpret(loaderArena.scope(), null));
+                                    .reinterpret(loaderArena, null));
         };
     }
 
@@ -261,7 +261,7 @@ public void cleanup() {
             return addr == 0L ?
                     Optional.empty() :
                     Optional.of(MemorySegment.ofAddress(addr)
-                            .reinterpret(libArena.scope(), null));
+                            .reinterpret(libArena, null));
         };
     }
 }

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

@@ -122,8 +122,10 @@ public MemorySegment asSlice(long offset, long newSize, long byteAlignment) {
 
     @Override
     @CallerSensitive
-    public final MemorySegment reinterpret(long newSize, Scope newScope, Consumer<MemorySegment> cleanup) {
-        return reinterpretInternal(Reflection.getCallerClass(), newSize, newScope, null);
+    public final MemorySegment reinterpret(long newSize, Arena arena, Consumer<MemorySegment> cleanup) {
+        Objects.requireNonNull(arena);
+        return reinterpretInternal(Reflection.getCallerClass(), newSize,
+                MemorySessionImpl.toMemorySession(arena), null);
     }
 
     @Override
@@ -134,13 +136,14 @@ public final MemorySegment reinterpret(long newSize) {
 
     @Override
     @CallerSensitive
-    public final MemorySegment reinterpret(Scope newScope, Consumer<MemorySegment> cleanup) {
-        return reinterpretInternal(Reflection.getCallerClass(), byteSize(), newScope, cleanup);
+    public final MemorySegment reinterpret(Arena arena, Consumer<MemorySegment> cleanup) {
+        Objects.requireNonNull(arena);
+        return reinterpretInternal(Reflection.getCallerClass(), byteSize(),
+                MemorySessionImpl.toMemorySession(arena), cleanup);
     }
 
     public MemorySegment reinterpretInternal(Class<?> callerClass, long newSize, Scope scope, Consumer<MemorySegment> cleanup) {
         Reflection.ensureNativeAccess(callerClass, MemorySegment.class, "reinterpret");
-        Objects.requireNonNull(scope);
         if (newSize < 0) {
             throw new IllegalArgumentException("newSize < 0");
         }

src/java.base/share/classes/jdk/internal/foreign/abi/Binding.java

@@ -695,7 +695,7 @@ public void interpret(Deque<Object> stack, StoreFunc storeFunc,
                               LoadFunc loadFunc, SegmentAllocator allocator) {
             MemorySegment segment = Utils.longToAddress((long) stack.pop(), size, align);
             if (needsScope) {
-                segment = segment.reinterpret(((Arena) allocator).scope(), null);
+                segment = segment.reinterpret((Arena) allocator, null);
             }
             stack.push(segment);
         }

src/java.base/share/classes/jdk/internal/foreign/abi/UpcallStubs.java

@@ -57,6 +57,6 @@ public void cleanup() {
                 freeUpcallStub(entry);
             }
         });
-        return MemorySegment.ofAddress(entry).reinterpret(arena.scope(), null);
+        return MemorySegment.ofAddress(entry).reinterpret(arena, null);
     }
 }

src/java.base/share/classes/jdk/internal/foreign/abi/fallback/FallbackLinker.java

@@ -191,16 +191,16 @@ private static void doUpcall(MethodHandle target, MemorySegment retPtr, MemorySe
         int numArgs = argLayouts.size();
         MemoryLayout retLayout = data.returnLayout();
         try (Arena upcallArena = Arena.ofConfined()) {
-            MemorySegment argsSeg = argPtrs.reinterpret(numArgs * ADDRESS.byteSize(), upcallArena.scope(), null);
+            MemorySegment argsSeg = argPtrs.reinterpret(numArgs * ADDRESS.byteSize(), upcallArena, null);
             MemorySegment retSeg = retLayout != null
-                ? retPtr.reinterpret(retLayout.byteSize(), upcallArena.scope(), null)
+                ? retPtr.reinterpret(retLayout.byteSize(), upcallArena, null)
                 : null;
 
             Object[] args = new Object[numArgs];
             for (int i = 0; i < numArgs; i++) {
                 MemoryLayout argLayout = argLayouts.get(i);
                 MemorySegment argPtr = argsSeg.getAtIndex(ADDRESS, i)
-                        .reinterpret(argLayout.byteSize(), upcallArena.scope(), null);
+                        .reinterpret(argLayout.byteSize(), upcallArena, null);
                 args[i] = readValue(argPtr, argLayout);
             }
 

src/java.base/share/classes/jdk/internal/foreign/abi/fallback/LibFallback.java

@@ -135,7 +135,7 @@ static MemorySegment createClosure(MemorySegment cif, MethodHandle target,
         long execPtr = ptrs[1];
         long globalTarget = ptrs[2];
 
-        return MemorySegment.ofAddress(execPtr).reinterpret(arena.scope(), unused -> freeClosure(closurePtr, globalTarget));
+        return MemorySegment.ofAddress(execPtr).reinterpret(arena, unused -> freeClosure(closurePtr, globalTarget));
     }
 
     private record UpcallData(MethodHandle target, Thread.UncaughtExceptionHandler handler) {}

test/jdk/java/foreign/TestNative.java

@@ -166,7 +166,7 @@ public void testDefaultAccessModes() {
         MemorySegment addr = allocateMemory(12);
         try (Arena arena = Arena.ofConfined()) {
             MemorySegment mallocSegment = addr.asSlice(0, 12)
-                    .reinterpret(arena.scope(), TestNative::freeMemory);
+                    .reinterpret(arena, TestNative::freeMemory);
             assertFalse(mallocSegment.isReadOnly());
         }
     }
@@ -177,7 +177,7 @@ public void testMallocSegment() {
         MemorySegment mallocSegment = null;
         try (Arena arena = Arena.ofConfined()) {
             mallocSegment = addr.asSlice(0, 12)
-                    .reinterpret(arena.scope(), TestNative::freeMemory);
+                    .reinterpret(arena, TestNative::freeMemory);
             assertEquals(mallocSegment.byteSize(), 12);
             //free here
         }
@@ -197,7 +197,7 @@ public void testBadResize() {
         try (Arena arena = Arena.ofConfined()) {
             MemorySegment segment = arena.allocate(4, 1);
             assertThrows(IllegalArgumentException.class, () -> segment.reinterpret(-1));
-            assertThrows(IllegalArgumentException.class, () -> segment.reinterpret(-1, Arena.ofAuto().scope(), null));
+            assertThrows(IllegalArgumentException.class, () -> segment.reinterpret(-1, Arena.ofAuto(), null));
         }
     }
 

test/jdk/java/foreign/TestNulls.java

@@ -108,8 +108,8 @@ public class TestNulls {
     };
 
     static final Set<String> EXCLUDE_LIST = Set.of(
-            "java.lang.foreign.MemorySegment/reinterpret(java.lang.foreign.MemorySegment$Scope,java.util.function.Consumer)/1/0",
-            "java.lang.foreign.MemorySegment/reinterpret(long,java.lang.foreign.MemorySegment$Scope,java.util.function.Consumer)/2/0"
+            "java.lang.foreign.MemorySegment/reinterpret(java.lang.foreign.Arena,java.util.function.Consumer)/1/0",
+            "java.lang.foreign.MemorySegment/reinterpret(long,java.lang.foreign.Arena,java.util.function.Consumer)/2/0"
     );
 
     static final Set<String> OBJECT_METHODS = Stream.of(Object.class.getMethods())

test/jdk/java/foreign/TestScopedOperations.java

@@ -193,7 +193,7 @@ enum SegmentFactory {
                     throw new AssertionError(ex);
                 }
             }),
-            UNSAFE(session -> MemorySegment.NULL.reinterpret(10, session.scope(), null));
+            UNSAFE(session -> MemorySegment.NULL.reinterpret(10, session, null));
 
             static {
                 try {

test/jdk/java/foreign/TestSegments.java

@@ -219,7 +219,7 @@ public Object[][] scopes() {
 
     @Test(dataProvider = "scopes")
     public void testIsAccessibleBy(Arena arena, boolean isConfined) {
-        MemorySegment segment = MemorySegment.NULL.reinterpret(arena.scope(), null);
+        MemorySegment segment = MemorySegment.NULL.reinterpret(arena, null);
         assertTrue(segment.isAccessibleBy(Thread.currentThread()));
         assertTrue(segment.isAccessibleBy(new Thread()) != isConfined);
     }

test/jdk/java/foreign/TestVarArgs.java

@@ -184,7 +184,7 @@ private static void check(int index, MemorySegment ptr, List<Arg> args) {
         MethodHandle getter = varArg.getter;
         try (Arena arena = Arena.ofConfined()) {
             MemorySegment seg = ptr.asSlice(0, layout)
-                    .reinterpret(arena.scope(), null);
+                    .reinterpret(arena, null);
             Object obj = getter.invoke(seg);
             varArg.check(obj);
         } catch (Throwable e) {

test/jdk/java/foreign/handles/lookup_module/handle/lookup/MethodHandleLookup.java

@@ -55,10 +55,10 @@ static Object[][] restrictedMethods() {
                             MethodType.methodType(MemorySegment.class, long.class)),
                             "MemorySegment::reinterpret/1" },
                     { MethodHandles.lookup().findVirtual(MemorySegment.class, "reinterpret",
-                            MethodType.methodType(MemorySegment.class, MemorySegment.Scope.class, Consumer.class)),
+                            MethodType.methodType(MemorySegment.class, Arena.class, Consumer.class)),
                             "MemorySegment::reinterpret/2" },
                     { MethodHandles.lookup().findVirtual(MemorySegment.class, "reinterpret",
-                            MethodType.methodType(MemorySegment.class, long.class, MemorySegment.Scope.class, Consumer.class)),
+                            MethodType.methodType(MemorySegment.class, long.class, Arena.class, Consumer.class)),
                             "MemorySegment::reinterpret/3" },
                     { MethodHandles.lookup().findStatic(SymbolLookup.class, "libraryLookup",
                             MethodType.methodType(SymbolLookup.class, String.class, Arena.class)),

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

@@ -72,7 +72,7 @@ public long segment_loop_addr_size_session() {
         long res = 0;
         for (int i = 0; i < ITERATIONS; i++) {
             res += MemorySegment.ofAddress(i)
-                    .reinterpret(i % 100, Arena.global().scope(), null).address();
+                    .reinterpret(i % 100, Arena.global(), null).address();
         }
         return res;
     }

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

@@ -87,7 +87,7 @@ public long ptr_to_long() throws Throwable {
 
     @Benchmark
     public long ptr_to_long_new_segment() throws Throwable {
-        MemorySegment newSegment = segment.reinterpret(100, arena.scope(), null);
+        MemorySegment newSegment = segment.reinterpret(100, arena, null);
         return (long)F_PTR_LONG.invokeExact(newSegment);
     }
 
@@ -103,7 +103,7 @@ public long ptr_to_ptr() throws Throwable {
 
     @Benchmark
     public long ptr_to_ptr_new_segment() throws Throwable {
-        MemorySegment newSegment = segment.reinterpret(100, arena.scope(), null);
+        MemorySegment newSegment = segment.reinterpret(100, arena, null);
         return ((MemorySegment)F_PTR_PTR.invokeExact(newSegment)).address();
     }
 }

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

@@ -194,7 +194,7 @@ class SlicingPoolAllocator implements Arena {
 
             public MemorySegment allocate(long byteSize, long byteAlignment) {
                 return slicing.allocate(byteSize, byteAlignment)
-                        .reinterpret(arena.scope(), null);
+                        .reinterpret(arena, null);
             }
 
             @Override

test/micro/org/openjdk/bench/jdk/incubator/vector/TestLoadStoreBytes.java

@@ -167,8 +167,8 @@ public void segmentNativeImplicit() {
   @Benchmark
   public void segmentNativeConfined() {
     try (final var arena = Arena.ofConfined()) {
-      final var srcSegmentConfined = srcSegment.reinterpret(arena.scope(), null);
-      final var dstSegmentConfined = dstSegment.reinterpret(arena.scope(), null);
+      final var srcSegmentConfined = srcSegment.reinterpret(arena, null);
+      final var dstSegmentConfined = dstSegment.reinterpret(arena, null);
 
       for (long i = 0; i < SPECIES.loopBound(srcArray.length); i += SPECIES.length()) {
         var v = ByteVector.fromMemorySegment(SPECIES, srcSegmentConfined, i, ByteOrder.nativeOrder());