Fix container for extension receiver types from Java

No tests are added since there are some already
(e.g. TypeEnhancement.testOverriddenExtensions)

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/descriptors/JavaMethodDescriptor.java

@@ -26,8 +26,15 @@ import org.jetbrains.kotlin.types.KotlinType;
 import org.jetbrains.kotlin.util.OperatorChecks;
 
 import java.util.List;
+import java.util.Map;
 
 public class JavaMethodDescriptor extends SimpleFunctionDescriptorImpl implements JavaCallableMemberDescriptor {
+    // TODO: It's only used to retrieve annotations from the original value parameter and can be removed when
+    // org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl.initialize accepts extension parameter descriptor
+    // instead of type
+    public static final UserDataKey<ValueParameterDescriptor> ORIGINAL_VALUE_PARAMETER_FOR_EXTENSION_RECEIVER =
+            new UserDataKey<ValueParameterDescriptor>() {};
+
     private enum ParameterNamesStatus {
         NON_STABLE_DECLARED(false, false),
         STABLE_DECLARED(true, false),
@@ -82,11 +89,12 @@ public class JavaMethodDescriptor extends SimpleFunctionDescriptorImpl implement
             @NotNull List<ValueParameterDescriptor> unsubstitutedValueParameters,
             @Nullable KotlinType unsubstitutedReturnType,
             @Nullable Modality modality,
-            @NotNull Visibility visibility
+            @NotNull Visibility visibility,
+            @Nullable Map<? extends UserDataKey<?>, ?> userData
     ) {
         SimpleFunctionDescriptorImpl descriptor = super.initialize(
                 receiverParameterType, dispatchReceiverParameter, typeParameters, unsubstitutedValueParameters,
-                unsubstitutedReturnType, modality, visibility);
+                unsubstitutedReturnType, modality, visibility, userData);
         setOperator(OperatorChecks.INSTANCE.check(descriptor).isSuccess());
         return descriptor;
     }

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaScope.kt

@@ -130,7 +130,11 @@ abstract class LazyJavaScope(protected val c: LazyJavaResolverContext) : MemberS
                 effectiveSignature.valueParameters,
                 effectiveSignature.returnType,
                 Modality.convertFromFlags(method.isAbstract, !method.isFinal),
-                method.visibility
+                method.visibility,
+                if (effectiveSignature.receiverType != null)
+                    mapOf(JavaMethodDescriptor.ORIGINAL_VALUE_PARAMETER_FOR_EXTENSION_RECEIVER to valueParameters.descriptors.first())
+                else
+                    emptyMap<FunctionDescriptor.UserDataKey<ValueParameterDescriptor>, ValueParameterDescriptor>()
         )
 
         functionDescriptorImpl.setParameterNamesStatus(effectiveSignature.hasStableParameterNames, valueParameters.hasSynthesizedNames)

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/typeEnhancement/signatureEnhancement.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.load.java.typeEnhancement
 
 import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotated
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.descriptors.annotations.composeAnnotations
@@ -35,6 +36,7 @@ import org.jetbrains.kotlin.types.asFlexibleType
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.types.isFlexible
 import org.jetbrains.kotlin.types.typeUtil.isTypeParameter
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import java.util.*
 
 class SignatureEnhancement(private val annotationTypeQualifierResolver: AnnotationTypeQualifierResolver) {
@@ -56,7 +58,12 @@ class SignatureEnhancement(private val annotationTypeQualifierResolver: Annotati
 
         val receiverTypeEnhancement =
                 if (extensionReceiverParameter != null)
-                    parts(typeContainer = null, isCovariant = false) { it.extensionReceiverParameter!!.type }.enhance()
+                    parts(
+                            typeContainer =
+                                this.safeAs<FunctionDescriptor>()
+                                    ?.getUserData(JavaMethodDescriptor.ORIGINAL_VALUE_PARAMETER_FOR_EXTENSION_RECEIVER),
+                            isCovariant = false
+                    ) { it.extensionReceiverParameter!!.type }.enhance()
                 else null