FIR: get rid of ANDAND / OROR by replacing with when

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.references.FirExplicitThisReference
 import org.jetbrains.kotlin.fir.references.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitBooleanTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitTypeRefImpl
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -158,8 +159,6 @@ internal fun IElementType.toFirOperation(): FirOperation =
         KtTokens.EXCLEQ -> FirOperation.NOT_EQ
         KtTokens.EQEQEQ -> FirOperation.IDENTITY
         KtTokens.EXCLEQEQEQ -> FirOperation.NOT_IDENTITY
-        KtTokens.ANDAND -> FirOperation.AND
-        KtTokens.OROR -> FirOperation.OR
         KtTokens.IN_KEYWORD -> FirOperation.IN
         KtTokens.NOT_IN -> FirOperation.NOT_IN
         KtTokens.RANGE -> FirOperation.RANGE
@@ -201,6 +200,85 @@ internal fun FirExpression.generateNotNullOrOther(
     }
 }
 
+internal fun FirExpression.generateLazyLogicalOperation(
+    session: FirSession, other: FirExpression, isAnd: Boolean, basePsi: KtElement
+): FirWhenExpression {
+    val terminalExpression = FirConstExpressionImpl(session, psi, IrConstKind.Boolean, !isAnd)
+    val terminalBlock = FirSingleExpressionBlock(session, terminalExpression)
+    val otherBlock = FirSingleExpressionBlock(session, other)
+    return FirWhenExpressionImpl(session, basePsi).apply {
+        branches += FirWhenBranchImpl(
+            session, psi, this@generateLazyLogicalOperation,
+            if (isAnd) otherBlock else terminalBlock
+        )
+        branches += FirWhenBranchImpl(
+            session, other.psi, FirElseIfTrueCondition(session, psi),
+            if (isAnd) terminalBlock else otherBlock
+        )
+        typeRef = FirImplicitBooleanTypeRef(session, basePsi)
+    }
+}
+
+internal fun KtWhenCondition.toFirWhenCondition(
+    session: FirSession,
+    convert: KtExpression?.(String) -> FirExpression,
+    toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef
+): FirExpression {
+    val firSubjectExpression = FirWhenSubjectExpression(session, this)
+    return when (this) {
+        is KtWhenConditionWithExpression -> {
+            FirOperatorCallImpl(
+                session,
+                expression,
+                FirOperation.EQ
+            ).apply {
+                arguments += firSubjectExpression
+                arguments += expression.convert("No expression in condition with expression")
+            }
+        }
+        is KtWhenConditionInRange -> {
+            FirOperatorCallImpl(
+                session,
+                rangeExpression,
+                if (isNegated) FirOperation.NOT_IN else FirOperation.IN
+            ).apply {
+                arguments += firSubjectExpression
+                arguments += rangeExpression.convert("No range in condition with range")
+            }
+        }
+        is KtWhenConditionIsPattern -> {
+            FirTypeOperatorCallImpl(
+                session, typeReference, if (isNegated) FirOperation.NOT_IS else FirOperation.IS,
+                typeReference.toFirOrErrorTypeRef()
+            ).apply {
+                arguments += firSubjectExpression
+            }
+        }
+        else -> {
+            FirErrorExpressionImpl(session, this, "Unsupported when condition: ${this.javaClass}")
+        }
+    }
+}
+
+internal fun Array<KtWhenCondition>.toFirWhenCondition(
+    session: FirSession,
+    basePsi: KtElement,
+    convert: KtExpression?.(String) -> FirExpression,
+    toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef
+): FirExpression {
+    var firCondition: FirExpression? = null
+    for (condition in this) {
+        val firConditionElement = condition.toFirWhenCondition(session, convert, toFirOrErrorTypeRef)
+        firCondition = when (firCondition) {
+            null -> firConditionElement
+            else -> firCondition.generateLazyLogicalOperation(
+                session, firConditionElement, false, basePsi
+            )
+        }
+    }
+    return firCondition!!
+}
+
 internal fun generateIncrementOrDecrementBlock(
     session: FirSession,
     baseExpression: KtUnaryExpression,

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

@@ -989,42 +989,6 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
             }
         }
 
