How to mimic `tableswitch` using `MethodHandle`?

Question

Context: I've been benchmarking the difference between using invokedynamic and manually generating bytecode (this is in the context of deciding whether a compiler targeting the JVM should emit more verbose "traditional" bytecode or just an invokedynamic call with a clever bootstrap method). In doing this, it has been pretty straightforward to map bytecode into MethodHandles combinators that are at least as fast, with the exception of tableswitch.

Question: Is there a trick to mimic tableswitch using MethodHandle? I tried mimicking it with a jump table: using a constant MethodHandle[], indexing into that with arrayElementGetter, then calling the found handle with MethodHandles.invoker. However, that ended up being around 50% slower than the original bytecode when I ran it through JMH.

Here's the code for producing the method handle:

private static MethodHandle makeProductElement(Class<?> receiverClass, List<MethodHandle> getters) {
    MethodHandle[] boxedGetters = getters
        .stream()
        .map(getter -> getter.asType(getter.type().changeReturnType(java.lang.Object.class)))
        .toArray(MethodHandle[]::new);

    MethodHandle getGetter = MethodHandles      // (I)H
        .arrayElementGetter(MethodHandle[].class)
        .bindTo(boxedGetters);
    MethodHandle invokeGetter = MethodHandles.permuteArguments( // (RH)O
        MethodHandles.invoker(MethodType.methodType(java.lang.Object.class, receiverClass)),
        MethodType.methodType(java.lang.Object.class, receiverClass, MethodHandle.class),
        1,
        0
    );

    return MethodHandles.filterArguments(invokeGetter, 1, getGetter);
}

Here's the initial bytecode (which I'm trying to replace with one invokedynamic call)

  public java.lang.Object productElement(int);
    descriptor: (I)Ljava/lang/Object;
    flags: (0x0001) ACC_PUBLIC
    Code:
      stack=3, locals=3, args_size=2
         0: iload_1
         1: istore_2
         2: iload_2
         3: tableswitch   { // 0 to 2
                       0: 28
                       1: 38
                       2: 45
                 default: 55
            }
        28: aload_0
        29: invokevirtual #62                 // Method i:()I
        32: invokestatic  #81                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
        35: goto          67
        38: aload_0
        39: invokevirtual #65                 // Method s:()Ljava/lang/String;
        42: goto          67
        45: aload_0
        46: invokevirtual #68                 // Method l:()J
        49: invokestatic  #85                 // Method java/lang/Long.valueOf:(J)Ljava/lang/Long;
        52: goto          67
        55: new           #87                 // class java/lang/IndexOutOfBoundsException
        58: dup
        59: iload_1
        60: invokestatic  #93                 // Method java/lang/Integer.toString:(I)Ljava/lang/String;
        63: invokespecial #96                 // Method java/lang/IndexOutOfBoundsException."<init>":(Ljava/lang/String;)V
        66: athrow
        67: areturn

Answer 1

The good thing about invokedynamic is that it allows to postpone the decision, how to implement the operation to the actual runtime. This is the trick behind LambdaMetafactory or StringConcatFactory which may return composed method handles, like in your example code, or dynamically generated code, at the particular implementation’s discretion.

There’s even a combined approach possible, generate classes which you compose to an operation, e.g. settling on the already existing LambdaMetafactory:

private static MethodHandle makeProductElement(
    MethodHandles.Lookup lookup, Class<?> receiverClass, List<MethodHandle> getters)
    throws Throwable {

    Function[] boxedGetters = new Function[getters.size()];
    MethodType factory = MethodType.methodType(Function.class);
    for(int ix = 0; ix < boxedGetters.length; ix++) {
        MethodHandle mh = getters.get(ix);
        MethodType actual = mh.type().wrap(), generic = actual.erase();
        boxedGetters[ix] = (Function)LambdaMetafactory.metafactory(lookup,
            "apply", factory, generic, mh, actual).getTarget().invokeExact();
    }

    Object switcher = new Object() {
        final Object get(Object receiver, int index) {
            return boxedGetters[index].apply(receiver);
        }
    };
    return lookup.bind(switcher, "get",
            MethodType.methodType(Object.class, Object.class, int.class))
        .asType(MethodType.methodType(Object.class, receiverClass, int.class));
}

