From 698f78570d9a7410b93d6de817984188fad50fbb Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Mon, 16 Mar 2020 10:45:28 +0300
Subject: [PATCH] [FIR] Add ability to pass `data=null` to transformInplace
 with data producer

---
 .../resolve/FirExpressionsResolveTransformer.kt     |  8 +++-----
 .../kotlin/fir/expressions/FirExpressionUtil.kt     | 13 +++++++++----
 .../kotlin/fir/visitors/FirTransformerUtil.kt       | 12 ++++++++++--
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
index 7d338090c3c..c9f9f07632a 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
@@ -37,10 +37,7 @@ import org.jetbrains.kotlin.fir.symbols.invoke
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
-import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
-import org.jetbrains.kotlin.fir.visitors.compose
-import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
@@ -188,7 +185,8 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         val numberOfStatements = block.statements.size
 
         block.transformStatementsIndexed(transformer) { index ->
-            if (index == numberOfStatements - 1) data else ResolutionMode.ContextIndependent
+            val value = if (index == numberOfStatements - 1) data else ResolutionMode.ContextIndependent
+            TransformData.Data(value)
         }
         block.transformAllStatementsExceptLast(integerLiteralTypeApproximator, builtinTypes.intType.type)
         block.transformOtherChildren(transformer, data)
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt
index 505d372de6b..95ee46fb12a 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt
@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.TransformData
 import org.jetbrains.kotlin.fir.visitors.transformInplace
 import org.jetbrains.kotlin.name.ClassId
 
@@ -64,19 +65,23 @@ fun buildErrorExpression(source: FirSourceElement?, diagnostic: FirDiagnostic):
     }
 }
 
-fun <D : Any> FirBlock.transformStatementsIndexed(transformer: FirTransformer<D>, dataProducer: (Int) -> D?): FirBlock {
+fun <D> FirBlock.transformStatementsIndexed(transformer: FirTransformer<D>, dataProducer: (Int) -> TransformData<D>): FirBlock {
     when (this) {
         is FirBlockImpl -> statements.transformInplace(transformer, dataProducer)
         is FirSingleExpressionBlock -> {
-            dataProducer(0)?.let { transformStatements(transformer, it) }
+            (dataProducer(0) as? TransformData.Data<D>)?.value?.let { transformStatements(transformer, it) }
         }
     }
     return this
 }
 
-fun <D : Any> FirBlock.transformAllStatementsExceptLast(transformer: FirTransformer<D>, data: D): FirBlock {
+fun <D> FirBlock.transformAllStatementsExceptLast(transformer: FirTransformer<D>, data: D): FirBlock {
     val threshold = statements.size - 1
     return transformStatementsIndexed(transformer) { index ->
-        data.takeIf { index < threshold }
+        if (index < threshold) {
+            TransformData.Data(data)
+        } else {
+            TransformData.Nothing
+        }
     }
 }
\ No newline at end of file
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/visitors/FirTransformerUtil.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/visitors/FirTransformerUtil.kt
index a6050a742d6..80589c546f1 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/visitors/FirTransformerUtil.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/visitors/FirTransformerUtil.kt
@@ -34,12 +34,20 @@ fun <T : FirElement, D> MutableList<T>.transformInplace(transformer: FirTransfor
     }
 }
 
-inline fun <T : FirElement, D : Any> MutableList<T>.transformInplace(transformer: FirTransformer<D>, dataProducer: (Int) -> D?) {
+sealed class TransformData<out D> {
+    class Data<D>(val value: D) : TransformData<D>()
+    object Nothing : TransformData<kotlin.Nothing>()
+}
+
+inline fun <T : FirElement, D> MutableList<T>.transformInplace(transformer: FirTransformer<D>, dataProducer: (Int) -> TransformData<D>) {
     val iterator = this.listIterator()
     var index = 0
     while (iterator.hasNext()) {
         val next = iterator.next() as FirPureAbstractElement
-        val data = dataProducer(index++) ?: continue
+        val data = when (val data = dataProducer(index++)) {
+            is TransformData.Data<D> -> data.value
+            TransformData.Nothing -> continue
+        }
         val result = next.transform<T, D>(transformer, data).single
         iterator.set(result)
     }