8331734: Atomic MemorySegment VarHandle operations fails for element layouts

Reviewed-by: pminborg, psandoz

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

@@ -26,6 +26,7 @@
 package jdk.internal.foreign;
 
 import jdk.internal.vm.annotation.ForceInline;
+
 import java.lang.foreign.AddressLayout;
 import java.lang.foreign.GroupLayout;
 import java.lang.foreign.MemoryLayout;
@@ -204,10 +205,7 @@ public class LayoutPath {
                     String.format("Path does not select a value layout: %s", breadcrumbs()));
         }
 
-        // If we have an enclosing layout, drop the alignment check for the accessed element,
+        VarHandle handle = valueLayout.varHandle();
-        // we check the root layout instead
-        ValueLayout accessedLayout = enclosing != null ? valueLayout.withByteAlignment(1) : valueLayout;
-        VarHandle handle = accessedLayout.varHandle();
         handle = MethodHandles.collectCoordinates(handle, 1, offsetHandle());
 
         // we only have to check the alignment of the root layout for the first dereference we do,

test/jdk/java/foreign/TestAccessModes.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2023, 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
@@ -29,11 +29,7 @@
  * @run testng/othervm -Djava.lang.invoke.VarHandle.VAR_HANDLE_GUARDS=false -Djava.lang.invoke.VarHandle.VAR_HANDLE_IDENTITY_ADAPT=true -Xverify:all TestAccessModes
  */
 
-import java.lang.foreign.AddressLayout;
+import java.lang.foreign.*;
-import java.lang.foreign.Arena;
-import java.lang.foreign.MemoryLayout;
-import java.lang.foreign.MemorySegment;
-import java.lang.foreign.ValueLayout;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
@@ -50,10 +46,12 @@ import static org.testng.Assert.*;
 public class TestAccessModes {
 
     @Test(dataProvider = "segmentsAndLayoutsAndModes")
-    public void testAccessModes(MemorySegment segment, ValueLayout layout, AccessMode mode) throws Throwable {
+    public void testAccessModes(MemorySegment segment, MemoryLayout layout, AccessMode mode) throws Throwable {
-        VarHandle varHandle = layout.varHandle();
+        VarHandle varHandle = layout instanceof ValueLayout ?
+                layout.varHandle() :
+                layout.varHandle(MemoryLayout.PathElement.groupElement(0));
         MethodHandle methodHandle = varHandle.toMethodHandle(mode);
-        boolean compatible = AccessModeKind.supportedModes(layout).contains(AccessModeKind.of(mode));
+        boolean compatible = AccessModeKind.supportedModes(accessLayout(layout)).contains(AccessModeKind.of(mode));
         try {
             Object o = methodHandle.invokeWithArguments(makeArgs(segment, varHandle.accessModeType(mode)));
             assertTrue(compatible);
@@ -61,10 +59,19 @@ public class TestAccessModes {
             assertFalse(compatible);
         } catch (IllegalArgumentException ex) {
             // access is unaligned, but access mode is supported
-            assertTrue(compatible);
+            assertTrue(compatible ||
+                    (layout instanceof GroupLayout && segment.maxByteAlignment() < layout.byteAlignment()));
         }
     }
 
+    static ValueLayout accessLayout(MemoryLayout layout) {
+        return switch (layout) {
+            case ValueLayout vl -> vl;
+            case GroupLayout gl -> accessLayout(gl.memberLayouts().get(0));
+            default -> throw new IllegalStateException();
+        };
+    }
+
     Object[] makeArgs(MemorySegment segment, MethodType type) throws Throwable {
         List<Object> args = new ArrayList<>();
         args.add(segment);
@@ -145,6 +152,7 @@ public class TestAccessModes {
         for (MemoryLayout layout : valueLayouts) {
             for (int align : new int[] { 1, 2, 4, 8 }) {
                 layouts.add(layout.withByteAlignment(align));
+                layouts.add(MemoryLayout.structLayout(layout.withByteAlignment(align)));
             }
         }
         return layouts.toArray(new MemoryLayout[0]);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2020, 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
@@ -23,6 +23,7 @@
 package org.openjdk.bench.java.lang.foreign;
 
 import java.lang.foreign.Arena;
+import java.lang.foreign.MemoryLayout;
 import java.lang.foreign.MemorySegment;
 
 import org.openjdk.jmh.annotations.Benchmark;
@@ -37,6 +38,8 @@ import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Warmup;
 import sun.misc.Unsafe;
 
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.util.concurrent.TimeUnit;
@@ -57,6 +60,13 @@ public class LoopOverNonConstant extends JavaLayouts {
     static final int CARRIER_SIZE = (int)JAVA_INT.byteSize();
     static final int ALLOC_SIZE = ELEM_SIZE * CARRIER_SIZE;
 
+    static final VarHandle VH_SEQ_INT = bindToZeroOffset(MemoryLayout.sequenceLayout(ELEM_SIZE, JAVA_INT).varHandle(PathElement.sequenceElement()));
+    static final VarHandle VH_SEQ_INT_UNALIGNED = bindToZeroOffset(MemoryLayout.sequenceLayout(ELEM_SIZE, JAVA_INT.withByteAlignment(1)).varHandle(PathElement.sequenceElement()));
+
+    static VarHandle bindToZeroOffset(VarHandle varHandle) {
+        return MethodHandles.insertCoordinates(varHandle, 1, 0L);
+    }
+
     Arena arena;
     MemorySegment segment;
     long unsafe_addr;
@@ -132,6 +142,24 @@ public class LoopOverNonConstant extends JavaLayouts {
         return sum;
     }
 
+    @Benchmark
+    public int segment_loop_nested() {
+        int sum = 0;
+        for (int i = 0; i < ELEM_SIZE; i++) {
+            sum += (int) VH_SEQ_INT.get(segment, (long) i);
+        }
+        return sum;
+    }
+
+    @Benchmark
+    public int segment_loop_nested_unaligned() {
+        int sum = 0;
+        for (int i = 0; i < ELEM_SIZE; i++) {
+            sum += (int) VH_SEQ_INT_UNALIGNED.get(segment, (long) i);
+        }
+        return sum;
+    }
+
     @Benchmark
     public int segment_loop_instance() {
         int sum = 0;