[FIR] Adapt FIR utilities after FirExpression.coneTypeOrNull introduction

#KT-59855

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

@@ -46,7 +46,7 @@ fun FirRegularClassBuilder.generateValuesFunction(
         source = sourceElement
         this.origin = origin
         this.moduleData = moduleData
-        returnTypeRef = buildResolvedTypeRef {
+        val returnTypeRef = buildResolvedTypeRef {
             source = sourceElement
             type = ConeClassLikeTypeImpl(
                 StandardClassIds.Array.toLookupTag(),
@@ -60,6 +60,7 @@ fun FirRegularClassBuilder.generateValuesFunction(
                 isNullable = false
             )
         }
+        this.returnTypeRef = returnTypeRef
         name = ENUM_VALUES
         this.status = createStatus(this@generateValuesFunction.status).apply {
             isStatic = true
@@ -68,7 +69,7 @@ fun FirRegularClassBuilder.generateValuesFunction(
         symbol = FirNamedFunctionSymbol(CallableId(packageFqName, classFqName, ENUM_VALUES))
         resolvePhase = this@generateValuesFunction.resolvePhase
         body = buildEmptyExpressionBlock().also {
-            it.replaceTypeRef(returnTypeRef)
+            it.replaceType(returnTypeRef.type)
         }
     }.apply {
         containingClassForStaticMemberAttr = this@generateValuesFunction.symbol.toLookupTag()
@@ -87,7 +88,7 @@ fun FirRegularClassBuilder.generateValueOfFunction(
         source = sourceElement
         this.origin = origin
         this.moduleData = moduleData
-        returnTypeRef = buildResolvedTypeRef {
+        val returnTypeRef = buildResolvedTypeRef {
             source = sourceElement
             type = ConeClassLikeTypeImpl(
                 this@generateValueOfFunction.symbol.toLookupTag(),
@@ -95,6 +96,7 @@ fun FirRegularClassBuilder.generateValueOfFunction(
                 isNullable = false
             )
         }
+        this.returnTypeRef = returnTypeRef
         name = ENUM_VALUE_OF
 
         status = createStatus(this@generateValueOfFunction.status).apply {
@@ -107,7 +109,7 @@ fun FirRegularClassBuilder.generateValueOfFunction(
             containingFunctionSymbol = this@buildSimpleFunction.symbol
             this.origin = origin
             this.moduleData = moduleData
-            returnTypeRef = buildResolvedTypeRef {
+            this.returnTypeRef = buildResolvedTypeRef {
                 source = sourceElement
                 type = ConeClassLikeTypeImpl(
                     StandardClassIds.String.toLookupTag(),
@@ -124,7 +126,7 @@ fun FirRegularClassBuilder.generateValueOfFunction(
         }
         resolvePhase = this@generateValueOfFunction.resolvePhase
         body = buildEmptyExpressionBlock().also {
-            it.replaceTypeRef(returnTypeRef)
+            it.replaceType(returnTypeRef.type)
         }
     }.apply {
         containingClassForStaticMemberAttr = this@generateValueOfFunction.symbol.toLookupTag()

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

@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneTypeOrNull
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitTypeRefImplWithoutSource
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
@@ -36,7 +38,7 @@ fun FirVariable.toQualifiedAccess(
         name = this@toQualifiedAccess.name
         resolvedSymbol = this@toQualifiedAccess.symbol
     }
-    this.typeRef = typeRef
+    this.coneTypeOrNull = typeRef.coneTypeOrNull
 }
 
 fun generateTemporaryVariable(
@@ -85,7 +87,12 @@ fun generateExplicitReceiverTemporaryVariable(
                 source = source,
                 name = SpecialNames.RECEIVER,
                 initializer = receiver,
-                typeRef = receiver.typeRef.copyWithNewSource(source),
+                typeRef = receiver.coneTypeOrNull?.let {
+                    buildResolvedTypeRef {
+                        type = it
+                        this.source = source
+                    }
+                }
             ).also { property ->
                 // Change the expression from x.a to <receiver>.a
                 val newReceiverAccess =

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt

@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.TransformData
 import org.jetbrains.kotlin.fir.visitors.transformInplace
@@ -85,9 +84,7 @@ fun buildErrorLoop(source: KtSourceElement?, diagnostic: ConeDiagnostic): FirErr
         this.source = source
         this.diagnostic = diagnostic
     }.also {
-        it.block.replaceTypeRef(buildErrorTypeRef {
-            this.diagnostic = diagnostic
-        })
+        it.block.replaceConeTypeOrNull(ConeErrorType(diagnostic))
     }
 }
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirTypeUtils.kt

@@ -37,13 +37,11 @@ val FirTypeRef.coneType: ConeKotlinType
 val FirTypeRef.coneTypeOrNull: ConeKotlinType?
     get() = coneTypeSafe()
 
-val FirExpression.coneType: ConeKotlinType get() = typeRef.coneType
+val FirExpression.coneType: ConeKotlinType get() = requireNotNull(coneTypeOrNull) { "Expected type to be resolved" }
 
-val FirExpression.coneTypeOrNull: ConeKotlinType? get() = typeRef.coneTypeOrNull
+inline fun <reified T : ConeKotlinType> FirExpression.coneTypeSafe(): T? = (coneTypeOrNull as? T)
 
-inline fun <reified T : ConeKotlinType> FirExpression.coneTypeSafe(): T? = typeRef.coneTypeSafe()
-
-inline fun <reified T : ConeKotlinType> FirExpression.coneTypeUnsafe(): T = typeRef.coneTypeUnsafe()
+inline fun <reified T : ConeKotlinType> FirExpression.coneTypeUnsafe(): T = coneTypeOrNull as T
 
 @RequiresOptIn(
     "This type check never expands type aliases. Use with care (probably Ok for expression & constructor types). " +