From 2cf8f75a907ff964048413cf0795390792d723fb Mon Sep 17 00:00:00 2001 From: Simon Ogorodnik Date: Wed, 19 Oct 2022 04:29:10 +0200 Subject: [PATCH] KT-53255. Fix StackOverflow during IR verification from K2 In FIR we desugar when with multiple conditions leading to same block as tree of OR expressions Given ``` when(some) { "a", "b", "c" -> {} else -> {} } ``` actually desugared into ``` when(val = some) { == "a" || == "b" || == "c" -> {} else -> {} } ``` There is a multiple ways of how we can organize such expressions in FIR Previously it was just nesting-chain of OR expressions While the most efficient way in terms of required stack depth is a balanced tree KT-53255 --- .../kotlin/fir/lightTree/fir/WhenEntry.kt | 17 ++----- .../kotlin/fir/builder/PsiConversionUtils.kt | 17 +++---- .../kotlin/fir/builder/ConversionUtils.kt | 23 ++++++++- .../codegen/box/when/longCondition.kt | 1 - .../ir/irText/expressions/when.fir.ir.txt | 48 +++++++++---------- .../ir/irText/expressions/when.fir.kt.txt | 22 ++++----- 6 files changed, 67 insertions(+), 61 deletions(-) diff --git a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/fir/WhenEntry.kt b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/fir/WhenEntry.kt index b5fae30afd1..a1c194c6d7e 100644 --- a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/fir/WhenEntry.kt +++ b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/fir/WhenEntry.kt @@ -6,7 +6,7 @@ package org.jetbrains.kotlin.fir.lightTree.fir import com.intellij.lang.LighterASTNode -import org.jetbrains.kotlin.fir.builder.generateLazyLogicalOperation +import org.jetbrains.kotlin.fir.builder.buildBalancedOrExpressionTree import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic import org.jetbrains.kotlin.fir.expressions.FirBlock @@ -20,19 +20,8 @@ data class WhenEntry( val isElse: Boolean = false ) { fun toFirWhenCondition(): FirExpression { - var firCondition: FirExpression? = null - for (condition in conditions) { - val firConditionElement = condition.toFirWhenCondition() - firCondition = when (firCondition) { - null -> firConditionElement - else -> firCondition.generateLazyLogicalOperation(firConditionElement, false, null) - } - } - return firCondition!! - } - - private fun FirExpression.toFirWhenCondition(): FirExpression { - return this + require(conditions.isNotEmpty()) + return buildBalancedOrExpressionTree(conditions) } fun toFirWhenConditionWithoutSubject(): FirExpression { diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiConversionUtils.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiConversionUtils.kt index acde1028bfc..be6147c605d 100644 --- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiConversionUtils.kt +++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiConversionUtils.kt @@ -72,17 +72,14 @@ internal fun Array.toFirWhenCondition( convert: KtExpression?.(String) -> FirExpression, toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef, ): FirExpression { - var firCondition: FirExpression? = null - for (condition in this) { - val firConditionElement = condition.toFirWhenCondition(subject, convert, toFirOrErrorTypeRef) - firCondition = when (firCondition) { - null -> firConditionElement - else -> firCondition.generateLazyLogicalOperation( - firConditionElement, false, condition.toKtPsiSourceElement(KtFakeSourceElementKind.WhenCondition), - ) - } + val conditions = this.map { condition -> + condition.toFirWhenCondition(subject, convert, toFirOrErrorTypeRef) } - return firCondition!! + + require(conditions.isNotEmpty()) + // We build balanced tree of OR expressions to ensure we won't run out of stack + // while processing huge conditions + return buildBalancedOrExpressionTree(conditions) } internal fun generateTemporaryVariable( diff --git a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt index 0580ac0f20a..0c695a23b2c 100644 --- a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt +++ b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt @@ -36,7 +36,6 @@ import org.jetbrains.kotlin.fir.symbols.constructStarProjectedType import org.jetbrains.kotlin.fir.symbols.impl.* import org.jetbrains.kotlin.fir.types.ConeClassLikeType import org.jetbrains.kotlin.fir.types.ConeStarProjection -import org.jetbrains.kotlin.fir.types.FirTypeRef import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef import org.jetbrains.kotlin.fir.types.impl.* @@ -588,3 +587,25 @@ data class CalleeAndReceiver( val receiverExpression: FirExpression? = null, val isImplicitInvoke: Boolean = false ) + +/** + * Creates balanced tree of OR expressions for given set of conditions + * We do so, to avoid too deep OR-expression structures, that can cause running out of stack while processing + * [conditions] should contain at least one element, otherwise it will cause StackOverflow + */ +fun buildBalancedOrExpressionTree(conditions: List, lower: Int = 0, upper: Int = conditions.lastIndex): FirExpression { + val size = upper - lower + 1 + val middle = size / 2 + lower + + if (lower == upper) { + return conditions[middle] + } + val leftNode = buildBalancedOrExpressionTree(conditions, lower, middle - 1) + val rightNode = buildBalancedOrExpressionTree(conditions, middle, upper) + + return leftNode.generateLazyLogicalOperation( + rightNode, + isAnd = false, + (leftNode.source ?: rightNode.source)?.fakeElement(KtFakeSourceElementKind.WhenCondition) + ) +} \ No newline at end of file diff --git a/compiler/testData/codegen/box/when/longCondition.kt b/compiler/testData/codegen/box/when/longCondition.kt index 69bc03339a7..ee3d81946a8 100644 --- a/compiler/testData/codegen/box/when/longCondition.kt +++ b/compiler/testData/codegen/box/when/longCondition.kt @@ -1,5 +1,4 @@ // TARGET_BACKEND: JVM_IR -// IGNORE_BACKEND_FIR: JVM_IR fun box(): String { return when ("foo") { diff --git a/compiler/testData/ir/irText/expressions/when.fir.ir.txt b/compiler/testData/ir/irText/expressions/when.fir.ir.txt index 7b3853e0656..8d91b77e736 100644 --- a/compiler/testData/ir/irText/expressions/when.fir.ir.txt +++ b/compiler/testData/ir/irText/expressions/when.fir.ir.txt @@ -102,50 +102,50 @@ FILE fqName: fileName:/when.kt BRANCH if: WHEN type=kotlin.Boolean origin=OROR BRANCH - if: WHEN type=kotlin.Boolean origin=OROR - BRANCH - if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ - arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null - arg1: CONST Int type=kotlin.Int value=1 - then: CONST Boolean type=kotlin.Boolean value=true - BRANCH - if: CONST Boolean type=kotlin.Boolean value=true - then: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ - arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null - arg1: CONST Int type=kotlin.Int value=2 + if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ + arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null + arg1: CONST Int type=kotlin.Int value=1 then: CONST Boolean type=kotlin.Boolean value=true BRANCH if: CONST Boolean type=kotlin.Boolean value=true then: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null - arg1: CONST Int type=kotlin.Int value=3 + arg1: CONST Int type=kotlin.Int value=2 then: CONST Boolean type=kotlin.Boolean value=true BRANCH if: CONST Boolean type=kotlin.Boolean value=true - then: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ - arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null - arg1: CONST Int type=kotlin.Int value=4 + then: WHEN type=kotlin.Boolean origin=OROR + BRANCH + if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ + arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null + arg1: CONST Int type=kotlin.Int value=3 + then: CONST Boolean type=kotlin.Boolean value=true + BRANCH + if: CONST Boolean type=kotlin.Boolean value=true + then: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ + arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null + arg1: CONST Int type=kotlin.Int value=4 then: CONST String type=kotlin.String value="1234" BRANCH if: WHEN type=kotlin.Boolean origin=OROR BRANCH - if: WHEN type=kotlin.Boolean origin=OROR + if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ + arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null + arg1: CONST Int type=kotlin.Int value=5 + then: CONST Boolean type=kotlin.Boolean value=true + BRANCH + if: CONST Boolean type=kotlin.Boolean value=true + then: WHEN type=kotlin.Boolean origin=OROR BRANCH if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null - arg1: CONST Int type=kotlin.Int value=5 + arg1: CONST Int type=kotlin.Int value=6 then: CONST Boolean type=kotlin.Boolean value=true BRANCH if: CONST Boolean type=kotlin.Boolean value=true then: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null - arg1: CONST Int type=kotlin.Int value=6 - then: CONST Boolean type=kotlin.Boolean value=true - BRANCH - if: CONST Boolean type=kotlin.Boolean value=true - then: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ - arg0: GET_VAR 'val tmp_1: kotlin.Int [val] declared in .testComma' type=kotlin.Int origin=null - arg1: CONST Int type=kotlin.Int value=7 + arg1: CONST Int type=kotlin.Int value=7 then: CONST String type=kotlin.String value="567" BRANCH if: WHEN type=kotlin.Boolean origin=OROR diff --git a/compiler/testData/ir/irText/expressions/when.fir.kt.txt b/compiler/testData/ir/irText/expressions/when.fir.kt.txt index e2c5261800d..76ec1fb96c5 100644 --- a/compiler/testData/ir/irText/expressions/when.fir.kt.txt +++ b/compiler/testData/ir/irText/expressions/when.fir.kt.txt @@ -38,20 +38,20 @@ fun testComma(x: Int): String { when { when { when { - when { - EQEQ(arg0 = tmp1_subject, arg1 = 1) -> true - else -> EQEQ(arg0 = tmp1_subject, arg1 = 2) - } -> true - else -> EQEQ(arg0 = tmp1_subject, arg1 = 3) + EQEQ(arg0 = tmp1_subject, arg1 = 1) -> true + else -> EQEQ(arg0 = tmp1_subject, arg1 = 2) } -> true - else -> EQEQ(arg0 = tmp1_subject, arg1 = 4) + else -> when { + EQEQ(arg0 = tmp1_subject, arg1 = 3) -> true + else -> EQEQ(arg0 = tmp1_subject, arg1 = 4) + } } -> "1234" when { - when { - EQEQ(arg0 = tmp1_subject, arg1 = 5) -> true - else -> EQEQ(arg0 = tmp1_subject, arg1 = 6) - } -> true - else -> EQEQ(arg0 = tmp1_subject, arg1 = 7) + EQEQ(arg0 = tmp1_subject, arg1 = 5) -> true + else -> when { + EQEQ(arg0 = tmp1_subject, arg1 = 6) -> true + else -> EQEQ(arg0 = tmp1_subject, arg1 = 7) + } } -> "567" when { EQEQ(arg0 = tmp1_subject, arg1 = 8) -> true