From ee23a52e540935ac52cfeb6105070477a940ed4f Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Wed, 2 Mar 2022 14:56:52 -0800
Subject: [PATCH] AA: avoid conversions to KtInitializerValue in constant
 evaluator

---
 .../KtFirCompileTimeConstantProvider.kt       | 39 ++++++++-----------
 1 file changed, 17 insertions(+), 22 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompileTimeConstantProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompileTimeConstantProvider.kt
index 1e7c658a435..d73803c4753 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompileTimeConstantProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompileTimeConstantProvider.kt
@@ -5,27 +5,24 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.components
 
-import org.jetbrains.kotlin.analysis.api.KtConstantInitializerValue
-import org.jetbrains.kotlin.analysis.api.KtInitializerValue
-import org.jetbrains.kotlin.analysis.api.KtNonConstantInitializerValue
 import org.jetbrains.kotlin.analysis.api.base.KtConstantValue
 import org.jetbrains.kotlin.analysis.api.components.KtCompileTimeConstantProvider
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.evaluate.FirCompileTimeConstantEvaluator
-import org.jetbrains.kotlin.analysis.api.fir.symbols.getKtConstantInitializer
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.utils.referredPropertySymbol
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhenBranch
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.resolvedSymbol
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 
 internal class KtFirCompileTimeConstantProvider(
@@ -34,20 +31,24 @@ internal class KtFirCompileTimeConstantProvider(
 ) : KtCompileTimeConstantProvider(), KtFirAnalysisSessionComponent {
 
     override fun evaluate(expression: KtExpression): KtConstantValue? = withValidityAssertion {
-        when (val fir = expression.getOrBuildFir(firResolveState)) {
+        evaluateFir(expression.getOrBuildFir(firResolveState), expression)
+    }
+
+    private fun evaluateFir(fir: FirElement?, context: KtExpression): KtConstantValue? = withValidityAssertion {
+        when (fir) {
             is FirPropertyAccessExpression -> {
-                fir.referredPropertySymbol?.toKtConstantValue()
+                fir.referredPropertySymbol?.toKtConstantValue(context)
             }
             is FirExpression -> {
                 try {
                     FirCompileTimeConstantEvaluator.evaluateAsKtConstantExpression(fir)
                 } catch (e: ArithmeticException) {
-                    KtConstantValue.KtErrorConstantValue(e.localizedMessage, fir.psi as? KtElement)
+                    KtConstantValue.KtErrorConstantValue(e.localizedMessage, context)
                 }
             }
             is FirNamedReference -> {
                 when (val resolvedSymbol = fir.resolvedSymbol) {
-                    is FirPropertySymbol -> resolvedSymbol.toKtConstantValue()
+                    is FirPropertySymbol -> resolvedSymbol.toKtConstantValue(context)
                     else -> null
                 }
             }
@@ -60,21 +61,15 @@ internal class KtFirCompileTimeConstantProvider(
             // `false` does not have a corresponding elements on the FIR side and hence the containing `FirWhenBranch` is returned. In this
             // case, we simply report null since FIR does not know about it.
             is FirWhenBranch -> null
-            else -> throwUnexpectedFirElementError(fir, expression)
+            else -> throwUnexpectedFirElementError(fir, context)
         }
     }
 
-    private fun FirPropertySymbol.toKtConstantValue(): KtConstantValue? = withValidityAssertion {
-        if (isVal) {
-            getKtConstantInitializer()?.toKtConstantValue()
+    private fun FirPropertySymbol.toKtConstantValue(context: KtExpression): KtConstantValue? = withValidityAssertion {
+        if (isVal && hasInitializer) {
+            // NB: the initializer could be [FirLazyExpression] in [BodyBuildingMode.LAZY_BODIES].
+            this.ensureResolved(FirResolvePhase.BODY_RESOLVE) // to unwrap lazy body
+            evaluateFir(fir.initializer, context)
         } else null
     }
-
-    private fun KtInitializerValue?.toKtConstantValue(): KtConstantValue? = withValidityAssertion {
-        when (this) {
-            null -> null
-            is KtConstantInitializerValue -> constant
-            is KtNonConstantInitializerValue -> initializerPsi?.let { evaluate(it) }
-        }
-    }
 }