This uses the LambdaMetafactory to generate a Function instance for each getter, similar to equivalent method references. Then, an actual class calling the right Function’s apply method is instantiated and a method handle to its get method returned.

This is a similar composition as your method handles, but with the reference implementation, no handles but fully materialized classes are used. I’d expect the composed handles and this approach to converge to the same performance for a very large number of invocations, but the materialized classes having a headstart for a medium number of invocations.

I added a first parameter MethodHandles.Lookup lookup which should be the lookup object received by the bootstrap method for the invokedynamic instruction. If used that way, the generated functions can access all methods the same way as the code containing the invokedynamic instruction, including private methods of that class.

Alternatively, you can generate a class containing a real switch instruction yourself. Using the ASM library, it may look like:

private static MethodHandle makeProductElement(
    MethodHandles.Lookup lookup, Class<?> receiverClass, List<MethodHandle> getters)
    throws ReflectiveOperationException {

    ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
    cw.visit(V1_8, ACC_INTERFACE|ACC_ABSTRACT,
        lookup.lookupClass().getName().replace('.', '/')+"$Switch", null,
        "java/lang/Object", null);
    MethodType type = MethodType.methodType(Object.class, receiverClass, int.class);
    MethodVisitor mv = cw.visitMethod(ACC_STATIC|ACC_PUBLIC, "get",
        type.toMethodDescriptorString(), null, null);
    mv.visitCode();

    Label defaultCase = new Label();
    Label[] cases = new Label[getters.size()];
    for(int ix = 0; ix < cases.length; ix++) cases[ix] = new Label();

    mv.visitVarInsn(ALOAD, 0);
    mv.visitVarInsn(ILOAD, 1);
    mv.visitTableSwitchInsn(0, cases.length - 1, defaultCase, cases);

    String owner = receiverClass.getName().replace('.', '/');

    for(int ix = 0; ix < cases.length; ix++) {
        mv.visitLabel(cases[ix]);
        MethodHandle mh = getters.get(ix);
        mv.visitMethodInsn(INVOKEVIRTUAL, owner, lookup.revealDirect(mh).getName(),
            mh.type().dropParameterTypes(0, 1).toMethodDescriptorString(), false);
        if(mh.type().returnType().isPrimitive()) {
            Class<?> boxed = mh.type().wrap().returnType();
            MethodType box = MethodType.methodType(boxed, mh.type().returnType());
            mv.visitMethodInsn(INVOKESTATIC, boxed.getName().replace('.', '/'),
                "valueOf", box.toMethodDescriptorString(), false);
        }
        mv.visitInsn(ARETURN);
    }
    mv.visitLabel(defaultCase);
    mv.visitTypeInsn(NEW, "java/lang/IndexOutOfBoundsException");
    mv.visitInsn(DUP);
    mv.visitVarInsn(ILOAD, 1);
    mv.visitMethodInsn(INVOKESTATIC, "java/lang/String",
        "valueOf", "(I)Ljava/lang/String;", false);
    mv.visitMethodInsn(INVOKESPECIAL, "java/lang/IndexOutOfBoundsException",
        "<init>", "(Ljava/lang/String;)V", false);
    mv.visitInsn(ATHROW);
    mv.visitMaxs(-1, -1);
    mv.visitEnd();
    cw.visitEnd();

    lookup = lookup.defineHiddenClass(
        cw.toByteArray(), true, MethodHandles.Lookup.ClassOption.NESTMATE);
    return lookup.findStatic(lookup.lookupClass(), "get", type);
}

This generates a new class with a static method containing the tableswitch instruction and the invocations (as well as the boxing conversions we now have to do ourselves). Also, it has the necessary code to create and throw an exception for out-of-bounds values. After generating the class, it returns a handle to that static method.

Answer 2

I don't know of your timeline. But it is likely there will be a MethodHandles.tableSwitch operation in Java 17. It is currently being integrated via https://github.com/openjdk/jdk/pull/3401/

Some more discussion about it here: https://mail.openjdk.java.net/pipermail/core-libs-dev/2021-April/076105.html

Answer 3

The things is, tableswitch isn't always compiled to a jump table. For a small number of labels, like in your example, it's likely to act as a binary search. Thus using a tree of regular "if-then" MethodHandles will be the closest equivalent.