Add MemorySegment::fill

Reviewed-by: jlaskey, jvernee, psandoz, mcimadamore

src/jdk.incubator.foreign/share/classes/jdk/incubator/foreign/MemorySegment.java

@@ -283,6 +283,31 @@
      */
     void close();
 
+    /**
+     * Fills a value into this memory segment.
+     * <p>
+     * More specifically, the given value is filled into each address of this
+     * segment. Equivalent to (but likely more efficient than) the following code:
+     *
+     * <blockquote><pre>
+     * byteHandle = MemoryLayout.ofSequence(MemoryLayouts.JAVA_BYTE)
+     *         .varHandle(byte.class, MemoryLayout.PathElement.sequenceElement());
+     * for (long l = 0; l < segment.byteSize(); l++) {
+     *     byteHandle.set(segment.baseAddress(), l, value);
+     * }</pre></blockquote>
+     * without any regard or guarantees on the ordering of particular memory
+     * elements being set.
+     * <p>
+     * Fill can be useful to initialize or reset the memory of a segment.
+     *
+     * @param value the value to fill into this segment
+     * @return this memory segment
+     * @throws IllegalStateException if this segment is not <em>alive</em>, or if access occurs from a thread other than the
+     * thread owning this segment
+     * @throws UnsupportedOperationException if this segment does not support the {@link #WRITE} access mode
+     */
+    MemorySegment fill(byte value);
+
     /**
      * Wraps this segment in a {@link ByteBuffer}. Some of the properties of the returned buffer are linked to
      * the properties of this segment. For instance, if this segment is <em>immutable</em>

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

@@ -32,6 +32,7 @@
 import jdk.internal.access.SharedSecrets;
 import jdk.internal.access.foreign.MemorySegmentProxy;
 import jdk.internal.access.foreign.UnmapperProxy;
+import jdk.internal.misc.Unsafe;
 import jdk.internal.vm.annotation.ForceInline;
 import sun.security.action.GetPropertyAction;
 
@@ -55,6 +56,8 @@
  */
 public abstract class AbstractMemorySegmentImpl implements MemorySegment, MemorySegmentProxy {
 
+    private static final Unsafe UNSAFE = Unsafe.getUnsafe();
+
     private static final boolean enableSmallSegments =
             Boolean.parseBoolean(GetPropertyAction.privilegedGetProperty("jdk.incubator.foreign.SmallSegments", "true"));
 
@@ -113,6 +116,13 @@ private AbstractMemorySegmentImpl asSliceNoCheck(long offset, long newSize) {
                 (AbstractMemorySegmentImpl)segment.withAccessModes(segment.accessModes() & ~CLOSE));
     }
 
