[FIR] Add ability to pass data=null to transformInplace with data producer

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)

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
+        }
     }
 }

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)
     }