From 7048184acfd93b13b52c36e9c2407d55a3a8337f Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Mon, 19 Feb 2024 16:14:13 +0100
Subject: [PATCH] [LL FIR] LLFirAnnotationArgumentsLazyResolver: add missed
 checks for context receivers

---
 .../LLFirAnnotationArgumentsLazyResolver.kt         | 13 ++++++-------
 .../low/level/api/fir/util/firCheckResolvedUtils.kt |  6 ++++++
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt
index 35944c37494..1acc3f6d335 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt
@@ -43,9 +43,7 @@ internal object LLFirAnnotationArgumentsLazyResolver : LLFirLazyResolver(FirReso
                     checkAnnotationsAreResolved(target, receiverParameter.typeRef)
                 }
 
-                for (contextReceiver in target.contextReceivers) {
-                    checkAnnotationsAreResolved(target, contextReceiver.typeRef)
-                }
+                checkAnnotationsAreResolved(target.contextReceivers, target)
             }
 
             is FirTypeParameter -> {
@@ -54,15 +52,16 @@ internal object LLFirAnnotationArgumentsLazyResolver : LLFirLazyResolver(FirReso
                 }
             }
 
-            is FirClass -> {
+            is FirRegularClass -> {
                 for (typeRef in target.superTypeRefs) {
                     checkAnnotationsAreResolved(target, typeRef)
                 }
+
+                checkAnnotationsAreResolved(target.contextReceivers, target)
             }
 
-            is FirTypeAlias -> {
-                checkAnnotationsAreResolved(target, target.expandedTypeRef)
-            }
+            is FirScript -> checkAnnotationsAreResolved(target.contextReceivers, target)
+            is FirTypeAlias -> checkAnnotationsAreResolved(target, target.expandedTypeRef)
         }
     }
 }
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/firCheckResolvedUtils.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/firCheckResolvedUtils.kt
index 24fed335e12..b80a5071db3 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/firCheckResolvedUtils.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/firCheckResolvedUtils.kt
@@ -214,6 +214,12 @@ internal fun checkAnnotationsAreResolved(owner: FirAnnotationContainer, typeRef:
     typeRef.accept(AnnotationChecker, owner)
 }
 
+internal fun checkAnnotationsAreResolved(contextReceivers: List<FirContextReceiver>, owner: FirDeclaration) {
+    for (contextReceiver in contextReceivers) {
+        checkAnnotationsAreResolved(owner, contextReceiver.typeRef)
+    }
+}
+
 internal fun FirAbstractBodyResolveTransformerDispatcher.checkAnnotationCallIsResolved(
     symbol: FirBasedSymbol<*>,
     annotationCall: FirAnnotationCall,