8069591: Customize LambdaForms which are invoked using MH.invoke/invokeExact

Reviewed-by: jrose, plevart, forax

src/share/classes/java/lang/invoke/DirectMethodHandle.java

@@ -31,7 +31,6 @@ import java.util.Arrays;
 import sun.invoke.util.VerifyAccess;
 import static java.lang.invoke.MethodHandleNatives.Constants.*;
 import static java.lang.invoke.LambdaForm.*;
-import static java.lang.invoke.LambdaForm.BasicType.*;
 import static java.lang.invoke.MethodTypeForm.*;
 import static java.lang.invoke.MethodHandleStatics.*;
 import java.lang.ref.WeakReference;
@@ -693,4 +692,10 @@ class DirectMethodHandle extends MethodHandle {
             }
         }
     }
+
+    @Override
+    void customize() {
+        assert(form.customized == null);
+        // No need to customize DMHs.
+    }
 }

src/share/classes/java/lang/invoke/InvokerBytecodeGenerator.java

@@ -56,9 +56,11 @@ class InvokerBytecodeGenerator {
     private static final String OBJ     = "java/lang/Object";
     private static final String OBJARY  = "[Ljava/lang/Object;";
 
+    private static final String MH_SIG  = "L" + MH + ";";
     private static final String LF_SIG  = "L" + LF + ";";
     private static final String LFN_SIG = "L" + LFN + ";";
     private static final String LL_SIG  = "(L" + OBJ + ";)L" + OBJ + ";";
+    private static final String LLV_SIG = "(L" + OBJ + ";L" + OBJ + ";)V";
     private static final String CLL_SIG = "(L" + CLS + ";L" + OBJ + ";)L" + OBJ + ";";
 
     /** Name of its super class*/
@@ -616,6 +618,15 @@ class InvokerBytecodeGenerator {
         return g.loadMethod(g.generateCustomizedCodeBytes());
     }
 
+    /** Generates code to check that actual receiver and LambdaForm matches */
+    private boolean checkActualReceiver() {
+        // Expects MethodHandle on the stack and actual receiver MethodHandle in slot #0
+        mv.visitInsn(Opcodes.DUP);
+        mv.visitVarInsn(Opcodes.ALOAD, localsMap[0]);
+        mv.visitMethodInsn(Opcodes.INVOKESTATIC, MHI, "assertSame", LLV_SIG, false);
+        return true;
+    }
+
     /**
      * Generate an invoker method for the passed {@link LambdaForm}.
      */
@@ -635,6 +646,17 @@ class InvokerBytecodeGenerator {
             mv.visitAnnotation("Ljava/lang/invoke/DontInline;", true);
         }
 
+        if (lambdaForm.customized != null) {
+            // Since LambdaForm is customized for a particular MethodHandle, it's safe to substitute
+            // receiver MethodHandle (at slot #0) with an embedded constant and use it instead.
+            // It enables more efficient code generation in some situations, since embedded constants
+            // are compile-time constants for JIT compiler.
+            mv.visitLdcInsn(constantPlaceholder(lambdaForm.customized));
+            mv.visitTypeInsn(Opcodes.CHECKCAST, MH);
+            assert(checkActualReceiver()); // expects MethodHandle on top of the stack
+            mv.visitVarInsn(Opcodes.ASTORE, localsMap[0]);
+        }
+
         // iterate over the form's names, generating bytecode instructions for each
         // start iterating at the first name following the arguments
         Name onStack = null;

src/share/classes/java/lang/invoke/Invokers.java

@@ -247,6 +247,7 @@ class Invokers {
         int nameCursor = OUTARG_LIMIT;
         final int MTYPE_ARG    = customized ? -1 : nameCursor++;  // might be last in-argument
         final int CHECK_TYPE   = nameCursor++;
+        final int CHECK_CUSTOM = (CUSTOMIZE_THRESHOLD >= 0) ? nameCursor++ : -1;
         final int LINKER_CALL  = nameCursor++;
         MethodType invokerFormType = mtype.invokerType();
         if (isLinker) {
@@ -279,6 +280,9 @@ class Invokers {
             // mh.invokeGeneric(a*):R => checkGenericType(mh, TYPEOF(a*:R)).invokeBasic(a*)
             outArgs[0] = names[CHECK_TYPE];
         }
+        if (CHECK_CUSTOM != -1) {
+            names[CHECK_CUSTOM] = new Name(NF_checkCustomized, names[CALL_MH]);
+        }
         names[LINKER_CALL] = new Name(outCallType, outArgs);
         lform = new LambdaForm(debugName, INARG_LIMIT, names);
         if (isLinker)
@@ -386,11 +390,32 @@ class Invokers {
         return ((CallSite)site).getTarget();
     }
 
+    /*non-public*/ static
+    @ForceInline
+    void checkCustomized(Object o) {
+        MethodHandle mh = (MethodHandle)o;
+        if (mh.form.customized == null) {
+            maybeCustomize(mh);
+        }
+    }
+
+    /*non-public*/ static
+    @DontInline
+    void maybeCustomize(MethodHandle mh) {
+        byte count = mh.customizationCount;
+        if (count >= CUSTOMIZE_THRESHOLD) {
+            mh.customize();
+        } else {
+            mh.customizationCount = (byte)(count+1);
+        }
+    }
+
     // Local constant functions:
     private static final NamedFunction
         NF_checkExactType,
         NF_checkGenericType,
-        NF_getCallSiteTarget;
+        NF_getCallSiteTarget,
+        NF_checkCustomized;
     static {
         try {
             NamedFunction nfs[] = {
@@ -399,7 +424,9 @@ class Invokers {
                 NF_checkGenericType = new NamedFunction(Invokers.class
                         .getDeclaredMethod("checkGenericType", Object.class, Object.class)),
                 NF_getCallSiteTarget = new NamedFunction(Invokers.class
-                        .getDeclaredMethod("getCallSiteTarget", Object.class))
+                        .getDeclaredMethod("getCallSiteTarget", Object.class)),
+                NF_checkCustomized = new NamedFunction(Invokers.class
+                        .getDeclaredMethod("checkCustomized", Object.class))
             };
             for (NamedFunction nf : nfs) {
                 // Each nf must be statically invocable or we get tied up in our bootstraps.

src/share/classes/java/lang/invoke/LambdaForm.java

@@ -120,12 +120,14 @@ class LambdaForm {
     final int arity;
     final int result;
     final boolean forceInline;
+    final MethodHandle customized;
     @Stable final Name[] names;
     final String debugName;
     MemberName vmentry;   // low-level behavior, or null if not yet prepared
     private boolean isCompiled;
 
-    volatile Object transformCache;  // managed by LambdaFormEditor
+    // Either a LambdaForm cache (managed by LambdaFormEditor) or a link to uncustomized version (for customized LF)
+    volatile Object transformCache;
 
     public static final int VOID_RESULT = -1, LAST_RESULT = -2;
 
@@ -244,16 +246,17 @@ class LambdaForm {
 
     LambdaForm(String debugName,
                int arity, Name[] names, int result) {
-        this(debugName, arity, names, result, /*forceInline=*/true);
+        this(debugName, arity, names, result, /*forceInline=*/true, /*customized=*/null);
     }
     LambdaForm(String debugName,
-               int arity, Name[] names, int result, boolean forceInline) {
+               int arity, Name[] names, int result, boolean forceInline, MethodHandle customized) {
         assert(namesOK(arity, names));
         this.arity = arity;
         this.result = fixResult(result, names);
         this.names = names.clone();
         this.debugName = fixDebugName(debugName);
         this.forceInline = forceInline;
+        this.customized = customized;
         int maxOutArity = normalize();
         if (maxOutArity > MethodType.MAX_MH_INVOKER_ARITY) {
             // Cannot use LF interpreter on very high arity expressions.
@@ -263,21 +266,21 @@ class LambdaForm {
     }
     LambdaForm(String debugName,
                int arity, Name[] names) {
-        this(debugName, arity, names, LAST_RESULT, /*forceInline=*/true);
+        this(debugName, arity, names, LAST_RESULT, /*forceInline=*/true, /*customized=*/null);
     }
     LambdaForm(String debugName,
                int arity, Name[] names, boolean forceInline) {
-        this(debugName, arity, names, LAST_RESULT, forceInline);
+        this(debugName, arity, names, LAST_RESULT, forceInline, /*customized=*/null);
     }
     LambdaForm(String debugName,
                Name[] formals, Name[] temps, Name result) {
         this(debugName,
-             formals.length, buildNames(formals, temps, result), LAST_RESULT, /*forceInline=*/true);
+             formals.length, buildNames(formals, temps, result), LAST_RESULT, /*forceInline=*/true, /*customized=*/null);
     }
     LambdaForm(String debugName,
                Name[] formals, Name[] temps, Name result, boolean forceInline) {
         this(debugName,
-             formals.length, buildNames(formals, temps, result), LAST_RESULT, forceInline);
+             formals.length, buildNames(formals, temps, result), LAST_RESULT, forceInline, /*customized=*/null);
     }
 
     private static Name[] buildNames(Name[] formals, Name[] temps, Name result) {
@@ -300,6 +303,7 @@ class LambdaForm {
         this.names = buildEmptyNames(arity, sig);
         this.debugName = "LF.zero";
         this.forceInline = true;
+        this.customized = null;
         assert(nameRefsAreLegal());
         assert(isEmpty());
         assert(sig.equals(basicTypeSignature())) : sig + " != " + basicTypeSignature();
@@ -371,6 +375,31 @@ class LambdaForm {
         return true;
     }
 
+    /** Customize LambdaForm for a particular MethodHandle */
+    LambdaForm customize(MethodHandle mh) {
+        LambdaForm customForm = new LambdaForm(debugName, arity, names, result, forceInline, mh);
+        if (COMPILE_THRESHOLD > 0 && isCompiled) {
+            // If shared LambdaForm has been compiled, compile customized version as well.
+            customForm.compileToBytecode();
+        }
+        customForm.transformCache = this; // LambdaFormEditor should always use uncustomized form.
+        return customForm;
+    }
+
+    /** Get uncustomized flavor of the LambdaForm */
+    LambdaForm uncustomize() {
+        if (customized == null) {
+            return this;
+        }
+        assert(transformCache != null); // Customized LambdaForm should always has a link to uncustomized version.
+        LambdaForm uncustomizedForm = (LambdaForm)transformCache;
+        if (COMPILE_THRESHOLD > 0 && isCompiled) {
+            // If customized LambdaForm has been compiled, compile uncustomized version as well.
+            uncustomizedForm.compileToBytecode();
+        }
+        return uncustomizedForm;
+    }
+
     /** Renumber and/or replace params so that they are interned and canonically numbered.
      *  @return maximum argument list length among the names (since we have to pass over them anyway)
      */
@@ -413,8 +442,8 @@ class LambdaForm {
             for (int i = arity; i < names.length; i++) {
                 names[i].internArguments();
             }
-            assert(nameRefsAreLegal());
         }
+        assert(nameRefsAreLegal());
         return maxOutArity;
     }
 

src/share/classes/java/lang/invoke/LambdaFormEditor.java

@@ -51,7 +51,10 @@ class LambdaFormEditor {
     static LambdaFormEditor lambdaFormEditor(LambdaForm lambdaForm) {
         // TO DO:  Consider placing intern logic here, to cut down on duplication.
         // lambdaForm = findPreexistingEquivalent(lambdaForm)
-        return new LambdaFormEditor(lambdaForm);
+
+        // Always use uncustomized version for editing.
+        // It helps caching and customized LambdaForms reuse transformCache field to keep a link to uncustomized version.
+        return new LambdaFormEditor(lambdaForm.uncustomize());
     }
 
     /** A description of a cached transform, possibly associated with the result of the transform.

src/share/classes/java/lang/invoke/MethodHandle.java

@@ -434,6 +434,8 @@ public abstract class MethodHandle {
     // form is not private so that invokers can easily fetch it
     /*private*/ MethodHandle asTypeCache;
     // asTypeCache is not private so that invokers can easily fetch it
+    /*non-public*/ byte customizationCount;
+    // customizationCount should be accessible from invokers
 
     /**
      * Reports the type of this method handle.
@@ -454,9 +456,9 @@ public abstract class MethodHandle {
         type.getClass();  // explicit NPE
         form.getClass();  // explicit NPE
         this.type = type;
-        this.form = form;
+        this.form = form.uncustomize();
 
-        form.prepare();  // TO DO:  Try to delay this step until just before invocation.
+        this.form.prepare();  // TO DO:  Try to delay this step until just before invocation.
     }
 
     /**
@@ -1425,12 +1427,24 @@ assertEquals("[three, thee, tee]", asListFix.invoke((Object)argv).toString());
      */
     /*non-public*/
     void updateForm(LambdaForm newForm) {
+        assert(newForm.customized == null || newForm.customized == this);
         if (form == newForm)  return;
         newForm.prepare();  // as in MethodHandle.<init>
         UNSAFE.putObject(this, FORM_OFFSET, newForm);
         UNSAFE.fullFence();
     }
 
+    /** Craft a LambdaForm customized for this particular MethodHandle */
+    /*non-public*/
+    void customize() {
+        if (form.customized == null) {
+            LambdaForm newForm = form.customize(this);
+            updateForm(newForm);
+        } else {
+            assert(form.customized == this);
+        }
+    }
+
     private static final long FORM_OFFSET;
     static {
         try {

src/share/classes/java/lang/invoke/MethodHandleImpl.java

@@ -1652,4 +1652,13 @@ import static java.lang.invoke.MethodHandles.Lookup.IMPL_LOOKUP;
         assert(elemType.isPrimitive());
         return Lazy.MH_copyAsPrimitiveArray.bindTo(Wrapper.forPrimitiveType(elemType));
     }
+
+    /*non-public*/ static void assertSame(Object mh1, Object mh2) {
+        if (mh1 != mh2) {
+            String msg = String.format("mh1 != mh2: mh1 = %s (form: %s); mh2 = %s (form: %s)",
+                    mh1, ((MethodHandle)mh1).form,
+                    mh2, ((MethodHandle)mh2).form);
+            throw newInternalError(msg);
+        }
+    }
 }

src/share/classes/java/lang/invoke/MethodHandleStatics.java

@@ -49,9 +49,10 @@ import sun.misc.Unsafe;
     static final int DONT_INLINE_THRESHOLD;
     static final int PROFILE_LEVEL;
     static final boolean PROFILE_GWT;
+    static final int CUSTOMIZE_THRESHOLD;
 
     static {
-        final Object[] values = new Object[8];
+        final Object[] values = new Object[9];
         AccessController.doPrivileged(new PrivilegedAction<Void>() {
                 public Void run() {
                     values[0] = Boolean.getBoolean("java.lang.invoke.MethodHandle.DEBUG_NAMES");
@@ -62,6 +63,7 @@ import sun.misc.Unsafe;
                     values[5] = Integer.getInteger("java.lang.invoke.MethodHandle.DONT_INLINE_THRESHOLD", 30);
                     values[6] = Integer.getInteger("java.lang.invoke.MethodHandle.PROFILE_LEVEL", 0);
                     values[7] = Boolean.parseBoolean(System.getProperty("java.lang.invoke.MethodHandle.PROFILE_GWT", "true"));
+                    values[8] = Integer.getInteger("java.lang.invoke.MethodHandle.CUSTOMIZE_THRESHOLD", 127);
                     return null;
                 }
             });
@@ -73,6 +75,11 @@ import sun.misc.Unsafe;
         DONT_INLINE_THRESHOLD     = (Integer) values[5];
         PROFILE_LEVEL             = (Integer) values[6];
         PROFILE_GWT               = (Boolean) values[7];
+        CUSTOMIZE_THRESHOLD       = (Integer) values[8];
+
+        if (CUSTOMIZE_THRESHOLD < -1 || CUSTOMIZE_THRESHOLD > 127) {
+            throw newInternalError("CUSTOMIZE_THRESHOLD should be in [-1...127] range");
+        }
     }
 
     /** Tell if any of the debugging switches are turned on.