+    @Override
+    public final MemorySegment fill(byte value){
+        checkRange(0, length, true);
+        UNSAFE.setMemory(base(), min(), length, value);
+        return this;
+    }
+
     @Override
     @ForceInline
     public final MemoryAddress baseAddress() {

test/jdk/java/foreign/TestSegments.java

@@ -32,19 +32,20 @@
 import jdk.incubator.foreign.MemorySegment;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
-
 import java.lang.invoke.VarHandle;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Spliterator;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.LongFunction;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
-
+import static jdk.incubator.foreign.MemorySegment.WRITE;
 import static org.testng.Assert.*;
 
 public class TestSegments {
@@ -161,6 +162,93 @@ public void testAccessModes(int accessModes) {
         }
     }
 
+    @DataProvider(name = "segmentFactories")
+    public Object[][] segmentFactories() {
+        List<Supplier<MemorySegment>> l = List.of(
+                () -> MemorySegment.ofArray(new byte[] { 0x00, 0x01, 0x02, 0x03 }),
+                () -> MemorySegment.ofArray(new char[] {'a', 'b', 'c', 'd' }),
+                () -> MemorySegment.ofArray(new double[] { 1d, 2d, 3d, 4d} ),
+                () -> MemorySegment.ofArray(new float[] { 1.0f, 2.0f, 3.0f, 4.0f }),
+                () -> MemorySegment.ofArray(new int[] { 1, 2, 3, 4 }),
+                () -> MemorySegment.ofArray(new long[] { 1l, 2l, 3l, 4l } ),
+                () -> MemorySegment.ofArray(new short[] { 1, 2, 3, 4 } ),
+                () -> MemorySegment.allocateNative(4),
+                () -> MemorySegment.allocateNative(4, 8),
+                () -> MemorySegment.allocateNative(MemoryLayout.ofValueBits(32, ByteOrder.nativeOrder()))
+        );
+        return l.stream().map(s -> new Object[] { s }).toArray(Object[][]::new);
+    }
+
+    @Test(dataProvider = "segmentFactories")
+    public void testFill(Supplier<MemorySegment> memorySegmentSupplier) {
+        VarHandle byteHandle = MemoryLayout.ofSequence(MemoryLayouts.JAVA_BYTE)
+                .varHandle(byte.class, MemoryLayout.PathElement.sequenceElement());
+
+        for (byte value : new byte[] {(byte) 0xFF, (byte) 0x00, (byte) 0x45}) {
+            try (MemorySegment segment = memorySegmentSupplier.get()) {
+                segment.fill(value);
+                for (long l = 0; l < segment.byteSize(); l++) {
+                    assertEquals((byte) byteHandle.get(segment.baseAddress(), l), value);
+                }
+
+                // fill a slice
+                var sliceSegment = segment.asSlice(1, segment.byteSize() - 2).fill((byte) ~value);
+                for (long l = 0; l < sliceSegment.byteSize(); l++) {
+                    assertEquals((byte) byteHandle.get(sliceSegment.baseAddress(), l), ~value);
+                }
+                // assert enclosing slice
+                assertEquals((byte) byteHandle.get(segment.baseAddress(), 0L), value);
+                for (long l = 1; l < segment.byteSize() - 2; l++) {
+                    assertEquals((byte) byteHandle.get(segment.baseAddress(), l), (byte) ~value);
+                }
+                assertEquals((byte) byteHandle.get(segment.baseAddress(), segment.byteSize() - 1L), value);
+            }
+        }
+    }
+
+    @Test(dataProvider = "segmentFactories", expectedExceptions = IllegalStateException.class)
+    public void testFillClosed(Supplier<MemorySegment> memorySegmentSupplier) {
+        MemorySegment segment = memorySegmentSupplier.get();
+        segment.close();
+        segment.fill((byte) 0xFF);
+    }
+
+    @Test(dataProvider = "segmentFactories", expectedExceptions = UnsupportedOperationException.class)
+    public void testFillIllegalAccessMode(Supplier<MemorySegment> memorySegmentSupplier) {
+        try (MemorySegment segment = memorySegmentSupplier.get()) {
+            segment.withAccessModes(segment.accessModes() & ~WRITE).fill((byte) 0xFF);
+        }
+    }
+
+    @Test(dataProvider = "segmentFactories")
+    public void testFillThread(Supplier<MemorySegment> memorySegmentSupplier) throws Exception {
+        try (MemorySegment segment = memorySegmentSupplier.get()) {
+            AtomicReference<RuntimeException> exception = new AtomicReference<>();
+            Runnable action = () -> {
+                try {
+                    segment.fill((byte) 0xBA);
+                } catch (RuntimeException e) {
+                    exception.set(e);
+                }
+            };
+            Thread thread = new Thread(action);
+            thread.start();
+            thread.join();
+
+            RuntimeException e = exception.get();
+            if (!(e instanceof IllegalStateException)) {
+                throw e;
+            }
+        }
+    }
+
+    @Test
+    public void testFillEmpty() {
+        MemorySegment.ofArray(new byte[] { }).fill((byte) 0xFF);
+        MemorySegment.ofArray(new byte[2]).asSlice(0, 0).fill((byte) 0xFF);
+        MemorySegment.ofByteBuffer(ByteBuffer.allocateDirect(0)).fill((byte) 0xFF);
+    }
+
     @Test(expectedExceptions = IllegalArgumentException.class)
     public void testBadWithAccessModes() {
         int[] arr = new int[1];
@@ -232,6 +320,7 @@ MemoryLayout make(long size) {
 
         final static List<String> CONFINED_NAMES = List.of(
                 "close",
+                "fill",
                 "toByteArray",
                 "withOwnerThread"
         );