classlib: support SwitchBootstraps (#764)

Java 21 supports switch pattern matching (as well as few earlier Java version as an experimental feature). Javac produces code with INVOKEDYNAMIC that relies on bootstrap methods from `java.lang.runtime.SwitchBootstraps`. This commit is intended to support for new Java language feature

classlib/src/main/java/org/teavm/classlib/impl/JCLPlugin.java

@@ -130,6 +130,22 @@ public class JCLPlugin implements TeaVMPlugin {
                         ValueType.object("java.lang.invoke.CallSite")),
                 stringConcatSubstitutor);
 
+        SwitchBootstrapSubstitutor switchBootstrapSubstitutor = new SwitchBootstrapSubstitutor();
+        host.add(new MethodReference("java.lang.runtime.SwitchBootstraps", "typeSwitch",
+                        ValueType.object("java.lang.invoke.MethodHandles$Lookup"),
+                        ValueType.object("java.lang.String"),
+                        ValueType.object("java.lang.invoke.MethodType"),
+                        ValueType.arrayOf(ValueType.object("java.lang.Object")),
+                        ValueType.object("java.lang.invoke.CallSite")),
+                switchBootstrapSubstitutor);
+        host.add(new MethodReference("java.lang.runtime.SwitchBootstraps", "enumSwitch",
+                        ValueType.object("java.lang.invoke.MethodHandles$Lookup"),
+                        ValueType.object("java.lang.String"),
+                        ValueType.object("java.lang.invoke.MethodType"),
+                        ValueType.arrayOf(ValueType.object("java.lang.Object")),
+                        ValueType.object("java.lang.invoke.CallSite")),
+                switchBootstrapSubstitutor);
+
         if (!isBootstrap()) {
             host.add(new ScalaHacks());
             host.add(new KotlinHacks());

classlib/src/main/java/org/teavm/classlib/impl/SwitchBootstrapSubstitutor.java

@@ -0,0 +1,120 @@
+/*
+ *  Copyright 2023 ihromant.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.classlib.impl;
+
+import java.util.List;
+import org.teavm.dependency.BootstrapMethodSubstitutor;
+import org.teavm.dependency.DynamicCallSite;
+import org.teavm.model.BasicBlock;
+import org.teavm.model.RuntimeConstant;
+import org.teavm.model.ValueType;
+import org.teavm.model.emit.PhiEmitter;
+import org.teavm.model.emit.ProgramEmitter;
+import org.teavm.model.emit.ValueEmitter;
+import org.teavm.model.instructions.SwitchInstruction;
+import org.teavm.model.instructions.SwitchTableEntry;
+
+public class SwitchBootstrapSubstitutor implements BootstrapMethodSubstitutor {
+    @Override
+    public ValueEmitter substitute(DynamicCallSite callSite, ProgramEmitter pe) {
+        boolean enumSwitch = callSite.getBootstrapMethod().getName().equals("enumSwitch");
+        List<RuntimeConstant> labels = callSite.getBootstrapArguments();
+        ValueEmitter target = callSite.getArguments().get(0);
+        ValueEmitter restartIdx = callSite.getArguments().get(1);
+        BasicBlock joint = pe.prepareBlock();
+        PhiEmitter result = pe.phi(ValueType.INTEGER, joint);
+        pe.when(() -> target.isNull()).thenDo(() -> {
+            pe.constant(-1).propagateTo(result);
+            pe.jump(joint);
+        });
+
+        var switchInsn = new SwitchInstruction();
+        switchInsn.setCondition(restartIdx.getVariable());
+        pe.addInstruction(switchInsn);
+
+        var block = pe.prepareBlock();
+        pe.enter(block);
+        ValueType.Object enumType = enumSwitch ? labels.stream()
+                .filter(l -> l.getKind() == RuntimeConstant.TYPE)
+                .findAny().map(vt -> (ValueType.Object) vt.getValueType()).orElseThrow() : null;
+        if (enumType != null) {
+            pe.initClass(enumType.getClassName());
+        }
+        for (var i = 0; i < labels.size(); ++i) {
+            var entry = new SwitchTableEntry();
+            entry.setCondition(i);
+            entry.setTarget(block);
+            switchInsn.getEntries().add(entry);
+
+            var label = labels.get(i);
+            emitFragment(target, i, label, pe, result, joint, enumType);
+
+            block = pe.prepareBlock();
+            pe.jump(block);
+            pe.enter(block);
+        }
+
+        switchInsn.setDefaultTarget(block);
+
+        pe.constant(callSite.getBootstrapArguments().size()).propagateTo(result);
+        pe.jump(joint);
+        pe.enter(joint);
+        return result.getValue();
+    }
+
+    private void emitFragment(ValueEmitter target, int idx, RuntimeConstant label, ProgramEmitter pe,
+            PhiEmitter result, BasicBlock exit, ValueType.Object enumType) {
+        switch (label.getKind()) {
+            case RuntimeConstant.TYPE:
+                ValueType type = label.getValueType();
+                pe.when(() -> target.instanceOf(type).isTrue())
+                        .thenDo(() -> {
+                            pe.constant(idx).propagateTo(result);
+                            pe.jump(exit);
+                        });
+                break;
+            case RuntimeConstant.INT:
+                int val = label.getInt();
+                pe.when(() -> target.instanceOf(ValueType.object("java.lang.Number")).isTrue()
+                        .and(() -> target.cast(Number.class)
+                                .invokeVirtual("intValue", int.class).isSame(pe.constant(val))))
+                        .thenDo(() -> {
+                            pe.constant(idx).propagateTo(result);
+                            pe.jump(exit);
+                        });
+                pe.when(() -> target.instanceOf(ValueType.object("java.lang.Character")).isTrue()
+                        .and(() -> target.cast(Character.class)
+                                .invokeSpecial("charValue", char.class).isSame(pe.constant(val))))
+                        .thenDo(() -> {
+                                    pe.constant(idx).propagateTo(result);
+                                    pe.jump(exit);
+                                });
+                break;
+            case RuntimeConstant.STRING:
+                String str = label.getString();
+                pe.when(enumType != null
+                                ? () -> pe.getField(enumType.getClassName(), str, enumType).isSame(target)
+                                : () -> pe.constant(str).isEqualTo(target))
+                        .thenDo(() -> {
+                            pe.constant(idx).propagateTo(result);
+                            pe.jump(exit);
+                        });
+                break;
+            default:
+                throw new IllegalArgumentException("Unsupported constant type: " + label.getKind());
+        }
+    }
+}

tests/src/test/java/org/teavm/vm/SwitchTest.java

@@ -0,0 +1,203 @@
+/*
+ *  Copyright 2023 ihromant.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.vm;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+import org.junit.runner.RunWith;
+import org.teavm.junit.TeaVMTestRunner;
+import org.testng.annotations.Test;
+
+@RunWith(TeaVMTestRunner.class)
+public class SwitchTest {
+    private static int switchWithLogic(Object o) {
+        return switch (o) {
+            case null -> -1;
+            case A a when (a.af & 31) == 5 -> 0;
+            case A a -> 1;
+            case B b -> {
+                b.bbf = 21;
+                yield b.f.length();
+            }
+            case C c -> c.cf.af;
+            case D(byte c, short d) when ((int) d & 31) == 31 -> c;
+            case D(byte c, short d) -> d;
+            case TestEnum te -> te.ordinal();
+            default -> throw new IllegalArgumentException();
+        };
+    }
+
+    @Test
+    public void genericSwitch() {
+        assertEquals(-1, switchWithLogic(null));
+        A a = new A();
+        assertEquals(1, switchWithLogic(a));
+        a.af = 5;
+        assertEquals(0, switchWithLogic(a));
+        B b = new B();
+        b.f = "abc";
+        assertEquals(3, switchWithLogic(b));
+        assertEquals(21, b.bbf);
+        C c = new C();
+        c.cf = a;
+        assertEquals(5, switchWithLogic(c));
+        assertEquals(Byte.MIN_VALUE, switchWithLogic(new D(Byte.MIN_VALUE, Short.MAX_VALUE)));
+        assertEquals(Short.MIN_VALUE, switchWithLogic(new D(Byte.MIN_VALUE, Short.MIN_VALUE)));
+        assertEquals(4, switchWithLogic(TestEnum.E));
+        try {
+            switchWithLogic(new Object());
+            fail();
+        } catch (IllegalArgumentException e) {
+            // ok
+        }
+    }
+
+    private static int enumSwitchWithLogic(TestEnum o) {
+        return switch (o) {
+            case A, B -> 1;
+            case TestEnum e when e.ordinal() % 3 == 0 -> 3;
+            case C, D, E, F -> 2;
+        };
+    }
+
+    @Test
+    public void enumSwitch() {
+        assertEquals(1, enumSwitchWithLogic(TestEnum.A));
+        assertEquals(2, enumSwitchWithLogic(TestEnum.C));
+        assertEquals(3, enumSwitchWithLogic(TestEnum.D));
+        assertEquals(2, enumSwitchWithLogic(TestEnum.F));
+    }
+
+    private static int integerSwitchWithLogic(Integer o) {
+        return switch (o) {
+            case 23 -> 1;
+            case Integer i when i < 10 -> 2;
+            case 42 -> 3;
+            default -> 4;
+        };
+    }
+
+    @Test
+    public void integerSwitch() {
+        assertEquals(1, integerSwitchWithLogic(23));
+        assertEquals(3, integerSwitchWithLogic(42));
+        assertEquals(4, integerSwitchWithLogic(11));
+        assertEquals(2, integerSwitchWithLogic(5));
+    }
+
+    private static int characterSwitchWithLogic(Character c) {
+        return switch (c) {
+            case Character ch when ch >= 'a' && ch <= 'z' -> 5;
+            case 'R' -> 1;
+            case 'T' -> 2;
+            default -> throw new IllegalArgumentException();
+        };
+    }
+
+    @Test
+    public void characterSwitch() {
+        assertEquals(5, characterSwitchWithLogic('a'));
+        assertEquals(1, characterSwitchWithLogic('R'));
+        assertEquals(2, characterSwitchWithLogic('T'));
+        try {
+            characterSwitchWithLogic('A');
+            fail();
+        } catch (IllegalArgumentException e) {
+            // ok
+        }
+    }
+
+    private static int stringSwitchWithLogic(String s) {
+        return switch (s) {
+            case String str when str.length() < 3 -> 0;
+            case "abc" -> 1;
+            default -> 2;
+        };
+    }
+
+    @Test
+    public void stringSwitch() {
+        assertEquals(0, stringSwitchWithLogic(""));
+        assertEquals(1, stringSwitchWithLogic("abc"));
+        assertEquals(2, stringSwitchWithLogic("bcd"));
+    }
+
+    private static int switchWithHierarchy(Object o) {
+        return switch (o) {
+            case Superclass s when s.x == 23 -> 1;
+            case SubclassA a -> 2;
+            case Superclass s when s.x == 24 -> 3;
+            case SubclassB b -> 4;
+            default -> 5;
+        };
+    }
+
+    @Test
+    public void hierarchySwitch() {
+        assertEquals(1, switchWithHierarchy(new SubclassA(23)));
+        assertEquals(2, switchWithHierarchy(new SubclassA(24)));
+        assertEquals(2, switchWithHierarchy(new SubclassA(1)));
+        assertEquals(1, switchWithHierarchy(new SubclassB(23)));
+        assertEquals(3, switchWithHierarchy(new SubclassB(24)));
+        assertEquals(4, switchWithHierarchy(new SubclassB(1)));
+        assertEquals(5, switchWithHierarchy("foo"));
+        assertEquals(5, switchWithHierarchy(new Superclass(1)));
+        assertEquals(1, switchWithHierarchy(new Superclass(23)));
+    }
+
+    private static class A {
+        private int af;
+    }
+
+    private static class B {
+        private String f;
+        private int bbf;
+    }
+
+    private static class C {
+        private A cf;
+        private long ccf;
+        private boolean cccf;
+    }
+
+    private record D(byte a, short b) {
+
+    }
+
+    private enum TestEnum {
+        A, B, C, D, E, F
+    }
+
+    private static class Superclass {
+        final int x;
+
+        public Superclass(int x) {
+            this.x = x;
+        }
+    }
+
+    private static class SubclassA extends Superclass {
+        public SubclassA(int x) {
+            super(x);
+        }
+    }
+
+    private static class SubclassB extends Superclass {
+        public SubclassB(int x) {
+            super(x);
+        }
+    }
+}