FIR: set super type reference properly

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

@@ -376,11 +376,15 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
                 qualifiedAccessExpression.resultType = FirResolvedTypeRefImpl(session, null, type, emptyList())
             }
             is FirSuperReference -> {
-                qualifiedAccessExpression.resultType =
-                    callee.superTypeRef as? FirResolvedTypeRef
-                        ?: implicitReceiverStack.filterIsInstance<ImplicitDispatchReceiverValue>().lastOrNull()
-                            ?.boundSymbol?.fir?.superTypeRefs?.firstOrNull()
-                                ?: FirErrorTypeRefImpl(session, qualifiedAccessExpression.psi, "No super type")
+                if (callee.superTypeRef is FirResolvedTypeRef) {
+                    qualifiedAccessExpression.resultType = callee.superTypeRef
+                } else {
+                    val superTypeRef = implicitReceiverStack.filterIsInstance<ImplicitDispatchReceiverValue>().lastOrNull()
+                        ?.boundSymbol?.fir?.superTypeRefs?.firstOrNull()
+                        ?: FirErrorTypeRefImpl(session, qualifiedAccessExpression.psi, "No super type")
+                    qualifiedAccessExpression.resultType = superTypeRef
+                    callee.replaceSuperTypeRef(superTypeRef)
+                }
             }
             is FirResolvedCallableReference -> {
                 if (qualifiedAccessExpression.typeRef !is FirResolvedTypeRef) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSuperReference.kt

@@ -11,6 +11,8 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitor
 interface FirSuperReference : FirReference {
     val superTypeRef: FirTypeRef
 
+    fun replaceSuperTypeRef(newSuperTypeRef: FirTypeRef)
+
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
         visitor.visitSuperReference(this, data)
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/references/FirExplicitSuperReference.kt

@@ -15,6 +15,10 @@ class FirExplicitSuperReference(
     psi: PsiElement?,
     override var superTypeRef: FirTypeRef
 ) : FirAbstractElement(session, psi), FirSuperReference {
+    override fun replaceSuperTypeRef(newSuperTypeRef: FirTypeRef) {
+        superTypeRef = newSuperTypeRef
+    }
+
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
         superTypeRef = superTypeRef.transformSingle(transformer, data)
         return this