8356126: Duplication handling and optimization of CaptureCallState

Reviewed-by: jvernee

Author: liach

src/hotspot/share/prims/downcallLinker.cpp

@@ -32,7 +32,7 @@
 
 // We call this from _thread_in_native, right after a downcall
 JVM_LEAF(void, DowncallLinker::capture_state(int32_t* value_ptr, int captured_state_mask))
-  // keep in synch with jdk.internal.foreign.abi.PreservableValues
+  // keep in synch with jdk.internal.foreign.abi.CapturableState
   enum PreservableValues {
     NONE = 0,
     GET_LAST_ERROR = 1,

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

@@ -34,11 +34,7 @@
 
 import java.lang.invoke.MethodHandle;
 import java.util.Map;
-import java.util.Objects;
-import java.util.Set;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 /**
  * A linker provides access to foreign functions from Java code, and access to Java code
@@ -859,12 +855,11 @@ static Option firstVariadicArg(int index) {
          * @see #captureStateLayout()
          */
         static Option captureCallState(String... capturedState) {
-            int set = Stream.of(Objects.requireNonNull(capturedState))
-                    .map(Objects::requireNonNull)
-                    .map(CapturableState::forName)
-                    .mapToInt(state -> 1 << state.ordinal())
-                    .sum();
-            return new LinkerOptions.CaptureCallState(set);
+            int mask = 0;
+            for (var state : capturedState) {
+                mask |= CapturableState.maskFromName(state);
+            }
+            return new LinkerOptions.CaptureCallState(mask);
         }
 
          /**

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

@@ -185,9 +185,7 @@ public boolean hasReturnBindings() {
     }
 
     public int capturedStateMask() {
-        return linkerOptions.capturedCallState()
-                .mapToInt(CapturableState::mask)
-                .reduce(0, (a, b) -> a | b);
+        return linkerOptions.capturedCallStateMask();
     }
 
     public boolean needsTransition() {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2025, 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,67 +29,68 @@
 
 import java.lang.foreign.MemoryLayout;
 import java.lang.foreign.StructLayout;
-import java.lang.foreign.ValueLayout;
-import java.util.List;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
+import java.util.ArrayList;
+import java.util.Map;
 
 import static java.lang.foreign.ValueLayout.JAVA_INT;
 
-public enum CapturableState {
-    GET_LAST_ERROR    ("GetLastError",    JAVA_INT, 1 << 0, OperatingSystem.isWindows()),
-    WSA_GET_LAST_ERROR("WSAGetLastError", JAVA_INT, 1 << 1, OperatingSystem.isWindows()),
-    ERRNO             ("errno",           JAVA_INT, 1 << 2, true);
+/**
+ * Utility class for the call states to capture.
+ */
+public final class CapturableState {
 
-    public static final StructLayout LAYOUT = MemoryLayout.structLayout(
-        supportedStates().map(CapturableState::layout).toArray(MemoryLayout[]::new));
-    public static final List<CapturableState> BY_ORDINAL = List.of(values());
+    public static final StructLayout LAYOUT;
+    // Keep in synch with DowncallLinker::capture_state in downcallLinker.cpp
+    private static final Map<String, Integer> MASKS;
 
     static {
-        assert (BY_ORDINAL.size() < Integer.SIZE); // Update LinkerOptions.CaptureCallState
-    }
-
-    private final String stateName;
-    private final ValueLayout layout;
-    private final int mask;
-    private final boolean isSupported;
-
-    CapturableState(String stateName, ValueLayout layout, int mask, boolean isSupported) {
-        this.stateName = stateName;
-        this.layout = layout.withName(stateName);
-        this.mask = mask;
-        this.isSupported = isSupported;
-    }
-
-    private static Stream<CapturableState> supportedStates() {
-        return Stream.of(values()).filter(CapturableState::isSupported);
-    }
-
-    public static CapturableState forName(String name) {
-        return Stream.of(values())
-                .filter(stl -> stl.stateName().equals(name))
-                .filter(CapturableState::isSupported)
-                .findAny()
-                .orElseThrow(() -> new IllegalArgumentException(
-                        "Unknown name: " + name +", must be one of: "
-                            + supportedStates()
-                                    .map(CapturableState::stateName)
-                                    .collect(Collectors.joining(", "))));
-    }
-
-    public String stateName() {
-        return stateName;
+        if (OperatingSystem.isWindows()) {
+            LAYOUT = MemoryLayout.structLayout(
+                    JAVA_INT.withName("GetLastError"),
+                    JAVA_INT.withName("WSAGetLastError"),
+                    JAVA_INT.withName("errno"));
+            MASKS = Map.of(
+                    "GetLastError",    1 << 0,
+                    "WSAGetLastError", 1 << 1,
+                    "errno",           1 << 2
+            );
+        } else {
+            LAYOUT = MemoryLayout.structLayout(
+                    JAVA_INT.withName("errno"));
+            MASKS = Map.of(
+                    "errno",           1 << 2
+            );
+        }
     }
 
-    public ValueLayout layout() {
-        return layout;
+    private CapturableState() {
     }
 
-    public int mask() {
-        return mask;
+    /**
+     * Returns the mask for a supported capturable state, or throw an
+     * IllegalArgumentException if no supported state with this name exists.
+     */
+    public static int maskFromName(String name) {
+        var ret = MASKS.get(name);
+        if (ret == null) {
+            throw new IllegalArgumentException(
+                    "Unknown name: " + name + ", must be one of: "
+                            + MASKS.keySet());
+        }
+        return ret;
     }
 
-    public boolean isSupported() {
-        return isSupported;
+    /**
+     * Returns a collection-like display string for a captured state mask.
+     * Enclosed with brackets.
+     */
+    public static String displayString(int mask) {
+        var displayList = new ArrayList<>(); // unordered
+        for (var e : MASKS.entrySet()) {
+            if ((mask & e.getValue()) != 0) {
+                displayList.add(e.getKey());
+            }
+        }
+        return displayList.toString();
     }
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2025, 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
@@ -84,9 +84,9 @@ public boolean hasCapturedCallState() {
         return getOption(CaptureCallState.class) != null;
     }
 
-    public Stream<CapturableState> capturedCallState() {
+    public int capturedCallStateMask() {
         CaptureCallState stl = getOption(CaptureCallState.class);
-        return stl == null ? Stream.empty() : stl.saved().stream();
+        return stl == null ? 0 : stl.mask();
     }
 
     public boolean isVariadicFunction() {
@@ -150,35 +150,25 @@ public boolean equals(Object obj) {
         }
     }
 
-    public record CaptureCallState(int compact) implements LinkerOptionImpl {
+    public record CaptureCallState(int mask) implements LinkerOptionImpl {
         @Override
         public void validateForDowncall(FunctionDescriptor descriptor) {
             // done during construction
         }
 
-        public Set<CapturableState> saved() {
-            var set = EnumSet.noneOf(CapturableState.class);
-            int mask = compact;
-            int i = 0;
-            while (mask != 0) {
-                if ((mask & 1) == 1) {
-                    set.add(CapturableState.BY_ORDINAL.get(i));
-                }
-                mask >>= 1;
-                i++;
-            }
-            return set;
-        }
-
-
         @Override
         public boolean equals(Object obj) {
-            return obj instanceof CaptureCallState that && compact == that.compact;
+            return obj instanceof CaptureCallState that && mask == that.mask;
         }
 
         @Override
         public int hashCode() {
-            return compact;
+            return mask;
+        }
+
+        @Override
+        public String toString() {
+            return "CaptureCallState" + CapturableState.displayString(mask);
         }
     }
 

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

@@ -84,9 +84,7 @@ protected MethodHandle arrangeDowncall(MethodType inferredMethodType, FunctionDe
         assertNotEmpty(function);
         MemorySegment cif = makeCif(inferredMethodType, function, options, Arena.ofAuto());
 
-        int capturedStateMask = options.capturedCallState()
-                .mapToInt(CapturableState::mask)
-                .reduce(0, (a, b) -> a | b);
+        int capturedStateMask = options.capturedCallStateMask();
         DowncallData invData = new DowncallData(cif, function.returnLayout().orElse(null),
                 function.argumentLayouts(), capturedStateMask, options.allowsHeapAccess());
 

test/jdk/java/foreign/capturecallstate/TestCaptureCallState.java

@@ -23,6 +23,7 @@
 
 /*
  * @test
+ * @bug 8356126
  * @library ../ /test/lib
  * @run testng/othervm/native --enable-native-access=ALL-UNNAMED TestCaptureCallState
  */
@@ -47,8 +48,7 @@
 import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
 import static java.lang.foreign.ValueLayout.JAVA_INT;
 import static java.lang.foreign.ValueLayout.JAVA_LONG;
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.*;
 
 public class TestCaptureCallState extends NativeTestHelper {
 
@@ -61,6 +61,17 @@ public class TestCaptureCallState extends NativeTestHelper {
         }
     }
 
+    // Basic sanity tests around Java API contracts
+    @Test
+    public void testApiContracts() {
+        assertThrows(IllegalArgumentException.class, () -> Linker.Option.captureCallState("Does not exist"));
+        var duplicateOpt = Linker.Option.captureCallState("errno", "errno"); // duplicates
+        var noDuplicateOpt = Linker.Option.captureCallState("errno");
+        assertEquals(duplicateOpt, noDuplicateOpt, "auto deduplication");
+        var display = duplicateOpt.toString();
+        assertTrue(display.contains("errno"), "toString should contain state name 'errno': " + display);
+    }
+
     private record SaveValuesCase(String nativeTarget, FunctionDescriptor nativeDesc, String threadLocalName,
                                   Consumer<Object> resultCheck, boolean critical) {}