-        private fun KtWhenCondition.toFirWhenCondition(firSubjectExpression: FirExpression): FirExpression {
-            return when (this) {
-                is KtWhenConditionWithExpression -> {
-                    FirOperatorCallImpl(
-                        session,
-                        expression,
-                        FirOperation.EQ
-                    ).apply {
-                        arguments += firSubjectExpression
-                        arguments += expression.toFirExpression("No expression in condition with expression")
-                    }
-                }
-                is KtWhenConditionInRange -> {
-                    FirOperatorCallImpl(
-                        session,
-                        rangeExpression,
-                        if (isNegated) FirOperation.NOT_IN else FirOperation.IN
-                    ).apply {
-                        arguments += firSubjectExpression
-                        arguments += rangeExpression.toFirExpression("No range in condition with range")
-                    }
-                }
-                is KtWhenConditionIsPattern -> {
-                    FirTypeOperatorCallImpl(
-                        session, typeReference, if (isNegated) FirOperation.NOT_IS else FirOperation.IS,
-                        typeReference.toFirOrErrorType()
-                    ).apply {
-                        arguments += firSubjectExpression
-                    }
-                }
-                else -> {
-                    FirErrorExpressionImpl(session, this, "Unsupported when condition: ${this.javaClass}")
-                }
-            }
-        }
-
         override fun visitWhenExpression(expression: KtWhenExpression, data: Unit): FirElement {
             val subjectExpression = expression.subjectExpression
             val subject = when (subjectExpression) {
@@ -1050,25 +1014,13 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
                     val branch = entry.expression.toFirBlock()
                     branches += if (!entry.isElse) {
                         if (hasSubject) {
-                            var firCondition: FirExpression? = null
+                            val firCondition = entry.conditions.toFirWhenCondition(
-                            for (condition in entry.conditions) {
+                                this@RawFirBuilder.session,
-                                val firConditionElement = condition.toFirWhenCondition(
+                                entry,
-                                    FirWhenSubjectExpression(this@RawFirBuilder.session, condition)
+                                { toFirExpression(it) },
-                                )
+                                { toFirOrErrorType() }
-                                when {
+                            )
-                                    firCondition == null -> firCondition = firConditionElement
+                            FirWhenBranchImpl(this@RawFirBuilder.session, entry, firCondition, branch)
-                                    firCondition is FirOperatorCallImpl && firCondition.operation == FirOperation.OR -> {
-                                        firCondition.arguments += firConditionElement
-                                    }
-                                    else -> {
-                                        firCondition = FirOperatorCallImpl(this@RawFirBuilder.session, entry, FirOperation.OR).apply {
-                                            arguments += firCondition!!
-                                            arguments += firConditionElement
-                                        }
-                                    }
-                                }
-                            }
-                            FirWhenBranchImpl(this@RawFirBuilder.session, entry, firCondition!!, branch)
                         } else {
                             val condition = entry.conditions.first() as? KtWhenConditionWithExpression
                             val firCondition = condition?.expression.toFirExpression("No expression in condition with expression")
@@ -1194,12 +1146,6 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
             return FirContinueExpressionImpl(session, expression).bindLabel(expression)
         }
 
-        private fun KtBinaryExpression.elvisToWhen(): FirWhenExpression {
-            val rightArgument = right.toFirExpression("No right operand")
-            val leftArgument = left.toFirExpression("No left operand")
-            return leftArgument.generateNotNullOrOther(session, rightArgument, "elvis", this)
-        }
-
         private fun KtUnaryExpression.bangBangToWhen(): FirWhenExpression {
             return baseExpression.toFirExpression("No operand").generateNotNullOrOther(
                 session,
@@ -1213,11 +1159,15 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
 
         override fun visitBinaryExpression(expression: KtBinaryExpression, data: Unit): FirElement {
             val operationToken = expression.operationToken
+            val leftArgument = expression.left.toFirExpression("No left operand")
             val rightArgument = expression.right.toFirExpression("No right operand")
             if (operationToken == ELVIS) {
-                return expression.elvisToWhen()
+                return leftArgument.generateNotNullOrOther(session, rightArgument, "elvis", expression)
             }
             val conventionCallName = operationToken.toBinaryName()
+            if (operationToken == ANDAND || operationToken == OROR) {
+                return leftArgument.generateLazyLogicalOperation(session, rightArgument, operationToken == ANDAND, expression)
+            }
             return if (conventionCallName != null || operationToken == IDENTIFIER) {
                 FirFunctionCallImpl(
                     session, expression
@@ -1237,7 +1187,7 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
                     FirOperatorCallImpl(session, expression, firOperation)
                 }
             }.apply {
-                arguments += expression.left.toFirExpression("No left operand")
+                arguments += leftArgument
                 arguments += rightArgument
             }
         }

compiler/fir/psi2fir/testData/rawBuilder/expressions/branches.txt

@@ -37,7 +37,15 @@ FILE: branches.kt
     }
     public? final? fun grade(g: Int): String {
         ^grade when (g#) {
-            ||(==($subj$, Int(6)), ==($subj$, Int(7))) ->  {
+            when () {
+                ==($subj$, Int(6)) ->  {
+                    Boolean(true)
+                }
+                else ->  {
+                    ==($subj$, Int(7))
+                }
+            }
+             ->  {
                 String(Outstanding)
             }
             ==($subj$, Int(5)) ->  {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt

@@ -293,11 +293,13 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
     }
 
     override fun transformWhenExpression(whenExpression: FirWhenExpression, data: Any?): CompositeTransformResult<FirStatement> {
-        val whenExpression = super.transformWhenExpression(whenExpression, data).single as FirWhenExpression
+        whenExpression.transformChildren(this, data)
-        val type = commonSuperType(whenExpression.branches.mapNotNull {
+        if (whenExpression.resultType !is FirResolvedTypeRef) {
-            it.result.resultType
+            val type = commonSuperType(whenExpression.branches.mapNotNull {
-        })
+                it.result.resultType
-        if (type != null) whenExpression.resultType = type
+            })
+            if (type != null) whenExpression.resultType = type
+        }
         return whenExpression.compose()
     }
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirOperation.kt

@@ -16,8 +16,6 @@ enum class FirOperation(val operator: String = "???") {
     GT(">"),
     LT_EQ("<="),
     GT_EQ(">="),
-    AND("&&"),
-    OR("||"),
     IN("in"),
     NOT_IN("!in"),
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirImplicitBuiltinTypeRef.kt

@@ -43,4 +43,9 @@ class FirImplicitEnumTypeRef(
 class FirImplicitAnnotationTypeRef(
     session: FirSession,
     psi: PsiElement?
-) : FirImplicitBuiltinTypeRef(session, psi, KotlinBuiltIns.FQ_NAMES.annotation.toUnsafe())
+) : FirImplicitBuiltinTypeRef(session, psi, KotlinBuiltIns.FQ_NAMES.annotation.toUnsafe())
+
+class FirImplicitBooleanTypeRef(
+    session: FirSession,
+    psi: PsiElement?
+) : FirImplicitBuiltinTypeRef(session, psi, KotlinBuiltIns.FQ_NAMES._boolean)