From 318f4b5f0236187517aac9a5dad1a1cfa27020d3 Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Fri, 23 Oct 2020 15:07:35 -0700
Subject: [PATCH] FIR2IR: refactor implicit cast insertion, part 3: implicit
 not_null

---
 .../fir/backend/Fir2IrImplicitCastInserter.kt | 29 ++++++++++++++++---
 .../kotlin/fir/backend/Fir2IrVisitor.kt       | 27 ++++-------------
 2 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrImplicitCastInserter.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrImplicitCastInserter.kt
index fd7ca27c800..eb2f769db42 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrImplicitCastInserter.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrImplicitCastInserter.kt
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.backend
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.FirThisReceiverExpression
 import org.jetbrains.kotlin.fir.references.FirReference
@@ -18,10 +19,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.types.ConeIntersectionType
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.expressions.*
@@ -43,6 +41,29 @@ class Fir2IrImplicitCastInserter(
         TODO("Should not be here: ${element.render()}")
     }
 
+    // TODO: can be private once this visitor becomes more comprehensive, also generalized
+    internal fun IrExpression.insertImplicitNotNullCastIfNeeded(expression: FirExpression): IrExpression {
+        if (this is IrTypeOperatorCall && this.operator == IrTypeOperator.IMPLICIT_NOTNULL) {
+            return this
+        }
+        // TODO: Other conditions to check?
+        // [TypeOperatorLowering] will retrieve the source (from start offset to end offset) as an assertion message.
+        // Avoid type casting if we can't determine the source for some reasons, e.g., implicit `this` receiver.
+        if (expression.source == null ||
+            expression.typeRef.coneTypeSafe<ConeKotlinType>()?.hasEnhancedNullability != true
+        ) {
+            return this
+        }
+        return IrTypeOperatorCallImpl(
+            this.startOffset,
+            this.endOffset,
+            this.type,
+            IrTypeOperator.IMPLICIT_NOTNULL,
+            this.type,
+            this
+        )
+    }
+
     // TODO: can be private once this visitor becomes more comprehensive
     internal fun IrContainerExpression.insertImplicitCasts(): IrContainerExpression {
         if (statements.isEmpty()) return this
diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
index 851584c80cf..da293eec1fc 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
@@ -469,29 +469,12 @@ class Fir2IrVisitor(
                     expression.accept(this, null) as IrExpression
                 }
             }
-        }.insertImplicitNotNullCastIfNeeded(expression)
-    }
-
-    private fun IrExpression.insertImplicitNotNullCastIfNeeded(expression: FirExpression): IrExpression {
-        if (this is IrTypeOperatorCall && this.operator == IrTypeOperator.IMPLICIT_NOTNULL) {
-            return this
+        }.let {
+            // TODO: expression(implicitCastInserter, it) as IrExpression
+            with(implicitCastInserter) {
+                it.insertImplicitNotNullCastIfNeeded(expression)
+            }
         }
-        // TODO: Other conditions to check?
-        // [TypeOperatorLowering] will retrieve the source (from start offset to end offset) as an assertion message.
-        // Avoid type casting if we can't determine the source for some reasons, e.g., implicit `this` receiver.
-        if (expression.source == null ||
-            expression.typeRef.coneTypeSafe<ConeKotlinType>()?.hasEnhancedNullability != true
-        ) {
-            return this
-        }
-        return IrTypeOperatorCallImpl(
-            this.startOffset,
-            this.endOffset,
-            this.type,
-            IrTypeOperator.IMPLICIT_NOTNULL,
-            this.type,
-            this
-        )
     }
 
     private fun convertToIrReceiverExpression(