From 27687602f88a48e9b4c35e471b00726a4b381ad6 Mon Sep 17 00:00:00 2001
From: Simon Ogorodnik <Simon.Ogorodnik@jetbrains.com>
Date: Thu, 14 Mar 2019 17:47:02 +0300
Subject: [PATCH] FIR: Add functions to change nullability/arguments of
 ConeTypes

---
 .../kotlin/fir/resolve/ResolveUtils.kt        | 45 ++++++++++++++++---
 1 file changed, 40 insertions(+), 5 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt
index bd5ecc0caac..ba4cc33213e 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt
@@ -6,18 +6,15 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
 import org.jetbrains.kotlin.fir.declarations.expandedConeType
-import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
-import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.scopes.impl.FirCompositeScope
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.symbols.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 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.ConeFunctionTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
@@ -93,3 +90,41 @@ private fun List<FirQualifierPart>.toTypeProjections(): Array<ConeKotlinTypeProj
 }.toTypedArray()
 
 
+fun <T : ConeKotlinType> T.withNullability(nullability: ConeNullability): T {
+    if (this.nullability == nullability) {
+        return this
+    }
+
+    return when (this) {
+        is ConeClassErrorType -> this
+        is ConeClassTypeImpl -> ConeClassTypeImpl(lookupTag, typeArguments, nullability.isNullable) as T
+        is ConeAbbreviatedTypeImpl -> ConeAbbreviatedTypeImpl(
+            abbreviationLookupTag,
+            typeArguments,
+            nullability.isNullable
+        ) as T
+        is ConeFunctionTypeImpl -> ConeFunctionTypeImpl(receiverType, parameterTypes, returnType, lookupTag, nullability.isNullable) as T
+        is ConeTypeParameterTypeImpl -> ConeTypeParameterTypeImpl(lookupTag, nullability.isNullable) as T
+        is ConeFlexibleType -> ConeFlexibleType(lowerBound.withNullability(nullability), upperBound.withNullability(nullability)) as T
+        else -> TODO("FIX KOTLIN COMPILER")
+    }
+}
+
+
+fun <T : ConeKotlinType> T.withArguments(arguments: Array<ConeKotlinTypeProjection>): T {
+    if (this.typeArguments === arguments) {
+        return this
+    }
+
+    return when (this) {
+        is ConeClassErrorType -> this
+        is ConeClassTypeImpl -> ConeClassTypeImpl(lookupTag, arguments, nullability.isNullable) as T
+        is ConeAbbreviatedTypeImpl -> ConeAbbreviatedTypeImpl(
+            abbreviationLookupTag,
+            arguments,
+            nullability.isNullable
+        ) as T
+        //is ConeFunctionTypeImpl -> ConeFunctionTypeImpl(receiverType, parameterTypes, returnType, lookupTag, nullability.isNullable) as T
+        else -> error("Not supported: $this: ${this.render()}")
+    }
+}
\ No newline at end of file