From 2cb6ee81a13114a9eb3fee5b7de4cbdb842cae52 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Wed, 23 Aug 2023 18:05:34 +0200
Subject: [PATCH] [LL FIR] add guards to type resolver

We should protect return types from value parameters/property accessors
and backing fields from repeatable modifications.
Such declarations can be modified during type resolution by
`transformTypeToArrayType` or property getter type propagation,
so we should be able to restore the original reference in the case of
exceptions (PCE for instance)

^KT-61422
---
 .../fir/transformers/LLFirTypeLazyResolver.kt | 51 +++++++++++++++++--
 1 file changed, 46 insertions(+), 5 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt
index 223b5c6b44e..20d9a758c0f 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt
@@ -103,6 +103,31 @@ private class LLFirTypeTargetResolver(
     }
 
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
+        when (target) {
+            is FirFunction -> resolve(target, TypeStateKeepers.FUNCTION)
+            is FirProperty -> resolve(target, TypeStateKeepers.PROPERTY)
+            is FirCallableDeclaration,
+            is FirDanglingModifierList,
+            is FirFileAnnotationsContainer,
+            is FirTypeAlias,
+            is FirScript,
+            is FirRegularClass,
+            -> rawResolve(target)
+
+            is FirFile, is FirAnonymousInitializer, is FirCodeFragment -> {}
+            else -> errorWithAttachment("Unknown declaration ${target::class.simpleName}") {
+                withFirEntry("declaration", target)
+            }
+        }
+    }
+
+    private fun <T : FirElementWithResolveState> resolve(target: T, keeper: StateKeeper<T, Unit>) {
+        resolveWithKeeper(target, Unit, keeper) {
+            rawResolve(target)
+        }
+    }
+
+    private fun rawResolve(target: FirElementWithResolveState) {
         when (target) {
             is FirConstructor -> {
                 // ConstructedTypeRef should be resolved only with type parameters, but not with nested classes and classes from supertypes
@@ -128,11 +153,9 @@ private class LLFirTypeTargetResolver(
             is FirDanglingModifierList, is FirFileAnnotationsContainer, is FirCallableDeclaration, is FirTypeAlias -> {
                 target.accept(transformer, null)
             }
-            is FirRegularClass -> {
-                resolveClassTypes(target)
-            }
-            is FirFile, is FirAnonymousInitializer, is FirCodeFragment -> {}
-            else -> errorWithAttachment("Unknown declaration ${target::class}") {
+
+            is FirRegularClass -> resolveClassTypes(target)
+            else -> errorWithAttachment("Unknown declaration ${target::class.simpleName}") {
                 withFirEntry("declaration", target)
             }
         }
@@ -157,3 +180,21 @@ private class LLFirTypeTargetResolver(
         }
     }
 }
+
+private object TypeStateKeepers {
+    val FUNCTION: StateKeeper<FirFunction, Unit> = stateKeeper { function, context ->
+        add(CALLABLE_DECLARATION, context)
+        entityList(function.valueParameters, CALLABLE_DECLARATION, context)
+    }
+
+    val PROPERTY: StateKeeper<FirProperty, Unit> = stateKeeper { property, context ->
+        add(CALLABLE_DECLARATION, context)
+        entity(property.getter, FUNCTION, context)
+        entity(property.setter, FUNCTION, context)
+        entity(property.backingField, CALLABLE_DECLARATION, context)
+    }
+
+    private val CALLABLE_DECLARATION: StateKeeper<FirCallableDeclaration, Unit> = stateKeeper { _, _ ->
+        add(FirCallableDeclaration::returnTypeRef, FirCallableDeclaration::replaceReturnTypeRef)
+    }
+}