Optimize runtime representation for callable reference subclasses

Instead of generating overrides for getOwner/getName/getSignature in
each anonymous class representing a callable reference, pass them to the
superclass' constructor and store as fields. This occupies some small
memory but helps to reduce the size of the generated class files, and
will be helpful for adding further runtime information to callable
references, such as information about implicit conversions this
reference has been subject to.

Represent owner as java.lang.Class + boolean instead of
KDeclarationContainer, so that the unnecessary wrapping Class->KClass
wouldn't happen before it's needed, and also to make sure all callable
references remain serializable.

Note that the argument type where the "is declaration container a class"
is passed is int instead of boolean. The plan is to pass the
aforementioned implicit conversion information as bits of this same
integer value.

 #KT-27362 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/ClosureCodegen.java

@@ -70,6 +70,7 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
     protected final Type asmType;
     protected final int visibilityFlag;
     private final boolean shouldHaveBoundReferenceReceiver;
+    private final boolean isOptimizedFunctionReference;
 
     private Method constructor;
     protected Type superClassAsmType;
@@ -119,6 +120,9 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
         assert closure != null : "Closure must be calculated for class: " + classDescriptor;
 
         this.shouldHaveBoundReferenceReceiver = CallableReferenceUtilKt.isForBoundCallableReference(closure);
+        this.isOptimizedFunctionReference =
+                functionReferenceTarget != null &&
+                superClassType.getConstructor().getDeclarationDescriptor() == state.getJvmRuntimeTypes().getFunctionReferenceImpl();
 
         this.asmType = typeMapper.mapClass(classDescriptor);
 
@@ -180,7 +184,7 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
     protected void generateClosureBody() {
         functionCodegen.generateMethod(JvmDeclarationOriginKt.OtherOrigin(element, funDescriptor), funDescriptor, strategy);
 
-        if (functionReferenceTarget != null) {
+        if (functionReferenceTarget != null && !isOptimizedFunctionReference) {
             generateFunctionReferenceMethods(functionReferenceTarget);
         }
 
@@ -285,7 +289,7 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
         );
     }
 
-    protected void generateBridge(
+    private void generateBridge(
             @NotNull Method bridge,
             @NotNull List<KotlinType> bridgeParameterKotlinTypes,
             @Nullable KotlinType bridgeReturnType,
@@ -396,7 +400,8 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
         }
     }
 
-    public static void generateCallableReferenceDeclarationContainer(
+    // Returns false if null was generated.
+    public static boolean generateCallableReferenceDeclarationContainerClass(
             @NotNull InstructionAdapter iv,
             @NotNull CallableDescriptor descriptor,
             @NotNull GenerationState state
@@ -420,15 +425,20 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
         }
         else {
             iv.aconst(null);
-            return;
+            return false;
         }
 
-        boolean isContainerPackage =
-                descriptor instanceof LocalVariableDescriptor
-                ? DescriptorUtils.getParentOfType(descriptor, ClassDescriptor.class) == null
-                : container instanceof PackageFragmentDescriptor;
+        return true;
+    }
 
