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

Fletcher

Does it have to be a “take it a leave it” approach? The bootstrap method could decide to generate bytecode containing a switch instruction or to return a composed method handle. The reference implementation of StringConcatFactory, to name a practical example, is capable of both, generating code ore composing handles and, of course, the compiler generated invokedynamic instruction is not affected by that runtime choice. – Hex

@Hex I hadn't considered that alternative at all, and if possible that would indeed be fine. Do you have a link to the relevant parts of the StringConcatFactory? I can't find any example of constructing a MethodHandle directly from bytecode. – Fletcher

It’s not “directly from bytecode” but rather creating a class from the bytecode and returning a handle to a method of that class. The preferred mechanism is defineHiddenClass which allows the generated class to invoke private methods of the lookup class. This has been made a standard functionality in JDK 15, previous versions used sun.misc.Unsafe for that. – Hex

@Hex thanks for the pointer! AFAICT, defineHiddenClass isn't used at all for the string concat stuff, but I did find one legit (and unsurprising) use in java.lang.invoke.InnerClassLambdaMetafactory. If you want to put together an example of doing this for some tableswitch instruction (even if it doesn't match exactly my question's function) I'll accept that – Fletcher

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.

Hex

Just wanted to add that I tried the second example almost verbatim (all I tweaked was switching the invokevirtual for a simpler getfield) and the performance is (unsurprisingly) the same as the initial boilerplate method. That's success! This is a great escape hatch to be aware of - thanks. – Fletcher

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

Moshemoshell

This is fantastic news, thanks for the link! I was looking for exactly a combinator like this (and indirectly also wondering where to bring up this feature request). My timeline is not cramped, but the codegen still needs to support Java 8 (I might use this if the user explicitly opts into a target of Java 17). I'm guessing this is also something that javac will need if they want to execute on their plan to transform switch compilation into something that uses indy. – Fletcher

Seeing this SO question actually reminded me to get that PR filed :D – Plasticity

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.

Leroy

In my example, I agree nested if-s would be better. That said, I'm trying to construct a generic bootstrap method that should be performant even for classes with many fields. I'm going to try this approach anyways to confirm how it behaves on large/small records. – Fletcher