Implement fake override mapping for functions

Related to KT-29636

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirAccessResolveTransformer.kt

@@ -14,15 +14,17 @@ import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedCallableReferenceImpl
+import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction.NEXT
-import org.jetbrains.kotlin.fir.scopes.impl.FirClassDeclaredMemberScope
-import org.jetbrains.kotlin.fir.scopes.impl.FirClassUseSiteScope
-import org.jetbrains.kotlin.fir.scopes.impl.FirCompositeScope
-import org.jetbrains.kotlin.fir.scopes.impl.FirTopLevelDeclaredMemberScope
+import org.jetbrains.kotlin.fir.scopes.impl.*
 import org.jetbrains.kotlin.fir.symbols.ConeCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassErrorType
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypedProjection
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
 
@@ -35,6 +37,21 @@ class FirAccessResolveTransformer : FirAbstractTreeTransformerWithSuperTypes(rev
         }
     }
 
+    private fun ConeClassLikeType.buildSubstitutionScope(
+        useSiteSession: FirSession,
+        unsubstituted: FirScope,
+        regularClass: FirRegularClass
+    ): FirClassSubstitutionScope? {
+        if (this.typeArguments.isEmpty()) return null
+
+        @Suppress("UNCHECKED_CAST")
+        val substitution = regularClass.typeParameters.zip(this.typeArguments) { typeParameter, typeArgument ->
+            typeParameter.symbol to (typeArgument as? ConeTypedProjection)?.type
+        }.filter { (_, type) -> type != null }.toMap() as Map<ConeTypeParameterSymbol, ConeKotlinType>
+
+        return FirClassSubstitutionScope(useSiteSession, unsubstituted, substitution, true)
+    }
+
     private fun FirRegularClass.buildUseSiteScope(useSiteSession: FirSession = session): FirClassUseSiteScope {
         val superTypeScope = FirCompositeScope(mutableListOf())
         val declaredScope = FirClassDeclaredMemberScope(this, useSiteSession)
@@ -43,7 +60,8 @@ class FirAccessResolveTransformer : FirAbstractTreeTransformerWithSuperTypes(rev
                 if (useSiteSuperType is ConeClassErrorType) return@mapNotNullTo null
                 val symbol = useSiteSuperType.symbol
                 if (symbol is FirClassSymbol) {
-                    symbol.fir.buildUseSiteScope(useSiteSession)
+                    val scope = symbol.fir.buildUseSiteScope(useSiteSession)
+                    useSiteSuperType.buildSubstitutionScope(useSiteSession, scope, symbol.fir) ?: scope
                 } else {
                     null
                 }
@@ -58,10 +76,10 @@ class FirAccessResolveTransformer : FirAbstractTreeTransformerWithSuperTypes(rev
         }
     }
 
-    var lookupFunctions = false
-    var lookupProperties = false
+    private var lookupFunctions = false
+    private var lookupProperties = false
 
-    inline fun <T> withNewSettings(block: () -> T): T {
+    private inline fun <T> withNewSettings(block: () -> T): T {
         val prevFunctions = lookupFunctions
         val prevProperties = lookupProperties
         val result = block()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirClassSubstitutionScope.kt

@@ -0,0 +1,149 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.scopes.impl
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirNamedFunction
+import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.declarations.impl.FirMemberFunctionImpl
+import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.ProcessorAction
+import org.jetbrains.kotlin.fir.symbols.*
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.impl.ConeAbbreviatedTypeImpl
+import org.jetbrains.kotlin.fir.types.impl.ConeClassTypeImpl
+import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeImpl
+import org.jetbrains.kotlin.name.Name
+
+class FirClassSubstitutionScope(
+    session: FirSession,
+    private val unsubstituted: FirScope,
+    private val substitution: Map<ConeTypeParameterSymbol, ConeKotlinType>,
+    lookupInFir: Boolean
+) : FirAbstractProviderBasedScope(session, lookupInFir) {
+
+    val fakeOverrides = mutableMapOf<ConeCallableSymbol, ConeCallableSymbol>()
+
+    private fun wrapProjection(old: ConeKotlinTypeProjection, newType: ConeKotlinType): ConeKotlinTypeProjection {
+        return when (old) {
+            is StarProjection -> old
+            is ConeKotlinTypeProjectionIn -> ConeKotlinTypeProjectionIn(newType)
+            is ConeKotlinTypeProjectionOut -> ConeKotlinTypeProjectionOut(newType)
+            is ConeKotlinType -> newType
+            else -> old
+        }
+    }
+
+    private fun ConeKotlinType.substitute(): ConeKotlinType? {
+        if (this is ConeTypeParameterType) return substitution[this]
+
+        val newArguments by lazy { arrayOfNulls<ConeKotlinTypeProjection>(typeArguments.size) }
+        var initialized = false
+        for ((index, typeArgument) in this.typeArguments.withIndex()) {
+            val type = (typeArgument as? ConeTypedProjection)?.type ?: continue
+            val newType = type.substitute()
+            if (newType != null) {
+                initialized = true
+                newArguments[index] = wrapProjection(typeArgument, newType)
+            }
+        }
+
+        if (initialized) {
+            for ((index, typeArgument) in this.typeArguments.withIndex()) {
+                if (newArguments[index] == null) {
+                    newArguments[index] = typeArgument
+                }
+            }
+            @Suppress("UNCHECKED_CAST")
+            return when (this) {
+                is ConeKotlinErrorType -> error("Trying to substitute arguments for error type")
+                is ConeTypeParameterType -> error("Trying to substitute arguments for type parameter")
+                is ConeClassTypeImpl -> ConeClassTypeImpl(symbol, newArguments as Array<ConeKotlinTypeProjection>)
+                is ConeAbbreviatedTypeImpl -> ConeAbbreviatedTypeImpl(
+                    abbreviationSymbol,
+                    newArguments as Array<ConeKotlinTypeProjection>,
+                    directExpansion.substitute() as? ConeClassLikeType ?: directExpansion
+                )
+                is ConeFunctionType -> TODO("Substitute function type properly")
+                is ConeClassLikeType -> error("Unknown class-like type to substitute: $this, ${this::class}")
+            }
+        }
+        return null
+    }
+
+
+    override fun processFunctionsByName(name: Name, processor: (ConeFunctionSymbol) -> ProcessorAction): ProcessorAction {
+        unsubstituted.processFunctionsByName(name) process@{ original ->
+
+            val function = fakeOverrides.getOrPut(original) { createFakeOverride(original, name) }
+            processor(function as ConeFunctionSymbol)
+        }
+
+
+        return super.processFunctionsByName(name, processor)
+    }
+
+    override fun processPropertiesByName(name: Name, processor: (ConePropertySymbol) -> ProcessorAction): ProcessorAction {
+        return unsubstituted.processPropertiesByName(name, processor)
+    }
+
+    private fun createFakeOverride(
+        original: ConeFunctionSymbol,
+        name: Name
+    ): FirFunctionSymbol {
+        val member = (original as FirBasedSymbol<*>).fir as? FirNamedFunction ?: error("Can't fake override for $original")
+        val receiverType = member.receiverType?.coneTypeUnsafe()
+        val newReceiverType = receiverType?.substitute()
+
+        val returnType = member.returnType.coneTypeUnsafe()
+        val newReturnType = returnType.substitute()
+
+        val newParameterTypes = member.valueParameters.map {
+            it.returnType.coneTypeUnsafe().substitute()
+        }
+
+        val symbol = FirFunctionSymbol(original.callableId, true)
+        with(member) {
+            // TODO: consider using here some light-weight functions instead of pseudo-real FirMemberFunctionImpl
+            // As second alternative, we can invent some light-weight kind of FirRegularClass
+            FirMemberFunctionImpl(
+                session,
+                psi,
+                symbol,
+                name,
+                member.receiverType?.withReplacedConeType(newReceiverType),
+                member.returnType.withReplacedConeType(newReturnType)
+            ).apply {
+                status = member.status as FirDeclarationStatusImpl
+                valueParameters += member.valueParameters.zip(newParameterTypes) { valueParameter, newType ->
+                    with(valueParameter) {
+                        FirValueParameterImpl(
+                            session, psi,
+                            name, this.returnType.withReplacedConeType(newType),
+                            defaultValue, isCrossinline, isNoinline, isVararg
+                        )
+                    }
+                }
+            }
+        }
+        return symbol
+    }
+}
+
+
+fun FirType.withReplacedConeType(newType: ConeKotlinType?): FirResolvedType {
+    require(this is FirResolvedType)
+    if (newType == null) return this
+
+    return FirResolvedTypeImpl(
+        session, psi, newType,
+        isNullable,
+        annotations
+    )
+
+}

compiler/fir/resolve/testData/resolve/overrides/simpleFakeOverride.kt

@@ -0,0 +1,15 @@
+
+open class A<T> {
+    fun foo(t: T): T {
+        return t
+    }
+}
+
+class Some
+
+class B : A<Some>() {
+    fun test() {
+        foo(Some())
+    }
+}
+

compiler/fir/resolve/testData/resolve/overrides/simpleFakeOverride.txt

@@ -0,0 +1,21 @@
+FILE: simpleFakeOverride.kt
+    <T> public open class A {
+        public constructor(): super<R|kotlin/Any|>()
+
+        public final function foo(t: R|T|): R|T| {
+            return@@@foo <Unresolved name: t>#
+        }
+
+    }
+    public final class Some {
+        public constructor(): super<R|kotlin/Any|>()
+
+    }
+    public final class B : R|A<Some>| {
+        public constructor(): super<R|A<Some>|>()
+
+        public final function test(): R|kotlin/Unit| {
+            R|FakeOverride</A.foo>|(<Unresolved name: Some>#())
+        }
+
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirResolveTestCaseGenerated.java

@@ -259,6 +259,11 @@ public class FirResolveTestCaseGenerated extends AbstractFirResolveTestCase {
         public void testSimple() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/overrides/simple.kt");
         }
+
+        @TestMetadata("simpleFakeOverride.kt")
+        public void testSimpleFakeOverride() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/overrides/simpleFakeOverride.kt");
+        }
     }
 
     @TestMetadata("compiler/fir/resolve/testData/resolve/references")