8332528: Generate code in SwitchBootstraps.generateTypeSwitch that require fewer adaptations

Reviewed-by: liach, jlahoda

Author: cl4es

make/jdk/src/classes/build/tools/classlist/HelloClasslist.java

@@ -151,6 +151,17 @@ public static void main(String ... args) throws Throwable {
 
         LOGGER.log(Level.FINE, "New Date: " + newDate + " - old: " + oldDate);
 
+        // Pull SwitchBootstraps and associated classes into the classlist
+        record A(int a) { }
+        record B(int b) { }
+        Object o = new A(4711);
+        int value = switch (o) {
+            case A a -> a.a;
+            case B b -> b.b;
+            default -> 17;
+        };
+        LOGGER.log(Level.FINE, "Value: " + value);
+
         // The Striped64$Cell is loaded rarely only when there's a contention among
         // multiple threads performing LongAdder.increment(). This results in
         // an inconsistency in the classlist between builds (see JDK-8295951).

src/java.base/share/classes/java/lang/runtime/SwitchBootstraps.java

@@ -26,6 +26,7 @@
 package java.lang.runtime;
 
 import java.lang.Enum.EnumDesc;
+import java.lang.classfile.ClassBuilder;
 import java.lang.classfile.CodeBuilder;
 import java.lang.constant.ClassDesc;
 import java.lang.constant.ConstantDesc;
@@ -48,6 +49,8 @@
 import java.lang.classfile.ClassFile;
 import java.lang.classfile.Label;
 import java.lang.classfile.instruction.SwitchCase;
+
+import jdk.internal.constant.ReferenceClassDescImpl;
 import jdk.internal.misc.PreviewFeatures;
 import jdk.internal.vm.annotation.Stable;
 
@@ -74,28 +77,38 @@ private SwitchBootstraps() {}
     private static final MethodHandles.Lookup LOOKUP = MethodHandles.lookup();
     private static final boolean previewEnabled = PreviewFeatures.isEnabled();
 
-    private static final MethodHandle NULL_CHECK;
-    private static final MethodHandle IS_ZERO;
-    private static final MethodHandle CHECK_INDEX;
-    private static final MethodHandle MAPPED_ENUM_LOOKUP;
+
+    private static final MethodType TYPES_SWITCH_TYPE = MethodType.methodType(int.class,
+            Object.class,
+            int.class,
+            BiPredicate.class,
+            List.class);
 
     private static final MethodTypeDesc TYPES_SWITCH_DESCRIPTOR =
             MethodTypeDesc.ofDescriptor("(Ljava/lang/Object;ILjava/util/function/BiPredicate;Ljava/util/List;)I");
-
-    static {
-        try {
-            NULL_CHECK = LOOKUP.findStatic(Objects.class, "isNull",
-                                           MethodType.methodType(boolean.class, Object.class));
-            IS_ZERO = LOOKUP.findStatic(SwitchBootstraps.class, "isZero",
-                                           MethodType.methodType(boolean.class, int.class));
-            CHECK_INDEX = LOOKUP.findStatic(Objects.class, "checkIndex",
-                                           MethodType.methodType(int.class, int.class, int.class));
-            MAPPED_ENUM_LOOKUP = LOOKUP.findStatic(SwitchBootstraps.class, "mappedEnumLookup",
-                                                   MethodType.methodType(int.class, Enum.class, MethodHandles.Lookup.class,
-                                                                         Class.class, EnumDesc[].class, EnumMap.class));
-        }
-        catch (ReflectiveOperationException e) {
-            throw new ExceptionInInitializerError(e);
+    private static final MethodTypeDesc CHECK_INDEX_DESCRIPTOR =
+            MethodTypeDesc.ofDescriptor("(II)I");
+
+    private static final ClassDesc CD_Objects = ReferenceClassDescImpl.ofValidated("Ljava/util/Objects;");
+
+    private static class StaticHolders {
+        private static final MethodHandle NULL_CHECK;
+        private static final MethodHandle IS_ZERO;
+        private static final MethodHandle MAPPED_ENUM_LOOKUP;
+
+        static {
+            try {
+                NULL_CHECK = LOOKUP.findStatic(Objects.class, "isNull",
+                                               MethodType.methodType(boolean.class, Object.class));
+                IS_ZERO = LOOKUP.findStatic(SwitchBootstraps.class, "isZero",
+                                               MethodType.methodType(boolean.class, int.class));
+                MAPPED_ENUM_LOOKUP = LOOKUP.findStatic(SwitchBootstraps.class, "mappedEnumLookup",
+                                                       MethodType.methodType(int.class, Enum.class, MethodHandles.Lookup.class,
+                                                                             Class.class, EnumDesc[].class, EnumMap.class));
+            }
+            catch (ReflectiveOperationException e) {
+                throw new ExceptionInInitializerError(e);
+            }
         }
     }
 
@@ -163,14 +176,13 @@ public static CallSite typeSwitch(MethodHandles.Lookup lookup,
             || (!invocationType.returnType().equals(int.class))
             || !invocationType.parameterType(1).equals(int.class))
             throw new IllegalArgumentException("Illegal invocation type " + invocationType);
-        requireNonNull(labels);
 
-        Stream.of(labels).forEach(l -> verifyLabel(l, selectorType));
+        for (Object l : labels) { // implicit null-check
+            verifyLabel(l, selectorType);
+        }
 
         MethodHandle target = generateTypeSwitch(lookup, selectorType, labels);
 
-        target = withIndexCheck(target, labels.length);
-
         return new ConstantCallSite(target);
     }
 
@@ -282,18 +294,17 @@ public static CallSite enumSwitch(MethodHandles.Lookup lookup,
             //else if (idx == 0) return mappingArray[selector.ordinal()]; //mapping array created lazily
             //else return "typeSwitch(labels)"
             MethodHandle body =
-                    MethodHandles.guardWithTest(MethodHandles.dropArguments(NULL_CHECK, 0, int.class),
+                    MethodHandles.guardWithTest(MethodHandles.dropArguments(StaticHolders.NULL_CHECK, 0, int.class),
                                                 MethodHandles.dropArguments(MethodHandles.constant(int.class, -1), 0, int.class, Object.class),
-                                                MethodHandles.guardWithTest(MethodHandles.dropArguments(IS_ZERO, 1, Object.class),
+                                                MethodHandles.guardWithTest(MethodHandles.dropArguments(StaticHolders.IS_ZERO, 1, Object.class),
                                                                             generateTypeSwitch(lookup, invocationType.parameterType(0), labels),
-                                                                            MethodHandles.insertArguments(MAPPED_ENUM_LOOKUP, 1, lookup, enumClass, labels, new EnumMap())));
+                                                                            MethodHandles.insertArguments(StaticHolders.MAPPED_ENUM_LOOKUP, 1, lookup, enumClass, labels, new EnumMap())));
             target = MethodHandles.permuteArguments(body, MethodType.methodType(int.class, Object.class, int.class), 1, 0);
         } else {
             target = generateTypeSwitch(lookup, invocationType.parameterType(0), labels);
         }
 
         target = target.asType(invocationType);
-        target = withIndexCheck(target, labels.length);
 
         return new ConstantCallSite(target);
     }
@@ -339,12 +350,6 @@ private static <T extends Enum<T>> int mappedEnumLookup(T value, MethodHandles.L
         return enumMap.map[value.ordinal()];
     }
 
-    private static MethodHandle withIndexCheck(MethodHandle target, int labelsCount) {
-        MethodHandle checkIndex = MethodHandles.insertArguments(CHECK_INDEX, 1, labelsCount + 1);
-
-        return MethodHandles.filterArguments(target, 1, checkIndex);
-    }
-
     private static final class ResolvedEnumLabels implements BiPredicate<Integer, Object> {
 
         private final MethodHandles.Lookup lookup;
@@ -407,6 +412,11 @@ private static Consumer<CodeBuilder> generateTypeSwitchSkeleton(Class<?> selecto
         int EXTRA_CLASS_LABELS  = 3;
 
         return cb -> {
+            // Objects.checkIndex(RESTART_IDX, labelConstants + 1)
+            cb.iload(RESTART_IDX);
+            cb.loadConstant(labelConstants.length + 1);
+            cb.invokestatic(CD_Objects, "checkIndex", CHECK_INDEX_DESCRIPTOR);
+            cb.pop();
             cb.aload(SELECTOR_OBJ);
             Label nonNullLabel = cb.newLabel();
             cb.if_nonnull(nonNullLabel);
@@ -621,7 +631,7 @@ private static MethodHandle generateTypeSwitch(MethodHandles.Lookup caller, Clas
         List<EnumDesc<?>> enumDescs = new ArrayList<>();
         List<Class<?>> extraClassLabels = new ArrayList<>();
 
-        byte[] classBytes = ClassFile.of().build(ClassDesc.of(typeSwitchClassName(caller.lookupClass())),
+        byte[] classBytes = ClassFile.of().build(ReferenceClassDescImpl.ofValidatedBinaryName(typeSwitchClassName(caller.lookupClass())),
                 clb -> {
                     clb.withFlags(AccessFlag.FINAL, AccessFlag.SUPER, AccessFlag.SYNTHETIC)
                        .withMethodBody("typeSwitch",
@@ -636,12 +646,8 @@ private static MethodHandle generateTypeSwitch(MethodHandles.Lookup caller, Clas
             lookup = caller.defineHiddenClass(classBytes, true, NESTMATE, STRONG);
             MethodHandle typeSwitch = lookup.findStatic(lookup.lookupClass(),
                                                         "typeSwitch",
-                                                        MethodType.methodType(int.class,
-                                                                              Object.class,
-                                                                              int.class,
-                                                                              BiPredicate.class,
-                                                                              List.class));
-            typeSwitch = MethodHandles.insertArguments(typeSwitch, 2, new ResolvedEnumLabels(caller, enumDescs.toArray(EnumDesc[]::new)),
+                                                        TYPES_SWITCH_TYPE);
+            typeSwitch = MethodHandles.insertArguments(typeSwitch, 2, new ResolvedEnumLabels(caller, enumDescs.toArray(new EnumDesc<?>[0])),
                                                        List.copyOf(extraClassLabels));
             typeSwitch = MethodHandles.explicitCastArguments(typeSwitch,
                                                              MethodType.methodType(int.class,

src/java.base/share/classes/jdk/internal/constant/ReferenceClassDescImpl.java

@@ -69,6 +69,17 @@ public static ReferenceClassDescImpl ofValidated(String descriptor) {
         return new ReferenceClassDescImpl(descriptor);
     }
 
+    /**
+     * Creates a {@linkplain ClassDesc} from a pre-validated descriptor string
+     * for a class or interface type or an array type.
+     *
+     * @param descriptor a field descriptor string for a class or interface type
+     * @jvms 4.3.2 Field Descriptors
+     */
+    public static ClassDesc ofValidatedBinaryName(String typeSwitchClassName) {
+        return ofValidated("L" + binaryToInternal(typeSwitchClassName) + ";");
+    }
+
     @Override
     public String descriptorString() {
         return descriptor;

test/micro/org/openjdk/bench/java/lang/runtime/SwitchSanity.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright (c) 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
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package org.openjdk.bench.java.lang.runtime;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.concurrent.TimeUnit;
+
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@State(Scope.Thread)
+@Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 2, timeUnit = TimeUnit.SECONDS)
+@Fork(3)
+public class SwitchSanity {
+
+    public record A(int a) { }
+    public record B(int b) { }
+    public record C(int c) { }
+
+    public Object[] inputs = new Object[10];
+    @Setup
+    public void setup() {
+        for (int i = 0; i < 10; i++) {
+            if (i % 2 == 0) {
+                inputs[i] = new A(i + 17);
+            } else if (i % 3 == 0) {
+                inputs[i] = new B(i + 4711);
+            } else {
+                inputs[i] = new C(i + 174711);
+            }
+        }
+    }
+
+    @Benchmark
+    public int switchSum() {
+        int sum = 0;
+        for (Object o : inputs) {
+            sum += switch (o) {
+                case A a -> a.a;
+                case B b -> b.b;
+                case C c -> c.c;
+                default -> 17;
+            };
+        }
+        return sum;
+    }
+
+    public static void main(String[] args) {
+        SwitchSanity s = new SwitchSanity();
+        s.setup();
+        System.out.println(s.switchSum());
+    }
+}