From 1239a8629eceab979481179a0f10c6f61976bcb4 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Thu, 24 Dec 2020 13:32:07 +0300
Subject: [PATCH] FIR Java: make type annotations lazy

---
 .../jetbrains/kotlin/fir/java/JavaUtils.kt    |  7 ++-
 .../kotlin/fir/types/jvm/FirJavaTypeRef.kt    | 60 ++++++++++++-------
 .../fir/types/impl/FirUserTypeRefImpl.kt      |  2 +-
 3 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt b/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt
index 4d2a66e9394..bfcb65d1548 100644
--- a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt
+++ b/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt
@@ -92,9 +92,12 @@ internal fun FirTypeRef.toConeKotlinTypeProbablyFlexible(
     }
 
 internal fun JavaType.toFirJavaTypeRef(session: FirSession, javaTypeParameterStack: JavaTypeParameterStack): FirJavaTypeRef {
-    val annotations = (this as? JavaClassifierType)?.annotations.orEmpty()
     return buildJavaTypeRef {
-        annotations.mapTo(this.annotations) { it.toFirAnnotationCall(session, javaTypeParameterStack) }
+        annotationBuilder = {
+            (this@toFirJavaTypeRef as? JavaClassifierType)?.annotations.orEmpty().map {
+                it.toFirAnnotationCall(session, javaTypeParameterStack)
+            }
+        }
         type = this@toFirJavaTypeRef
     }
 }
diff --git a/compiler/fir/jvm/src/org/jetbrains/kotlin/fir/types/jvm/FirJavaTypeRef.kt b/compiler/fir/jvm/src/org/jetbrains/kotlin/fir/types/jvm/FirJavaTypeRef.kt
index 2d8b05fa664..edf658289f1 100644
--- a/compiler/fir/jvm/src/org/jetbrains/kotlin/fir/types/jvm/FirJavaTypeRef.kt
+++ b/compiler/fir/jvm/src/org/jetbrains/kotlin/fir/types/jvm/FirJavaTypeRef.kt
@@ -5,39 +5,59 @@
 
 package org.jetbrains.kotlin.fir.types.jvm
 
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.types.FirQualifierPart
-import org.jetbrains.kotlin.fir.types.builder.FirUserTypeRefBuilder
-import org.jetbrains.kotlin.fir.types.impl.FirUserTypeRefImpl
+import org.jetbrains.kotlin.fir.types.FirTypeProjection
+import org.jetbrains.kotlin.fir.types.FirUserTypeRef
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.fir.visitors.transformInplace
 import org.jetbrains.kotlin.load.java.structure.JavaType
 
 class FirJavaTypeRef(
     val type: JavaType,
-    annotations: MutableList<FirAnnotationCall>,
-    qualifier: MutableList<FirQualifierPart>
-) : FirUserTypeRefImpl(
-    source = null,
-    isMarkedNullable = false,
-    qualifier,
-    annotations
-)
+    annotationBuilder: () -> List<FirAnnotationCall>,
+    override val qualifier: MutableList<FirQualifierPart>
+) : FirUserTypeRef(), FirAnnotationContainer {
+    override val isMarkedNullable: Boolean
+        get() = false
 
-@FirBuilderDsl
-class FirJavaTypeRefBuilder : FirUserTypeRefBuilder() {
-    override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
-    lateinit var type: JavaType
+    override val source: FirSourceElement?
+        get() = null
 
-    override fun build(): FirJavaTypeRef {
-        return FirJavaTypeRef(type, annotations, qualifier)
+    override val annotations: List<FirAnnotationCall> by lazy { annotationBuilder() }
+
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        for (part in qualifier) {
+            part.typeArgumentList.typeArguments.forEach { it.accept(visitor, data) }
+        }
+        annotations.forEach { it.accept(visitor, data) }
     }
 
-    @Deprecated("Modification of 'source' has no impact for FirJavaTypeRefBuilder", level = DeprecationLevel.HIDDEN)
-    override var source: FirSourceElement? = null
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirUserTypeRef {
+        for (part in qualifier) {
+            (part.typeArgumentList.typeArguments as MutableList<FirTypeProjection>).transformInplace(transformer, data)
+        }
+        return this
+    }
 
-    @Deprecated("Modification of 'isMarkedNullable' has no impact for FirJavaTypeRefBuilder", level = DeprecationLevel.HIDDEN)
-    override var isMarkedNullable: Boolean = false
+    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirUserTypeRef {
+        return this
+    }
+}
+
+@FirBuilderDsl
+class FirJavaTypeRefBuilder {
+    lateinit var annotationBuilder: () -> List<FirAnnotationCall>
+    lateinit var type: JavaType
+    val qualifier: MutableList<FirQualifierPart> = mutableListOf()
+
+    fun build(): FirJavaTypeRef {
+        return FirJavaTypeRef(type, annotationBuilder, qualifier)
+    }
 }
 
 inline fun buildJavaTypeRef(init: FirJavaTypeRefBuilder.() -> Unit): FirJavaTypeRef {
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirUserTypeRefImpl.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirUserTypeRefImpl.kt
index 3479d2782ad..7f6740bcfc7 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirUserTypeRefImpl.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirUserTypeRefImpl.kt
@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.fir.visitors.transformInplace
 
-open class FirUserTypeRefImpl(
+class FirUserTypeRefImpl(
     override val source: FirSourceElement?,
     override val isMarkedNullable: Boolean,
     override val qualifier: MutableList<FirQualifierPart>,