-        if (isContainerPackage) {
+    public static void generateCallableReferenceDeclarationContainer(
+            @NotNull InstructionAdapter iv,
+            @NotNull CallableDescriptor descriptor,
+            @NotNull GenerationState state
+    ) {
+        if (!generateCallableReferenceDeclarationContainerClass(iv, descriptor, state)) return;
+
+        if (isTopLevelCallableReference(descriptor)) {
             // Note that this name is not used in reflection. There should be the name of the referenced declaration's module instead,
             // but there's no nice API to obtain that name here yet
             // TODO: write the referenced declaration's module name and use it in reflection
@@ -441,6 +451,12 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
         }
     }
 
+    public static boolean isTopLevelCallableReference(@NotNull CallableDescriptor descriptor) {
+        return descriptor instanceof LocalVariableDescriptor
+               ? DescriptorUtils.getParentOfType(descriptor, ClassDescriptor.class) == null
+               : descriptor.getContainingDeclaration() instanceof PackageFragmentDescriptor;
+    }
+
     @NotNull
     protected Method generateConstructor() {
         List<FieldInfo> args = calculateConstructorParameters(typeMapper, state.getLanguageVersionSettings(), closure, asmType);
@@ -476,24 +492,38 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
 
             iv.load(0, superClassAsmType);
 
-            String superClassConstructorDescriptor;
+            List<Type> superCtorArgTypes = new ArrayList<>();
             if (superClassAsmType.equals(LAMBDA) || superClassAsmType.equals(FUNCTION_REFERENCE) ||
+                superClassAsmType.equals(FUNCTION_REFERENCE_IMPL) ||
                 CoroutineCodegenUtilKt.isCoroutineSuperClass(state.getLanguageVersionSettings(), superClassAsmType.getInternalName())) {
                 int arity = calculateArity();
                 iv.iconst(arity);
+                superCtorArgTypes.add(Type.INT_TYPE);
                 if (shouldHaveBoundReferenceReceiver) {
-                    CallableReferenceUtilKt.loadBoundReferenceReceiverParameter(iv, boundReceiverParameterIndex, boundReceiverType, boundReceiverKotlinType);
-                    superClassConstructorDescriptor = "(ILjava/lang/Object;)V";
+                    CallableReferenceUtilKt.loadBoundReferenceReceiverParameter(
+                            iv, boundReceiverParameterIndex, boundReceiverType, boundReceiverKotlinType
+                    );
+                    superCtorArgTypes.add(OBJECT_TYPE);
                 }
-                else {
-                    superClassConstructorDescriptor = "(I)V";
+                if (isOptimizedFunctionReference) {
+                    assert functionReferenceTarget != null : "No function reference target: " + funDescriptor;
+                    generateCallableReferenceDeclarationContainerClass(iv, functionReferenceTarget, state);
+                    iv.aconst(functionReferenceTarget.getName().asString());
+                    PropertyReferenceCodegen.generateCallableReferenceSignature(iv, functionReferenceTarget, state);
+                    iv.aconst(isTopLevelCallableReference(functionReferenceTarget) ? 1 : 0);
+                    superCtorArgTypes.add(JAVA_CLASS_TYPE);
+                    superCtorArgTypes.add(JAVA_STRING_TYPE);
+                    superCtorArgTypes.add(JAVA_STRING_TYPE);
+                    superCtorArgTypes.add(Type.INT_TYPE);
                 }
             }
             else {
                 assert !shouldHaveBoundReferenceReceiver : "Unexpected bound reference with supertype " + superClassAsmType;
-                superClassConstructorDescriptor = "()V";
             }
-            iv.invokespecial(superClassAsmType.getInternalName(), "<init>", superClassConstructorDescriptor, false);
+            iv.invokespecial(
+                    superClassAsmType.getInternalName(), "<init>",
+                    Type.getMethodDescriptor(Type.VOID_TYPE, superCtorArgTypes.toArray(new Type[0])), false
+            );
 
             iv.visitInsn(RETURN);
 

compiler/backend/src/org/jetbrains/kotlin/codegen/JvmRuntimeTypes.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.builtins.createFunctionType
 import org.jetbrains.kotlin.codegen.coroutines.coroutinesJvmInternalPackageFqName
 import org.jetbrains.kotlin.codegen.coroutines.getOrCreateJvmSuspendFunctionView
 import org.jetbrains.kotlin.codegen.coroutines.isSuspendLambdaOrLocalFunction
+import org.jetbrains.kotlin.config.ApiVersion
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.config.isReleaseCoroutines
@@ -30,14 +31,22 @@ class JvmRuntimeTypes(module: ModuleDescriptor, private val languageVersionSetti
     private val kotlinCoroutinesJvmInternalPackage =
         MutablePackageFragmentDescriptor(module, languageVersionSettings.coroutinesJvmInternalPackageFqName())
 
+    val generateOptimizedCallableReferenceSuperClasses = languageVersionSettings.apiVersion >= ApiVersion.KOTLIN_1_4
+
     private fun internal(className: String, packageFragment: PackageFragmentDescriptor = kotlinJvmInternalPackage): Lazy<ClassDescriptor> =
         lazy { createClass(packageFragment, className) }
 
     private fun coroutinesInternal(name: String): Lazy<ClassDescriptor> =
         lazy { createCoroutineSuperClass(name) }
 
+    private fun propertyClasses(prefix: String, suffix: String): Lazy<List<ClassDescriptor>> =
+        lazy { (0..2).map { i -> createClass(kotlinJvmInternalPackage, prefix + i + suffix) } }
+
     private val lambda: ClassDescriptor by internal("Lambda")
+
     private val functionReference: ClassDescriptor by internal("FunctionReference")
+    val functionReferenceImpl: ClassDescriptor by internal("FunctionReferenceImpl")
+
     private val localVariableReference: ClassDescriptor by internal("LocalVariableReference")
     private val mutableLocalVariableReference: ClassDescriptor by internal("MutableLocalVariableReference")
 
@@ -60,13 +69,10 @@ class JvmRuntimeTypes(module: ModuleDescriptor, private val languageVersionSetti
             coroutineImpl
     }
 
-    private val propertyReferences: List<ClassDescriptor> by lazy {
-        (0..2).map { i -> createClass(kotlinJvmInternalPackage, "PropertyReference$i") }
-    }
-
-    private val mutablePropertyReferences: List<ClassDescriptor> by lazy {
-        (0..2).map { i -> createClass(kotlinJvmInternalPackage, "MutablePropertyReference$i") }
-    }
+    private val propertyReferences: List<ClassDescriptor> by propertyClasses("PropertyReference", "")
+    private val mutablePropertyReferences: List<ClassDescriptor> by propertyClasses("MutablePropertyReference", "")
+    private val propertyReferenceImpls: List<ClassDescriptor> by propertyClasses("PropertyReference", "Impl")
+    private val mutablePropertyReferenceImpls: List<ClassDescriptor> by propertyClasses("MutablePropertyReference", "Impl")
 
     private fun createClass(
         packageFragment: PackageFragmentDescriptor,
@@ -144,7 +150,8 @@ class JvmRuntimeTypes(module: ModuleDescriptor, private val languageVersionSetti
         )
 
         val suspendFunctionType = if (referencedFunction.isSuspend) suspendFunctionInterface?.defaultType else null
-        return listOfNotNull(functionReference.defaultType, functionType, suspendFunctionType)
+        val superClass = if (generateOptimizedCallableReferenceSuperClasses) functionReferenceImpl else functionReference
+        return listOfNotNull(superClass.defaultType, functionType, suspendFunctionType)
     }
 
     fun getSupertypeForPropertyReference(descriptor: VariableDescriptorWithAccessors, isMutable: Boolean, isBound: Boolean): KotlinType {
@@ -157,6 +164,11 @@ class JvmRuntimeTypes(module: ModuleDescriptor, private val languageVersionSetti
                     (if (descriptor.dispatchReceiverParameter != null) 1 else 0) -
                     if (isBound) 1 else 0
 
-        return (if (isMutable) mutablePropertyReferences else propertyReferences)[arity].defaultType
+        val classes = when {
+            generateOptimizedCallableReferenceSuperClasses -> if (isMutable) mutablePropertyReferenceImpls else propertyReferenceImpls
+            else -> if (isMutable) mutablePropertyReferences else propertyReferences
+        }
+
+        return classes[arity].defaultType
     }
 }

compiler/backend/src/org/jetbrains/kotlin/codegen/MemberCodegen.java

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.codegen.inline.SourceMapper;
 import org.jetbrains.kotlin.codegen.state.GenerationState;
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper;
 import org.jetbrains.kotlin.codegen.state.TypeMapperUtilsKt;
+import org.jetbrains.kotlin.config.ApiVersion;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.annotations.AnnotatedImpl;
 import org.jetbrains.kotlin.descriptors.annotations.Annotations;
@@ -642,6 +643,8 @@ public abstract class MemberCodegen<T extends KtPureElement/* TODO: & KtDeclarat
 
         if (!state.getClassBuilderMode().generateBodies) return;
 
+        boolean generateClassIntCtorCall = state.getJvmRuntimeTypes().getGenerateOptimizedCallableReferenceSuperClasses();
+
         InstructionAdapter iv = createOrGetClInitCodegen().v;
         iv.iconst(delegatedProperties.size());
         iv.newarray(K_PROPERTY_TYPE);
@@ -658,14 +661,29 @@ public abstract class MemberCodegen<T extends KtPureElement/* TODO: & KtDeclarat
             iv.anew(implType);
             iv.dup();
 
-            // TODO: generate the container once and save to a local field instead (KT-10495)
-            ClosureCodegen.generateCallableReferenceDeclarationContainer(iv, property, state);
+            List<Type> superCtorArgTypes = new ArrayList<>();
+            if (generateClassIntCtorCall) {
+                ClosureCodegen.generateCallableReferenceDeclarationContainerClass(iv, property, state);
+                superCtorArgTypes.add(JAVA_CLASS_TYPE);
+            } else {
+                // TODO: generate the container once and save to a local field instead (KT-10495)
+                ClosureCodegen.generateCallableReferenceDeclarationContainer(iv, property, state);
+                superCtorArgTypes.add(K_DECLARATION_CONTAINER_TYPE);
+            }
+
             iv.aconst(property.getName().asString());
             PropertyReferenceCodegen.generateCallableReferenceSignature(iv, property, state);
+            superCtorArgTypes.add(JAVA_STRING_TYPE);
+            superCtorArgTypes.add(JAVA_STRING_TYPE);
+
+            if (generateClassIntCtorCall) {
+                iv.aconst(ClosureCodegen.isTopLevelCallableReference(property) ? 1 : 0);
+                superCtorArgTypes.add(Type.INT_TYPE);
+            }
 
             iv.invokespecial(
                     implType.getInternalName(), "<init>",
-                    Type.getMethodDescriptor(Type.VOID_TYPE, K_DECLARATION_CONTAINER_TYPE, JAVA_STRING_TYPE, JAVA_STRING_TYPE), false
+                    Type.getMethodDescriptor(Type.VOID_TYPE, superCtorArgTypes.toArray(new Type[0])), false
             );
             Method wrapper = PropertyReferenceCodegen.getWrapperMethodForPropertyReference(property, receiverCount);
             iv.invokestatic(REFLECTION, wrapper.getName(), wrapper.getDescriptor(), false);

compiler/backend/src/org/jetbrains/kotlin/codegen/PropertyReferenceCodegen.kt

@@ -114,16 +114,16 @@ class PropertyReferenceCodegen(
 
         generateConstructor()
 
-        generateMethod("property reference getName", ACC_PUBLIC, method("getName", JAVA_STRING_TYPE)) {
-            aconst(target.name.asString())
-        }
-
-        generateMethod("property reference getSignature", ACC_PUBLIC, method("getSignature", JAVA_STRING_TYPE)) {
-            generateCallableReferenceSignature(this, target, state)
-        }
-
-        generateMethod("property reference getOwner", ACC_PUBLIC, method("getOwner", K_DECLARATION_CONTAINER_TYPE)) {
-            ClosureCodegen.generateCallableReferenceDeclarationContainer(this, target, state)
+        if (!isOptimizedPropertyReferenceSupertype(superAsmType)) {
+            generateMethod("property reference getName", ACC_PUBLIC, method("getName", JAVA_STRING_TYPE)) {
+                aconst(target.name.asString())
+            }
+            generateMethod("property reference getSignature", ACC_PUBLIC, method("getSignature", JAVA_STRING_TYPE)) {
+                generateCallableReferenceSignature(this, target, state)
+            }
+            generateMethod("property reference getOwner", ACC_PUBLIC, method("getOwner", K_DECLARATION_CONTAINER_TYPE)) {
+                ClosureCodegen.generateCallableReferenceDeclarationContainer(this, target, state)
+            }
         }
 
         if (!isLocalDelegatedProperty) {
@@ -136,16 +136,31 @@ class PropertyReferenceCodegen(
             val shouldHaveBoundReferenceReceiver = closure.isForBoundCallableReference()
             val receiverIndexAndFieldInfo = generateClosureFieldsInitializationFromParameters(closure, constructorArgs)
 
-            if (receiverIndexAndFieldInfo == null) {
-                assert(!shouldHaveBoundReferenceReceiver) { "No bound reference receiver in constructor parameters: $constructorArgs" }
-                load(0, OBJECT_TYPE)
-                invokespecial(superAsmType.internalName, "<init>", "()V", false)
-            } else {
+            load(0, OBJECT_TYPE)
+            val superCtorArgTypes = mutableListOf<Type>()
+            if (receiverIndexAndFieldInfo != null) {
                 val (receiverIndex, receiverFieldInfo) = receiverIndexAndFieldInfo
-                load(0, OBJECT_TYPE)
                 loadBoundReferenceReceiverParameter(receiverIndex, receiverFieldInfo.fieldType, receiverFieldInfo.fieldKotlinType)
-                invokespecial(superAsmType.internalName, "<init>", "(Ljava/lang/Object;)V", false)
+                superCtorArgTypes.add(OBJECT_TYPE)
+            } else {
+                assert(!shouldHaveBoundReferenceReceiver) { "No bound reference receiver in constructor parameters: $constructorArgs" }
             }
+
+            if (isOptimizedPropertyReferenceSupertype(superAsmType)) {
+                ClosureCodegen.generateCallableReferenceDeclarationContainerClass(this, target, state)
+                aconst(target.name.asString())
+                generateCallableReferenceSignature(this, target, state)
+                aconst(if (ClosureCodegen.isTopLevelCallableReference(target)) 1 else 0)
+                superCtorArgTypes.add(JAVA_CLASS_TYPE)
+                superCtorArgTypes.add(JAVA_STRING_TYPE)
+                superCtorArgTypes.add(JAVA_STRING_TYPE)
+                superCtorArgTypes.add(Type.INT_TYPE)
+            }
+
+            invokespecial(
+                superAsmType.internalName, "<init>",
+                Type.getMethodDescriptor(Type.VOID_TYPE, *superCtorArgTypes.toTypedArray()), false
+            )
         }
     }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/LambdaInfo.kt

@@ -141,7 +141,7 @@ abstract class DefaultLambda(
                 interfaces: Array<out String>?
             ) {
                 isPropertyReference = superName in PROPERTY_REFERENCE_SUPER_CLASSES
-                isFunctionReference = superName == FUNCTION_REFERENCE.internalName
+                isFunctionReference = superName == FUNCTION_REFERENCE.internalName || superName == FUNCTION_REFERENCE_IMPL.internalName
 
                 super.visit(version, access, name, signature, superName, interfaces)
             }
@@ -204,10 +204,12 @@ abstract class DefaultLambda(
     protected abstract fun mapAsmSignature(sourceCompiler: SourceCompilerForInline): Method
 
     private companion object {
-        val PROPERTY_REFERENCE_SUPER_CLASSES = listOf(
-            PROPERTY_REFERENCE0, PROPERTY_REFERENCE1, PROPERTY_REFERENCE2,
-            MUTABLE_PROPERTY_REFERENCE0, MUTABLE_PROPERTY_REFERENCE1, MUTABLE_PROPERTY_REFERENCE2
-        ).mapTo(HashSet(), Type::getInternalName)
+        val PROPERTY_REFERENCE_SUPER_CLASSES =
+            listOf(
+                PROPERTY_REFERENCE0, PROPERTY_REFERENCE1, PROPERTY_REFERENCE2,
+                MUTABLE_PROPERTY_REFERENCE0, MUTABLE_PROPERTY_REFERENCE1, MUTABLE_PROPERTY_REFERENCE2
+            ).plus(OPTIMIZED_PROPERTY_REFERENCE_SUPERTYPES)
+                .mapTo(HashSet(), Type::getInternalName)
     }
 }