Avoiding megamorphic callsites in Function.andThen

Question

Looking at Function.andThen:

    default <V> Function<T, V> andThen(Function<? super R, ? extends V> after) {
        Objects.requireNonNull(after);
        return (T t) -> after.apply(apply(t));
    }

As this function offers great value it's used extensively which turns its returned function into a pretty bad megamorphic callsite (in the current OpenJDK11 implementation). Even when resorting to manually instantiating a lambda via LamdaMetafactory (see the accepted answer here) all the generated lambdas share the same body and therefor the same byte-code-index which is (as I am aware of) used to keep the type-profile data of callsites which makes it megamorphic again.

Is there any reasonable way to achieve a non-megamorphic andThen without:

• compile-time code-generation via build-system plugins
• runtime code-generation that is not supported on J9, OpenJDK Hotspot and Graal

Answer 1

Of course, the simplest and least resource consuming solution is to wait until JDK-8015416 gets fixed.

As a work-around, we may produce new distinct classes.

If we limit the approach to certain well-known operations like combining two Function instances via andThen, spinning a new class for each request is quiet easy to achieve, with standard APIs and no need to perform bytecode magic:

public final class FunctionCombinator<T,U,R> implements Function<T,R> {
    final Function<? super T,? extends U> first;
    final Function<? super U,? extends R> second;

    public FunctionCombinator(
        Function<? super T,? extends U> f1, Function<? super U,? extends R> f2) {
        first = f1;
        second = f2;
    }

    @Override
    public R apply(T t) {
        return second.apply(first.apply(t));
    }

    @Override
    public <V> Function<T, V> andThen(Function<? super R, ? extends V> after) {
        return newCombinator(this, after);
    }

    @Override
    public <V> Function<V, R> compose(Function<? super V, ? extends T> before) {
        return newCombinator(before, this);
    }

    @SuppressWarnings("unchecked")
    public static <A,B,C> Function<A,C> newCombinator(
        Function<? super A,? extends B> f1, Function<? super B,? extends C> f2) {

        Objects.requireNonNull(f1);
        Objects.requireNonNull(f2);

        URL u = FunctionCombinator.class.getProtectionDomain()
            .getCodeSource().getLocation();
        try(URLClassLoader cl = new URLClassLoader(new URL[] { u }, null)) {
            return cl.loadClass(FunctionCombinator.class.getName())
                .asSubclass(Function.class)
                .getConstructor(Function.class, Function.class)
                .newInstance(f1, f2);
        }
        catch(IOException | ReflectiveOperationException ex) {
            throw new IllegalStateException(ex);
        }
    }
}

This code simply re-reads the class definition of FunctionCombinator in a new class loader having the bootstrap loader as its parent, which prevents it from resolving to the already existing class.

The hidden costs associated with such a class loader are implementation dependent. Note that this code closes the loader after loading the single class to cut down the allocated resources. If the particular JVM supports class unloading, the loader of a function combinator may get garbage collected when the function is not used anymore.

Of course, environments with static compilation, not supporting the addition of new classes won’t support this. For those environments, you have to rely on the capabilities of static code analysis anyway. You may make this class generation optional, to avoid failures in these environments:

static final boolean CREATE_NEW_CLASSES = Boolean.getBoolean("generateNewClassForAndThen");

@SuppressWarnings("unchecked")
public static <A,B,C> Function<A,C> newCombinator(
    Function<? super A,? extends B> f1, Function<? super B,? extends C> f2) {

    Objects.requireNonNull(f1);
    Objects.requireNonNull(f2);

    if(!CREATE_NEW_CLASSES) return new FunctionCombinator<>(f1, f2);

    URL u = FunctionCombinator.class.getProtectionDomain().getCodeSource().getLocation();
    try(URLClassLoader cl = new URLClassLoader(new URL[] { u }, null)) {
        return cl.loadClass(FunctionCombinator.class.getName()).asSubclass(Function.class)
            .getConstructor(Function.class, Function.class)
            .newInstance(f1, f2);
    }
    catch(IOException | ReflectiveOperationException ex) {
        throw new IllegalStateException(ex);
    }
}

Then, you need to specify -DgenerateNewClassForAndThen=true on the command line to activate this feature.

If you want to consume less resources, you’d have to resort to Unsafe…