FIR2IR: refactor implicit cast insertion, part 2: coerce-to-Unit

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrImplicitCastInserter.kt

@@ -24,10 +24,10 @@ import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.coerceToUnitIfNeeded
 import org.jetbrains.kotlin.name.Name
 
 class Fir2IrImplicitCastInserter(
@@ -43,6 +43,35 @@ class Fir2IrImplicitCastInserter(
         TODO("Should not be here: ${element.render()}")
     }
 
+    // TODO: can be private once this visitor becomes more comprehensive
+    internal fun IrContainerExpression.insertImplicitCasts(): IrContainerExpression {
+        if (statements.isEmpty()) return this
+
+        val lastIndex = statements.lastIndex
+        statements.forEachIndexed { i, irStatement ->
+            if (irStatement !is IrErrorCallExpression && irStatement is IrExpression) {
+                if (i != lastIndex) {
+                    statements[i] = irStatement.coerceToUnitIfNeeded(irStatement.type, irBuiltIns)
+                } else {
+                    // TODO: for the last statement, need to cast to the return type if mismatched
+                }
+            }
+        }
+
+        return this
+    }
+
+    internal fun IrBlockBody.insertImplicitCasts(): IrBlockBody {
+        if (statements.isEmpty()) return this
+
+        statements.forEachIndexed { i, irStatement ->
+            if (irStatement !is IrErrorCallExpression && irStatement is IrExpression) {
+                statements[i] = irStatement.coerceToUnitIfNeeded(irStatement.type, irBuiltIns)
+            }
+        }
+        return this
+    }
+
     override fun visitExpressionWithSmartcast(expressionWithSmartcast: FirExpressionWithSmartcast, data: IrElement): IrExpression {
         return implicitCastOrExpression(data as IrExpression, expressionWithSmartcast.typeRef)
     }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -36,7 +36,6 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.util.coerceToUnitIfNeeded
 import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.parentClassOrNull
@@ -553,7 +552,11 @@ class Fir2IrVisitor(
                 } else {
                     emptyList()
                 }
-            ).insertImplicitCasts()
+            ).also {
+                with(implicitCastInserter) {
+                    it.insertImplicitCasts()
+                }
+            }
         }
     }
 
@@ -571,42 +574,19 @@ class Fir2IrVisitor(
                 IrCompositeImpl(
                     startOffset, endOffset, type, origin,
                     mapToIrStatements().filterNotNull()
-                ).insertImplicitCasts()
+                )
             } else {
                 IrBlockImpl(
                     startOffset, endOffset, type, origin,
                     mapToIrStatements().filterNotNull()
-                ).insertImplicitCasts()
+                )
-            }
+            }.also {
-        }
+                // TODO: can remove this once implicit cast inserter visits more expression kinds directly
-    }
+                with(implicitCastInserter) {
-
+                    it.insertImplicitCasts()
-    private fun IrBlockBody.insertImplicitCasts(): IrBlockBody {
-        if (statements.isEmpty()) return this
-
-        statements.forEachIndexed { i, irStatement ->
-            if (irStatement !is IrErrorCallExpression && irStatement is IrExpression) {
-                statements[i] = irStatement.coerceToUnitIfNeeded(irStatement.type, irBuiltIns)
-            }
-        }
-        return this
-    }
-
-    private fun IrContainerExpression.insertImplicitCasts(): IrContainerExpression {
-        if (statements.isEmpty()) return this
-
-        val lastIndex = statements.lastIndex
-        statements.forEachIndexed { i, irStatement ->
-            if (irStatement !is IrErrorCallExpression && irStatement is IrExpression) {
-                if (i != lastIndex) {
-                    statements[i] = irStatement.coerceToUnitIfNeeded(irStatement.type, irBuiltIns)
-                } else {
-                    // TODO: for the last statement, need to cast to the return type if mismatched
                 }
             }
         }
-
-        return this
     }
 
     override fun visitErrorExpression(errorExpression: FirErrorExpression, data: Any?): IrElement {