Simplify zeroing allocator

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

@@ -30,7 +30,6 @@
 import java.util.Objects;
 
 import jdk.internal.foreign.AbstractMemorySegmentImpl;
-import jdk.internal.foreign.ArenaImpl;
 import jdk.internal.foreign.SlicingAllocator;
 import jdk.internal.foreign.StringSupport;
 import jdk.internal.foreign.ZeroingAllocator;
@@ -735,22 +734,22 @@ private static void assertWritable(MemorySegment segment) {
 
     @ForceInline
     private MemorySegment allocateNoInit(long byteSize) {
-        return this instanceof ArenaImpl arenaImpl ?
-                arenaImpl.allocateNoInit(byteSize, 1) :
+        return this instanceof ZeroingAllocator zeroingAllocator ?
+                zeroingAllocator.allocateRaw(byteSize, 1) :
                 allocate(byteSize);
     }
 
     @ForceInline
     private MemorySegment allocateNoInit(MemoryLayout layout) {
-        return this instanceof ArenaImpl arenaImpl ?
-                arenaImpl.allocateNoInit(layout.byteSize(), layout.byteAlignment()) :
+        return this instanceof ZeroingAllocator zeroingAllocator ?
+                zeroingAllocator.allocateRaw(layout.byteSize(), layout.byteAlignment()) :
                 allocate(layout);
     }
 
     @ForceInline
     private MemorySegment allocateNoInit(MemoryLayout layout, long size) {
-        return this instanceof ArenaImpl arenaImpl ?
-                arenaImpl.allocateNoInit(layout.byteSize() * size, layout.byteAlignment()) :
+        return this instanceof ZeroingAllocator zeroingAllocator ?
+                zeroingAllocator.allocateRaw(layout.byteSize() * size, layout.byteAlignment()) :
                 allocate(layout, size);
     }
 }

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

@@ -26,10 +26,8 @@
 package jdk.internal.foreign;
 
 import java.lang.foreign.Arena;
-import java.lang.foreign.MemoryLayout;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.MemorySegment.Scope;
-import java.util.Objects;
 
 public final class ArenaImpl implements Arena, ZeroingAllocator {
 
@@ -50,14 +48,14 @@ public void close() {
         session.close();
     }
 
-    public MemorySegment allocateNoInit(long byteSize, long byteAlignment) {
+    @Override
+    public MemorySegment allocateRaw(long byteSize, long byteAlignment) {
         Utils.checkAllocationSizeAndAlign(byteSize, byteAlignment);
         return SegmentFactories.allocateSegment(byteSize, byteAlignment, session, shouldReserveMemory);
     }
 
-    @Override
     public MemorySegment allocate(long byteSize, long byteAlignment) {
-        MemorySegment segment = allocateNoInit(byteSize, byteAlignment);
-        return segment.fill((byte)0);
+        MemorySegment segment = allocateRaw(byteSize, byteAlignment);
+        return segment.fill((byte) 0);
     }
 }

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

@@ -1,55 +1,23 @@
 package jdk.internal.foreign;
 
-import jdk.internal.vm.annotation.Stable;
-
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
-import java.lang.foreign.ValueLayout;
 
 /**
- * Marker interface for allocators that provide zeroing semantics.
+ * Internal interface for allocators that provide zeroing semantics.
  */
 public interface ZeroingAllocator extends SegmentAllocator {
 
-    static ZeroingAllocator of(SegmentAllocator allocator) {
-        if (allocator instanceof ZeroingAllocator zeroingAllocator) {
-            // already a zeroing allocator
-            return zeroingAllocator;
-        } else if (allocator instanceof MemorySegment ||
-                   allocator instanceof SlicingAllocator) {
-            // prefix and slicing allocators do not zero
-            return (byteSize, byteAlignment) -> allocator.allocate(byteSize, byteAlignment).fill((byte)0);
-        } else {
-            return new ZeroingWrapper(allocator);
-        }
-    }
-
-    class ZeroingWrapper implements ZeroingAllocator {
-        static final MemorySegment ZERO = MemorySegment.ofArray(new long[1024]);
-        @Stable
-        static final ValueLayout[] ALIGNED_LAYOUTS;
+    MemorySegment allocateRaw(long byteSize, long byteAlign);
 
-        final SegmentAllocator allocator;
-
-        ZeroingWrapper(SegmentAllocator allocator) {
-            this.allocator = allocator;
-        }
-
-        static {
-            ALIGNED_LAYOUTS = new ValueLayout[8];
-            ALIGNED_LAYOUTS[0] = ValueLayout.JAVA_BYTE;
-            ALIGNED_LAYOUTS[1] = ValueLayout.JAVA_SHORT;
-            ALIGNED_LAYOUTS[3] = ValueLayout.JAVA_INT;
-            ALIGNED_LAYOUTS[7] = ValueLayout.JAVA_LONG;
-        }
+    @Override
+    default MemorySegment allocate(long byteSize, long byteAlign) {
+        return allocateRaw(byteSize, byteAlign).fill((byte)0);
+    }
 
-        public MemorySegment allocate(long size, long align) {
-            if (align > 8 || size > ZERO.byteSize()) {
-                // slow path (might perform double zeroing)
-                return allocator.allocate(size, align).fill((byte) 0);
-            } else {
-                return allocator.allocateFrom(ALIGNED_LAYOUTS[(int)align - 1], ZERO, ALIGNED_LAYOUTS[(int)align - 1], 0, size);
-            }
-        };
+    static ZeroingAllocator of(SegmentAllocator allocator) {
+        return (allocator instanceof ZeroingAllocator zeroingAllocator) ?
+                zeroingAllocator : // already a zeroing allocator
+                allocator::allocate;
     }
 }