Extract bangbang and bindLabel methods (raw FIR)

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.fir.lightTree.converter.ConverterUtil.extractArgumen
 import org.jetbrains.kotlin.fir.lightTree.converter.ConverterUtil.getAsStringWithoutBacktick
 import org.jetbrains.kotlin.fir.lightTree.converter.ConverterUtil.isExpression
 import org.jetbrains.kotlin.fir.lightTree.converter.ConverterUtil.nameAsSafeName
-import org.jetbrains.kotlin.fir.lightTree.converter.utils.bangBangToWhen
 import org.jetbrains.kotlin.fir.lightTree.converter.utils.generateDestructuringBlock
 import org.jetbrains.kotlin.fir.lightTree.converter.utils.getOperationSymbol
 import org.jetbrains.kotlin.fir.lightTree.converter.utils.qualifiedAccessTokens
@@ -313,7 +312,7 @@ class ExpressionsConverter(
 
         val operationToken = operationTokenName.getOperationSymbol()
         if (operationToken == EXCLEXCL) {
-            return bangBangToWhen(session, getAsFirExpression(argument, "No operand"))
+            return argument.bangBangToWhen(null) { toFirExpression(it) }
         }
 
         val conventionCallName = operationToken.toUnaryName()
@@ -875,34 +874,14 @@ class ExpressionsConverter(
      */
     private fun convertLoopJump(jump: LighterASTNode): FirExpression {
         var isBreak = true
-        var labelName: String? = null
         jump.forEachChildren {
             when (it.tokenType) {
                 CONTINUE_KEYWORD -> isBreak = false
                 //BREAK -> isBreak = true
-                LABEL_QUALIFIER -> labelName = it.asText.replace("@", "")
             }
         }
 
-        return (if (isBreak) FirBreakExpressionImpl(session, null) else FirContinueExpressionImpl(session, null)).apply {
+        return (if (isBreak) FirBreakExpressionImpl(session, null) else FirContinueExpressionImpl(session, null)).bindLabel(jump)
-            target = FirLoopTarget(labelName)
-            val lastLoop = context.firLoops.lastOrNull()
-            if (labelName == null) {
-                if (lastLoop != null) {
-                    target.bind(lastLoop)
-                } else {
-                    target.bind(FirErrorLoop(this@ExpressionsConverter.session, null, "Cannot bind unlabeled jump to a loop"))
-                }
-            } else {
-                for (firLoop in context.firLoops.asReversed()) {
-                    if (firLoop.label?.name == labelName) {
-                        target.bind(firLoop)
-                        return this
-                    }
-                }
-                target.bind(FirErrorLoop(this@ExpressionsConverter.session, null, "Cannot bind label $labelName to a loop"))
-            }
-        }
     }
 
     /**

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/utils/ExpressionUtil.kt

@@ -63,14 +63,3 @@ fun generateDestructuringBlock(
         }
     }
 }
-
-fun bangBangToWhen(session: FirSession, baseExpression: FirExpression): FirWhenExpression {
-    return baseExpression.generateNotNullOrOther(
-        session,
-        FirThrowExpressionImpl(
-            session, null, FirFunctionCallImpl(session, null).apply {
-                calleeReference = FirSimpleNamedReference(session, null, RawFirBuilder.KNPE)
-            }
-        ), "bangbang", null
-    )
-}

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -9,11 +9,15 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.KtNodeTypes.*
 import org.jetbrains.kotlin.fir.FirFunctionTarget
-import org.jetbrains.kotlin.fir.FirLabel
+import org.jetbrains.kotlin.fir.FirLoopTarget
 import org.jetbrains.kotlin.fir.FirReference
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.declarations.FirNamedFunction
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirErrorFunction
+import org.jetbrains.kotlin.fir.declarations.impl.FirErrorLoop
 import org.jetbrains.kotlin.fir.declarations.impl.FirTypeParameterImpl
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.*
@@ -30,10 +34,11 @@ import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.lexer.KtTokens.CLOSING_QUOTE
 import org.jetbrains.kotlin.lexer.KtTokens.OPEN_QUOTE
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtStringTemplateExpression
+import org.jetbrains.kotlin.psi.KtUnaryExpression
 import org.jetbrains.kotlin.resolve.constants.evaluate.*
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -89,6 +94,10 @@ abstract class BaseFirBuilder<T>(val session: FirSession, val context: Context =
     }
 
     /**** Common utils ****/
+    companion object {
+        val KNPE = Name.identifier("KotlinNullPointerException")
+    }
+
     fun FirExpression.toReturn(basePsi: PsiElement? = psi, labelName: String? = null): FirReturnExpression {
         return FirReturnExpressionImpl(
             this@BaseFirBuilder.session,
@@ -149,6 +158,17 @@ abstract class BaseFirBuilder<T>(val session: FirSession, val context: Context =
             ?: emptyList()
     }
 
+    fun T?.bangBangToWhen(parent: KtUnaryExpression?, convert: T?.(String) -> FirExpression): FirWhenExpression {
+        return this.convert("No operand").generateNotNullOrOther(
+            session,
+            FirThrowExpressionImpl(
+                session, parent.getPsiOrNull(), FirFunctionCallImpl(session, parent.getPsiOrNull()).apply {
+                    calleeReference = FirSimpleNamedReference(this@BaseFirBuilder.session, parent, KNPE)
+                }
+            ), "bangbang", parent
+        )
+    }
+
     fun FirAbstractLoop.configure(generateBlock: () -> FirBlock): FirAbstractLoop {
         label = context.firLabels.pop()
         context.firLoops += this
@@ -157,6 +177,28 @@ abstract class BaseFirBuilder<T>(val session: FirSession, val context: Context =
         return this
     }
 
+    fun FirLoopJump.bindLabel(expression: T): FirLoopJump {
+        val labelName = expression.getLabelName()
+        target = FirLoopTarget(labelName)
+        val lastLoop = context.firLoops.lastOrNull()
+        if (labelName == null) {
+            if (lastLoop != null) {
+                target.bind(lastLoop)
+            } else {
+                target.bind(FirErrorLoop(this@BaseFirBuilder.session, expression.getPsiOrNull(), "Cannot bind unlabeled jump to a loop"))
+            }
+        } else {
+            for (firLoop in context.firLoops.asReversed()) {
+                if (firLoop.label?.name == labelName) {
+                    target.bind(firLoop)
+                    return this
+                }
+            }
+            target.bind(FirErrorLoop(this@BaseFirBuilder.session, expression.getPsiOrNull(), "Cannot bind label $labelName to a loop"))
+        }
+        return this
+    }
+
     /**** Conversion utils ****/
     private fun <T> T.getPsiOrNull(): PsiElement? {
         return if (this is PsiElement) this else null
@@ -181,7 +223,9 @@ abstract class BaseFirBuilder<T>(val session: FirSession, val context: Context =
                     )
                 } else if (convertedText is Number) {
                     // TODO: support byte / short
-                    FirConstExpressionImpl(session, expression.getPsiOrNull(), IrConstKind.Int, convertedText.toInt(), "Incorrect int: $text")
+                    FirConstExpressionImpl(
+                        session, expression.getPsiOrNull(), IrConstKind.Int, convertedText.toInt(), "Incorrect int: $text"
+                    )
                 } else {
                     FirErrorExpressionImpl(session, expression.getPsiOrNull(), reason = "Incorrect constant expression: $text")
                 }

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1035,28 +1035,6 @@ class RawFirBuilder(session: FirSession, val stubMode: Boolean) : BaseFirBuilder
             }
         }
 
-        private fun FirLoopJump.bindLabel(expression: KtExpressionWithLabel): FirLoopJump {
-            val labelName = expression.getLabelName()
-            target = FirLoopTarget(labelName)
-            val lastLoop = context.firLoops.lastOrNull()
-            if (labelName == null) {
-                if (lastLoop != null) {
-                    target.bind(lastLoop)
-                } else {
-                    target.bind(FirErrorLoop(psi, "Cannot bind unlabeled jump to a loop"))
-                }
-            } else {
-                for (firLoop in context.firLoops.asReversed()) {
-                    if (firLoop.label?.name == labelName) {
-                        target.bind(firLoop)
-                        return this
-                    }
-                }
-                target.bind(FirErrorLoop(psi, "Cannot bind label $labelName to a loop"))
-            }
-            return this
-        }
-
         override fun visitBreakExpression(expression: KtBreakExpression, data: Unit): FirElement {
             return FirBreakExpressionImpl(expression).bindLabel(expression)
         }
@@ -1065,17 +1043,6 @@ class RawFirBuilder(session: FirSession, val stubMode: Boolean) : BaseFirBuilder
             return FirContinueExpressionImpl(expression).bindLabel(expression)
         }
 
-        private fun KtUnaryExpression.bangBangToWhen(): FirWhenExpression {
-            return baseExpression.toFirExpression("No operand").generateNotNullOrOther(
-                session,
-                FirThrowExpressionImpl(
-                    this, FirFunctionCallImpl(this).apply {
-                        calleeReference = FirSimpleNamedReference(this@bangBangToWhen, KNPE)
-                    }
-                ), "bangbang", this
-            )
-        }
-
         override fun visitBinaryExpression(expression: KtBinaryExpression, data: Unit): FirElement {
             val operationToken = expression.operationToken
             val leftArgument = expression.left.toFirExpression("No left operand")
@@ -1137,7 +1104,7 @@ class RawFirBuilder(session: FirSession, val stubMode: Boolean) : BaseFirBuilder
             val operationToken = expression.operationToken
             val argument = expression.baseExpression
             if (operationToken == EXCLEXCL) {
-                return expression.bangBangToWhen()
+                return expression.baseExpression.bangBangToWhen(expression) { (this as KtExpression).toFirExpression(it) }
             }
             val conventionCallName = operationToken.toUnaryName()
             return if (conventionCallName != null) {
@@ -1298,8 +1265,4 @@ class RawFirBuilder(session: FirSession, val stubMode: Boolean) : BaseFirBuilder
             return FirExpressionStub(expression)
         }
     }
-
-    companion object {
-        val KNPE = Name.identifier("KotlinNullPointerException")
-    }
 }