Extract type preparation during type checking to a separate component

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -393,22 +393,6 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
         return ConeTypeIntersector.intersectTypes(this as ConeInferenceContext, types as List<ConeKotlinType>)
     }
 
-    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
-        return when (type) {
-            is ConeClassLikeType -> type.fullyExpandedType(session)
-            is ConeFlexibleType -> {
-                val lowerBound = prepareType(type.lowerBound)
-                if (lowerBound === type.lowerBound) return type
-
-                ConeFlexibleType(
-                    lowerBound as ConeKotlinType,
-                    prepareType(type.upperBound) as ConeKotlinType
-                )
-            }
-            else -> type
-        }
-    }
-
     override fun KotlinTypeMarker.isNullableType(): Boolean {
         require(this is ConeKotlinType)
         if (this.isMarkedNullable)
@@ -551,7 +535,8 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 class ConeTypeCheckerContext(
     override val isErrorTypeEqualsToAnything: Boolean,
     override val isStubTypeEqualsToAnything: Boolean,
-    override val typeSystemContext: ConeInferenceContext
+    override val typeSystemContext: ConeInferenceContext,
+    val kotlinTypePreparator: ConeTypePreparator = ConeTypePreparator.getDefault(typeSystemContext.session),
 ) : AbstractTypeCheckerContext() {
 
     val session: FirSession = typeSystemContext.session
@@ -585,7 +570,11 @@ class ConeTypeCheckerContext(
     }
 
     override fun refineType(type: KotlinTypeMarker): KotlinTypeMarker {
-        return typeSystemContext.prepareType(type)
+        return kotlinTypePreparator.prepareType(type)
+    }
+
+    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+        return kotlinTypePreparator.prepareType(type)
     }
 
     override val KotlinTypeMarker.isAllowedTypeVariable: Boolean

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypePreparator.kt

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.types
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.types.AbstractTypePreparator
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+
+abstract class ConeTypePreparator(val session: FirSession) : AbstractTypePreparator {
+    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+        return when (type) {
+            is ConeClassLikeType -> type.fullyExpandedType(session)
+            is ConeFlexibleType -> {
+                val lowerBound = prepareType(type.lowerBound)
+                if (lowerBound === type.lowerBound) return type
+
+                ConeFlexibleType(
+                    lowerBound as ConeKotlinType,
+                    prepareType(type.upperBound) as ConeKotlinType
+                )
+            }
+            else -> type
+        }
+    }
+
+    companion object {
+        fun getDefault(session: FirSession) = object : ConeTypePreparator(session) {}
+    }
+}