Implement caching module annotations per class id

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/modules/CliJavaModuleResolver.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModule
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
+import java.util.concurrent.ConcurrentHashMap
 
 class CliJavaModuleResolver(
     private val moduleGraph: JavaModuleGraph,
@@ -44,8 +45,9 @@ class CliJavaModuleResolver(
 
     private val virtualFileFinder by lazy { VirtualFileFinder.getInstance(project) }
 
-    override fun getModuleAnnotations(classId: ClassId): List<JavaAnnotation>? {
+    override fun getAnnotationsForModuleOwnerOfClass(classId: ClassId): List<JavaAnnotation>? {
         val virtualFile = virtualFileFinder.findSourceOrBinaryVirtualFile(classId) ?: return null
+
         return (findJavaModule(virtualFile) as? JavaModule.Explicit)?.moduleInfo?.annotations
     }
 
@@ -89,4 +91,8 @@ class CliJavaModuleResolver(
 
         return null
     }
+
+    companion object {
+        private const val MODULE_ANNOTATIONS_CACHE_SIZE = 10000
+    }
 }

core/descriptors.jvm/src/org/jetbrains/kotlin/load/java/lazy/JavaModuleAnnotationsProvider.kt

@@ -9,5 +9,5 @@ import org.jetbrains.kotlin.load.java.structure.JavaAnnotation
 import org.jetbrains.kotlin.name.ClassId
 
 interface JavaModuleAnnotationsProvider {
-    fun getModuleAnnotations(classId: ClassId): List<JavaAnnotation>?
+    fun getAnnotationsForModuleOwnerOfClass(classId: ClassId): List<JavaAnnotation>?
 }

core/descriptors.jvm/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaClassDescriptor.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.load.java.toDescriptorVisibility
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.isValidJavaFqName
 import org.jetbrains.kotlin.resolve.constants.StringValue
+import org.jetbrains.kotlin.resolve.descriptorUtil.classId
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.resolveTopLevelClass
@@ -67,6 +68,10 @@ class LazyJavaClassDescriptor(
         }
     }
 
+    val moduleAnnotations by lazy {
+        classId?.let { outerContext.components.javaModuleResolver.getAnnotationsForModuleOwnerOfClass(it) }
+    }
+
     private val kind = when {
         jClass.isAnnotationType -> ClassKind.ANNOTATION_CLASS
         jClass.isInterface -> ClassKind.INTERFACE

core/descriptors.jvm/src/org/jetbrains/kotlin/load/java/typeEnhancement/signatureEnhancement.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.load.java.descriptors.*
 import org.jetbrains.kotlin.load.java.lazy.LazyJavaResolverContext
 import org.jetbrains.kotlin.load.java.lazy.copyWithNewDefaultTypeQualifiers
 import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaAnnotationDescriptor
+import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaClassDescriptor
 import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaTypeParameterDescriptor
 import org.jetbrains.kotlin.load.java.lazy.descriptors.isJavaField
 import org.jetbrains.kotlin.load.kotlin.SignatureBuildingComponents
@@ -150,8 +151,7 @@ class SignatureEnhancement(
 
     private fun <D : CallableMemberDescriptor> D.getDefaultAnnotations(c: LazyJavaResolverContext): Annotations {
         val topLevelClassifier = getTopLevelContainingClassifier() ?: return annotations
-        val classId = topLevelClassifier.classId ?: return annotations
+        val moduleAnnotations = (topLevelClassifier as? LazyJavaClassDescriptor)?.moduleAnnotations
-        val moduleAnnotations = c.components.javaModuleResolver.getModuleAnnotations(classId)
 
         if (moduleAnnotations.isNullOrEmpty()) return annotations
 

core/descriptors.runtime/src/org/jetbrains/kotlin/descriptors/runtime/components/RuntimeModuleData.kt

@@ -123,7 +123,7 @@ fun makeLazyJavaPackageFragmentFromClassLoaderProvider(
         SignatureEnhancement(annotationTypeQualifierResolver, JavaTypeEnhancementState.DISABLED_JSR_305, JavaTypeEnhancement(JavaResolverSettings.Default)),
         JavaClassesTracker.Default, JavaResolverSettings.Default, NewKotlinTypeChecker.Default, javaTypeEnhancementState,
         object : JavaModuleAnnotationsProvider {
-            override fun getModuleAnnotations(classId: ClassId): List<JavaAnnotation>? = null
+            override fun getAnnotationsForModuleOwnerOfClass(classId: ClassId): List<JavaAnnotation>? = null
         }
     )
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/modules/IdeJavaModuleResolver.kt

@@ -16,16 +16,27 @@ import org.jetbrains.kotlin.load.java.structure.impl.convert
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinder
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.resolve.jvm.modules.JavaModule
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
+import java.util.concurrent.ConcurrentHashMap
 
 class IdeJavaModuleResolver(private val project: Project) : JavaModuleResolver {
     private val virtualFileFinder by lazy { VirtualFileFinder.getInstance(project) }
 
-    override fun getModuleAnnotations(classId: ClassId): List<JavaAnnotation>? {
+    private val modulesAnnotationCache = ConcurrentHashMap<ClassId, List<JavaAnnotation>>()
-        val virtualFile = virtualFileFinder.findVirtualFileWithHeader(classId) ?: return null
 
-        return findJavaModule(virtualFile)?.annotations?.convert(::JavaAnnotationImpl)
+    override fun getAnnotationsForModuleOwnerOfClass(classId: ClassId): List<JavaAnnotation>? {
+        if (modulesAnnotationCache.containsKey(classId)) {
+            return modulesAnnotationCache[classId]
+        }
+
+        val virtualFile = virtualFileFinder.findVirtualFileWithHeader(classId) ?: return null
+        val moduleAnnotations = findJavaModule(virtualFile)?.annotations?.convert(::JavaAnnotationImpl)
+
+        if (moduleAnnotations != null && moduleAnnotations.size < MODULE_ANNOTATIONS_CACHE_SIZE) {
+            modulesAnnotationCache[classId] = moduleAnnotations
+        }
+
+        return moduleAnnotations
     }
 
     private fun findJavaModule(file: VirtualFile): PsiJavaModule? = ModuleHighlightUtil2.getModuleDescriptor(file, project)
@@ -65,4 +76,8 @@ class IdeJavaModuleResolver(private val project: Project) : JavaModuleResolver {
     // Returns whether or not [source] exports [packageName] to [target]
     private fun exports(source: PsiJavaModule, packageName: String, target: PsiJavaModule): Boolean =
         source is LightJavaModule || JavaModuleGraphUtil.exports(source, packageName, target)
+
+    companion object {
+        private const val MODULE_ANNOTATIONS_CACHE_SIZE = 10000
+    }
 }