FIR Java: support Kotlin extension overriding: at last, #KT-29937 Fixed

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/scopes/JavaClassEnhancementScope.kt

@@ -81,7 +81,7 @@ class JavaClassEnhancementScope(
         val firField = (original as FirBasedSymbol<*>).fir as? FirJavaField ?: error("Can't make enhancement for $original")
 
         val memberContext = context.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, jsr305State, firField.annotations)
-        val newReturnTypeRef = enhanceReturnType(firField, memberContext, null)
+        val newReturnTypeRef = enhanceReturnType(firField, emptyList(), memberContext, null)
 
         val symbol = FirPropertySymbol(original.callableId)
         with(firField) {
@@ -124,20 +124,26 @@ class JavaClassEnhancementScope(
             }
         }
 
-        val newReceiverTypeRef = if (firMethod is FirJavaMethod && firMethod.receiverTypeRef != null) {
-            enhanceReceiverType(firMethod, memberContext)
+        val overriddenMembers = firMethod.overriddenMembers()
+        val hasReceiver = overriddenMembers.any { it.receiverTypeRef != null }
+
+        val newReceiverTypeRef = if (firMethod is FirJavaMethod && hasReceiver) {
+            enhanceReceiverType(firMethod, overriddenMembers, memberContext)
         } else null
         val newReturnTypeRef = if (firMethod is FirJavaConstructor) {
             firMethod.returnTypeRef
         } else {
-            enhanceReturnType(firMethod, memberContext, predefinedEnhancementInfo)
+            enhanceReturnType(firMethod, overriddenMembers, memberContext, predefinedEnhancementInfo)
         }
 
         val newValueParameterInfo = mutableListOf<EnhanceValueParameterResult>()
 
         for ((index, valueParameter) in firMethod.valueParameters.withIndex()) {
+            if (hasReceiver && index == 0) continue
             newValueParameterInfo += enhanceValueParameter(
-                firMethod, memberContext, predefinedEnhancementInfo, valueParameter as FirJavaValueParameter, index
+                firMethod, overriddenMembers, hasReceiver,
+                memberContext, predefinedEnhancementInfo, valueParameter as FirJavaValueParameter,
+                if (hasReceiver) index - 1 else index
             )
         }
 
@@ -213,10 +219,12 @@ class JavaClassEnhancementScope(
 
     private fun enhanceReceiverType(
         ownerFunction: FirJavaMethod,
+        overriddenMembers: List<FirCallableMember>,
         memberContext: FirJavaEnhancementContext
     ): FirResolvedTypeRef {
         val signatureParts = ownerFunction.partsForValueParameter(
             typeQualifierResolver,
+            overriddenMembers,
             // TODO: check me
             parameterContainer = ownerFunction,
             methodContext = memberContext,
@@ -229,6 +237,8 @@ class JavaClassEnhancementScope(
 
     private fun enhanceValueParameter(
         ownerFunction: FirCallableMember,
+        overriddenMembers: List<FirCallableMember>,
+        hasReceiver: Boolean,
         memberContext: FirJavaEnhancementContext,
         predefinedEnhancementInfo: PredefinedFunctionEnhancementInfo?,
         ownerParameter: FirJavaValueParameter,
@@ -236,9 +246,10 @@ class JavaClassEnhancementScope(
     ): EnhanceValueParameterResult {
         val signatureParts = ownerFunction.partsForValueParameter(
             typeQualifierResolver,
+            overriddenMembers,
             parameterContainer = ownerParameter,
             methodContext = memberContext,
-            typeInSignature = TypeInSignature.ValueParameter(index)
+            typeInSignature = TypeInSignature.ValueParameter(hasReceiver, index)
         ).enhance(session, jsr305State, predefinedEnhancementInfo?.parametersInfo?.getOrNull(index))
         val firResolvedTypeRef = signatureParts.type
         val defaultValue = ownerParameter.getDefaultValueFromAnnotation()
@@ -252,11 +263,13 @@ class JavaClassEnhancementScope(
 
     private fun enhanceReturnType(
         owner: FirCallableMember,
+        overriddenMembers: List<FirCallableMember>,
         memberContext: FirJavaEnhancementContext,
         predefinedEnhancementInfo: PredefinedFunctionEnhancementInfo?
     ): FirResolvedTypeRef {
         val signatureParts = owner.parts(
             typeQualifierResolver,
+            overriddenMembers,
             typeContainer = owner, isCovariant = true,
             containerContext = memberContext,
             containerApplicabilityType =
@@ -297,22 +310,32 @@ class JavaClassEnhancementScope(
         }
 
         object Receiver : TypeInSignature() {
-            override fun getTypeRef(member: FirCallableMember): FirTypeRef = member.receiverTypeRef!!
+            override fun getTypeRef(member: FirCallableMember): FirTypeRef {
+                if (member is FirJavaMethod) return member.valueParameters[0].returnTypeRef
+                return member.receiverTypeRef!!
+            }
         }
 
-        class ValueParameter(val index: Int) : TypeInSignature() {
-            override fun getTypeRef(member: FirCallableMember): FirTypeRef = (member as FirFunction).valueParameters[index].returnTypeRef
+        class ValueParameter(val hasReceiver: Boolean, val index: Int) : TypeInSignature() {
+            override fun getTypeRef(member: FirCallableMember): FirTypeRef {
+                if (hasReceiver && member is FirJavaMethod) {
+                    return (member as FirFunction).valueParameters[index + 1].returnTypeRef
+                }
+                return (member as FirFunction).valueParameters[index].returnTypeRef
+            }
         }
     }
 
     private fun FirCallableMember.partsForValueParameter(
         typeQualifierResolver: FirAnnotationTypeQualifierResolver,
+        overriddenMembers: List<FirCallableMember>,
         // TODO: investigate if it's really can be a null (check properties' with extension overrides in Java)
         parameterContainer: FirAnnotationContainer?,
         methodContext: FirJavaEnhancementContext,
         typeInSignature: TypeInSignature
     ) = parts(
         typeQualifierResolver,
+        overriddenMembers,
         parameterContainer, false,
         parameterContainer?.let {
             methodContext.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, jsr305State, it.annotations)
@@ -323,6 +346,7 @@ class JavaClassEnhancementScope(
 
     private fun FirCallableMember.parts(
         typeQualifierResolver: FirAnnotationTypeQualifierResolver,
+        overriddenMembers: List<FirCallableMember>,
         typeContainer: FirAnnotationContainer?,
         isCovariant: Boolean,
         containerContext: FirJavaEnhancementContext,
@@ -334,7 +358,7 @@ class JavaClassEnhancementScope(
             typeQualifierResolver,
             typeContainer,
             typeRef as FirJavaTypeRef,
-            this.overriddenMembers().map {
+            overriddenMembers.map {
                 typeInSignature.getTypeRef(it)
             },
             isCovariant,

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/scopes/JavaClassUseSiteScope.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirNamedFunction
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.java.declarations.FirJavaMethod
 import org.jetbrains.kotlin.fir.java.toNotNullConeKotlinType
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
@@ -54,29 +55,30 @@ class JavaClassUseSiteScope(
     private fun isEqualTypes(a: FirTypeRef, b: FirTypeRef) =
         isEqualTypes(a.toNotNullConeKotlinType(session), b.toNotNullConeKotlinType(session))
 
-    private fun isOverriddenFunCheck(member: FirNamedFunction, self: FirNamedFunction): Boolean {
-        return member.valueParameters.size == self.valueParameters.size &&
-                member.valueParameters.zip(self.valueParameters).all { (memberParam, selfParam) ->
-                    isEqualTypes(memberParam.returnTypeRef, selfParam.returnTypeRef)
-                }
+    private fun isOverriddenFunCheck(overriddenInJava: FirJavaMethod, base: FirNamedFunction): Boolean {
+        val receiverTypeRef = base.receiverTypeRef
+        if (receiverTypeRef == null) {
+            return overriddenInJava.valueParameters.size == base.valueParameters.size &&
+                    overriddenInJava.valueParameters.zip(base.valueParameters).all { (memberParam, selfParam) ->
+                        isEqualTypes(memberParam.returnTypeRef, selfParam.returnTypeRef)
+                    }
+        } else {
+            if (overriddenInJava.valueParameters.size != base.valueParameters.size + 1) return false
+            val baseParameterTypes = listOf(receiverTypeRef) + base.valueParameters.map { it.returnTypeRef }
+            val javaParameterTypes = overriddenInJava.valueParameters.map { it.returnTypeRef }
+            return baseParameterTypes.zip(javaParameterTypes).all { (baseType, typeFromJava) ->
+                isEqualTypes(baseType, typeFromJava)
+            }
+        }
     }
 
     internal fun ConeFunctionSymbol.getOverridden(candidates: Set<ConeFunctionSymbol>): ConeCallableSymbol? {
         if (overriddenByBase.containsKey(this)) return overriddenByBase[this]
 
-        fun sameReceivers(memberTypeRef: FirTypeRef?, selfTypeRef: FirTypeRef?): Boolean {
-            return when {
-                memberTypeRef != null && selfTypeRef != null -> isEqualTypes(memberTypeRef, selfTypeRef)
-                else -> memberTypeRef == null && selfTypeRef == null
-            }
-        }
-
         val self = (this as FirFunctionSymbol).fir as FirNamedFunction
         val overriding = candidates.firstOrNull {
-            val member = (it as FirFunctionSymbol).fir as FirNamedFunction
-            self.modality != Modality.FINAL
-                    && sameReceivers(member.receiverTypeRef, self.receiverTypeRef)
-                    && isOverriddenFunCheck(member, self)
+            val member = (it as FirFunctionSymbol).fir as FirJavaMethod
+            self.modality != Modality.FINAL && isOverriddenFunCheck(member, self)
         } // TODO: two or more overrides for one fun?
         overriddenByBase[this] = overriding
         return overriding

idea/testData/fir/multiModule/javaInheritsKotlinExtension/extraDump.java.txt

@@ -0,0 +1,4 @@
+public open class Inheritor : R|Second| {
+    public open operator function foo R|First|.(first: R|kotlin/String|, s: R|kotlin/Int|): R|kotlin/Unit|
+
+}

idea/testData/fir/multiModule/javaInheritsKotlinExtension/jvm/Base.kt

@@ -0,0 +1,5 @@
+interface First
+
+open class Second {
+    open fun First.foo(s: String, i: Int) {}
+}

idea/testData/fir/multiModule/javaInheritsKotlinExtension/jvm/Base.txt

@@ -0,0 +1,10 @@
+FILE: Base.kt
+    public abstract interface First {
+    }
+    public open class Second {
+        public constructor(): super<R|kotlin/Any|>()
+
+        public open function foo R|First|.(s: R|kotlin/String|, i: R|kotlin/Int|): R|kotlin/Unit| {
+        }
+
+    }

idea/testData/fir/multiModule/javaInheritsKotlinExtension/jvm/Inheritor.java

@@ -0,0 +1,3 @@
+public class Inheritor extends Second {
+    public void foo(First first, String s, int i) {}
+}

idea/testData/fir/multiModule/javaInheritsKotlinExtension/jvm/Test.kt

@@ -0,0 +1,5 @@
+class Tester : Inheritor(), First {
+    fun test() {
+        foo("abc", 456)
+    }
+}

idea/testData/fir/multiModule/javaInheritsKotlinExtension/jvm/Test.txt

@@ -0,0 +1,9 @@
+FILE: Test.kt
+    public final class Tester : R|Inheritor|, R|First| {
+        public constructor(): super<R|Inheritor|>()
+
+        public final function test(): R|kotlin/Unit| {
+            R|/Inheritor.foo|(String(abc), Int(456))
+        }
+
+    }

idea/tests/org/jetbrains/kotlin/idea/fir/FirMultiModuleResolveTestGenerated.java

@@ -64,6 +64,11 @@ public class FirMultiModuleResolveTestGenerated extends AbstractFirMultiModuleRe
         runTest("idea/testData/fir/multiModule/fullWithJavaPredefinedSignature/");
     }
 
+    @TestMetadata("javaInheritsKotlinExtension")
+    public void testJavaInheritsKotlinExtension() throws Exception {
+        runTest("idea/testData/fir/multiModule/javaInheritsKotlinExtension/");
+    }
+
     @TestMetadata("javaInheritsRawKotlin")
     public void testJavaInheritsRawKotlin() throws Exception {
         runTest("idea/testData/fir/multiModule/javaInheritsRawKotlin/");