8257184: Upstream 8252504: Add a method to MemoryLayout which returns a offset-computing method handle

Reviewed-by: mcimadamore, chegar

Author: JornVernee

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

@@ -31,10 +31,11 @@
 
 import java.lang.constant.Constable;
 import java.lang.constant.DynamicConstantDesc;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 import java.lang.invoke.VarHandle;
 import java.nio.ByteOrder;
 import java.util.EnumSet;
-import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.OptionalLong;
@@ -173,6 +174,19 @@
  * it follows that the memory access var handle {@code valueHandle} will feature an <em>additional</em> {@code long}
  * access coordinate.
  *
+ * <p>A layout path with free dimensions can also be used to create an offset-computing method handle, using the
+ * {@link #bitOffset(PathElement...)} or {@link #byteOffsetHandle(PathElement...)} method. Again, free dimensions are
+ * translated into {@code long} parameters of the created method handle. The method handle can be used to compute the
+ * offsets of elements of a sequence at different indices, by supplying these indices when invoking the method handle.
+ * For instance:
+ *
+ * <blockquote><pre>{@code
+MethodHandle offsetHandle = taggedValues.byteOffsetHandle(PathElement.sequenceElement(),
+                                                          PathElement.groupElement("kind"));
+long offset1 = (long) offsetHandle.invokeExact(1L); // 8
+long offset2 = (long) offsetHandle.invokeExact(2L); // 16
+ * }</pre></blockquote>
+ *
  * <h2>Layout attributes</h2>
  *
  * Layouts can be optionally associated with one or more <em>attributes</em>. A layout attribute forms a <em>name/value</em>
@@ -337,9 +351,6 @@ default long byteAlignment() {
      * Computes the offset, in bits, of the layout selected by a given layout path, where the path is considered rooted in this
      * layout.
      *
-     * @apiNote if the layout path has one (or more) free dimensions,
-     * the offset is computed as if all the indices corresponding to such dimensions were set to {@code 0}.
-     *
      * @param elements the layout path elements.
      * @return The offset, in bits, of the layout selected by the layout path in {@code elements}.
      * @throws IllegalArgumentException if the layout path does not select any layout nested in this layout, or if the
@@ -348,16 +359,46 @@ default long byteAlignment() {
      * @throws UnsupportedOperationException if one of the layouts traversed by the layout path has unspecified size.
      */
     default long bitOffset(PathElement... elements) {
-        return computePathOp(LayoutPath.rootPath(this, MemoryLayout::bitSize), LayoutPath::offset, EnumSet.of(PathKind.SEQUENCE_ELEMENT, PathKind.SEQUENCE_RANGE), elements);
+        return computePathOp(LayoutPath.rootPath(this, MemoryLayout::bitSize), LayoutPath::offset,
+                EnumSet.of(PathKind.SEQUENCE_ELEMENT, PathKind.SEQUENCE_RANGE), elements);
+    }
+
+    /**
+     * Creates a method handle that can be used to compute the offset, in bits, of the layout selected
+     * by a given layout path, where the path is considered rooted in this layout.
+     *
+     * <p>The returned method handle has a return type of {@code long}, and features as many {@code long}
+     * parameter types as there are free dimensions in the provided layout path (see {@link PathElement#sequenceElement()},
+     * where the order of the parameters corresponds to the order of the path elements.
+     * The returned method handle can be used to compute a layout offset similar to {@link #bitOffset(PathElement...)},
+     * but where some sequence indices are specified only when invoking the method handle.
+     *
+     * <p>The final offset returned by the method handle is computed as follows:
+     *
+     * <blockquote><pre>{@code
+    offset = c_1 + c_2 + ... + c_m + (x_1 * s_1) + (x_2 * s_2) + ... + (x_n * s_n)
+     * }</pre></blockquote>
+     *
+     * where {@code x_1}, {@code x_2}, ... {@code x_n} are <em>dynamic</em> values provided as {@code long}
+     * arguments, whereas {@code c_1}, {@code c_2}, ... {@code c_m} and {@code s_0}, {@code s_1}, ... {@code s_n} are
+     * <em>static</em> stride constants which are derived from the layout path.
+     *
+     * @param elements the layout path elements.
+     * @return a method handle that can be used to compute the bit offset of the layout element
+     * specified by the given layout path elements, when supplied with the missing sequence element indices.
+     * @throws IllegalArgumentException if the layout path contains one or more path elements that select
+     * multiple sequence element indices (see {@link PathElement#sequenceElement(long, long)}).
+     * @throws UnsupportedOperationException if one of the layouts traversed by the layout path has unspecified size.
+     */
+    default MethodHandle bitOffsetHandle(PathElement... elements) {
+        return computePathOp(LayoutPath.rootPath(this, MemoryLayout::bitSize), LayoutPath::offsetHandle,
+                EnumSet.of(PathKind.SEQUENCE_RANGE), elements);
     }
 
     /**
      * Computes the offset, in bytes, of the layout selected by a given layout path, where the path is considered rooted in this
      * layout.
      *
-     * @apiNote if the layout path has one (or more) free dimensions,
-     * the offset is computed as if all the indices corresponding to such dimensions were set to {@code 0}.
-     *
      * @param elements the layout path elements.
      * @return The offset, in bytes, of the layout selected by the layout path in {@code elements}.
      * @throws IllegalArgumentException if the layout path does not select any layout nested in this layout, or if the
@@ -367,8 +408,44 @@ default long bitOffset(PathElement... elements) {
      * or if {@code bitOffset(elements)} is not a multiple of 8.
      */
     default long byteOffset(PathElement... elements) {
-        return Utils.bitsToBytesOrThrow(bitOffset(elements),
-                () -> new UnsupportedOperationException("Cannot compute byte offset; bit offset is not a multiple of 8"));
+        return Utils.bitsToBytesOrThrow(bitOffset(elements), Utils.bitsToBytesThrowOffset);
+    }
+
+    /**
+     * Creates a method handle that can be used to compute the offset, in bytes, of the layout selected
+     * by a given layout path, where the path is considered rooted in this layout.
+     *
+     * <p>The returned method handle has a return type of {@code long}, and features as many {@code long}
+     * parameter types as there are free dimensions in the provided layout path (see {@link PathElement#sequenceElement()},
+     * where the order of the parameters corresponds to the order of the path elements.
+     * The returned method handle can be used to compute a layout offset similar to {@link #byteOffset(PathElement...)},
+     * but where some sequence indices are specified only when invoking the method handle.
+     *
+     * <p>The final offset returned by the method handle is computed as follows:
+     *
+     * <blockquote><pre>{@code
+    bitOffset = c_1 + c_2 + ... + c_m + (x_1 * s_1) + (x_2 * s_2) + ... + (x_n * s_n)
+    offset = bitOffset / 8
+     * }</pre></blockquote>
+     *
+     * where {@code x_1}, {@code x_2}, ... {@code x_n} are <em>dynamic</em> values provided as {@code long}
+     * arguments, whereas {@code c_1}, {@code c_2}, ... {@code c_m} and {@code s_0}, {@code s_1}, ... {@code s_n} are
+     * <em>static</em> stride constants which are derived from the layout path.
+     *
+     * <p>The method handle will throw an {@link UnsupportedOperationException} if the computed
+     * offset in bits is not a multiple of 8.
+     *
+     * @param elements the layout path elements.
+     * @return a method handle that can be used to compute the byte offset of the layout element
+     * specified by the given layout path elements, when supplied with the missing sequence element indices.
+     * @throws IllegalArgumentException if the layout path contains one or more path elements that select
+     * multiple sequence element indices (see {@link PathElement#sequenceElement(long, long)}).
+     * @throws UnsupportedOperationException if one of the layouts traversed by the layout path has unspecified size.
+     */
+    default MethodHandle byteOffsetHandle(PathElement... elements) {
+        MethodHandle mh = bitOffsetHandle(elements);
+        mh = MethodHandles.filterReturnValue(mh, Utils.MH_bitsToBytesOrThrowForOffset);
+        return mh;
     }
 
     /**

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

@@ -61,11 +61,17 @@ public class LayoutPath {
     private static final JavaLangInvokeAccess JLI = SharedSecrets.getJavaLangInvokeAccess();
 
     private static final MethodHandle ADD_STRIDE;
+    private static final MethodHandle MH_ADD_SCALED_OFFSET;
+
+    private static final int UNSPECIFIED_ELEM_INDEX = -1;
 
     static {
         try {
-            ADD_STRIDE = MethodHandles.lookup().findStatic(LayoutPath.class, "addStride",
+            MethodHandles.Lookup lookup = MethodHandles.lookup();
+            ADD_STRIDE = lookup.findStatic(LayoutPath.class, "addStride",
                     MethodType.methodType(long.class, MemorySegment.class, long.class, long.class, long.class));
+            MH_ADD_SCALED_OFFSET = lookup.findStatic(LayoutPath.class, "addScaledOffset",
+                    MethodType.methodType(long.class, long.class, long.class, long.class));
         } catch (Throwable ex) {
             throw new ExceptionInInitializerError(ex);
         }
@@ -93,7 +99,7 @@ public LayoutPath sequenceElement() {
         check(SequenceLayout.class, "attempting to select a sequence element from a non-sequence layout");
         SequenceLayout seq = (SequenceLayout)layout;
         MemoryLayout elem = seq.elementLayout();
-        return LayoutPath.nestedPath(elem, offset, addStride(sizeFunc.applyAsLong(elem)), -1, this);
+        return LayoutPath.nestedPath(elem, offset, addStride(sizeFunc.applyAsLong(elem)), UNSPECIFIED_ELEM_INDEX, this);
     }
 
     public LayoutPath sequenceElement(long start, long step) {
@@ -102,7 +108,8 @@ public LayoutPath sequenceElement(long start, long step) {
         checkSequenceBounds(seq, start);
         MemoryLayout elem = seq.elementLayout();
         long elemSize = sizeFunc.applyAsLong(elem);
-        return LayoutPath.nestedPath(elem, offset + (start * elemSize), addStride(elemSize * step), -1, this);
+        return LayoutPath.nestedPath(elem, offset + (start * elemSize), addStride(elemSize * step),
+                UNSPECIFIED_ELEM_INDEX, this);
     }
 
     public LayoutPath sequenceElement(long index) {
@@ -177,6 +184,22 @@ public VarHandle dereferenceHandle(Class<?> carrier) {
         return handle;
     }
 
+    private static long addScaledOffset(long base, long index, long stride) {
+        return base + (stride * index);
+    }
+
+    public MethodHandle offsetHandle() {
+        MethodHandle mh = MethodHandles.identity(long.class);
+        for (int i = strides.length - 1; i >=0; i--) {
+            MethodHandle collector = MethodHandles.insertArguments(MH_ADD_SCALED_OFFSET, 2, strides[i]);
+            // (J, ...) -> J to (J, J, ...) -> J
+            // i.e. new coord is prefixed. Last coord will correspond to innermost layout
+            mh = MethodHandles.collectArguments(mh, 0, collector);
+        }
+        mh = MethodHandles.insertArguments(mh, 0, offset);
+        return mh;
+    }
+
     public MemoryLayout layout() {
         return layout;
     }

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

@@ -53,11 +53,21 @@ public final class Utils {
             .orElse("deny");
 
     private static final MethodHandle SEGMENT_FILTER;
+    public static final MethodHandle MH_bitsToBytesOrThrowForOffset;
+
+    public static final Supplier<RuntimeException> bitsToBytesThrowOffset
+        = () -> new UnsupportedOperationException("Cannot compute byte offset; bit offset is not a multiple of 8");
 
     static {
         try {
-            SEGMENT_FILTER = MethodHandles.lookup().findStatic(Utils.class, "filterSegment",
+            MethodHandles.Lookup lookup = MethodHandles.lookup();
+            SEGMENT_FILTER = lookup.findStatic(Utils.class, "filterSegment",
                     MethodType.methodType(MemorySegmentProxy.class, MemorySegment.class));
+            MH_bitsToBytesOrThrowForOffset = MethodHandles.insertArguments(
+                lookup.findStatic(Utils.class, "bitsToBytesOrThrow",
+                    MethodType.methodType(long.class, long.class, Supplier.class)),
+                1,
+                bitsToBytesThrowOffset);
         } catch (Throwable ex) {
             throw new ExceptionInInitializerError(ex);
         }