Introduce runtime ClassValue-based cache for typeOf machinery

* Cache KType instances constructor from the given classifier
* For generics, cache KTypes with already substituted arguments
* It significantly speeds up all typeOf-based APIs, both accesses to typeOf and its related properties (i.e. classifier)

#KT-53508

Merge-request: KT-MR-6818
Merged-by: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -44330,6 +44330,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/reflection/typeOf/annotatedType.kt");
             }
 
+            @Test
+            @TestMetadata("caching.kt")
+            public void testCaching() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/typeOf/caching.kt");
+            }
+
             @Test
             @TestMetadata("classes.kt")
             public void testClasses() throws Exception {

compiler/testData/codegen/box/reflection/typeOf/caching.kt

@@ -0,0 +1,24 @@
+// TARGET_BACKEND: JVM_IR
+// WITH_REFLECT
+
+import kotlin.test.assertEquals
+import kotlin.test.assertSame
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
+
+private inline fun <reified T> check(isNullable: Boolean = false) {
+    val t1 = typeOf<T>()
+    val t2 = typeOf<T>()
+    assertSame(t1.classifier, t2.classifier)
+    assertEquals(isNullable, t1.isMarkedNullable)
+    assertEquals(isNullable, t2.isMarkedNullable)
+}
+
+fun box(): String {
+    check<Int>()
+    check<Int?>(true)
+    check<List<Int>?>(true)
+    check<List<Int>>()
+    check<List<Int?>?>(true)
+    return "OK"
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -44330,6 +44330,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/reflection/typeOf/annotatedType.kt");
             }
 
+            @Test
+            @TestMetadata("caching.kt")
+            public void testCaching() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/typeOf/caching.kt");
+            }
+
             @Test
             @TestMetadata("classes.kt")
             public void testClasses() throws Exception {

core/reflection.jvm/src/kotlin/reflect/jvm/internal/ClassCaches.kt

@@ -1,23 +0,0 @@
-/*
- * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package kotlin.reflect.jvm.internal
-
-import kotlin.reflect.KDeclarationContainer
-
-
-private val K_CLASS_CACHE = createCache { KClassImpl(it) }
-private val K_PACKAGE_CACHE = createCache { KPackageImpl(it) }
-
-// This function is invoked on each reflection access to Java classes, properties, etc. Performance is critical here.
-@Suppress("UNCHECKED_CAST")
-internal fun <T : Any> getOrCreateKotlinClass(jClass: Class<T>): KClassImpl<T> = K_CLASS_CACHE.get(jClass) as KClassImpl<T>
-
-internal fun <T : Any> getOrCreateKotlinPackage(jClass: Class<T>): KDeclarationContainer = K_PACKAGE_CACHE.get(jClass)
-
-internal fun clearCaches() {
-    K_CLASS_CACHE.clear()
-    K_PACKAGE_CACHE.clear()
-}

core/reflection.jvm/src/kotlin/reflect/jvm/internal/ReflectionFactoryImpl.java

@@ -33,17 +33,17 @@ public class ReflectionFactoryImpl extends ReflectionFactory {
     @Override
     public KDeclarationContainer getOrCreateKotlinPackage(Class javaClass, String moduleName) {
         // moduleName is unused deliberately and only left in public ABI
-        return ClassCachesKt.getOrCreateKotlinPackage(javaClass);
+        return CachesKt.getOrCreateKotlinPackage(javaClass);
     }
 
     @Override
     public KClass getOrCreateKotlinClass(Class javaClass) {
-        return ClassCachesKt.getOrCreateKotlinClass(javaClass);
+        return CachesKt.getOrCreateKotlinClass(javaClass);
     }
 
     @Override
     public KClass getOrCreateKotlinClass(Class javaClass, String internalName) {
-        return ClassCachesKt.getOrCreateKotlinClass(javaClass);
+        return CachesKt.getOrCreateKotlinClass(javaClass);
     }
 
     @Override
@@ -111,6 +111,17 @@ public class ReflectionFactoryImpl extends ReflectionFactory {
 
     @Override
     public KType typeOf(KClassifier klass, List<KTypeProjection> arguments, boolean isMarkedNullable) {
+        /*
+         * We control how this method is called and ensure that `typeOf` is invoked mostly (in scenarios that
+         * bother us performance-wise) on the result of `getOrCreateKotlinClass(klass)`, thus we do downcast
+         * and extract `java.lang.Class` in a zero-cost manner.
+         * If that's our case, we go to caching code-path that caches relatively slow `createType()` call,
+         * and, what's more important, all member-based computations on this KType (e.g. `classifier`)
+         * are properly cached as well.
+         */
+        if (klass instanceof ClassBasedDeclarationContainer) {
+            return CachesKt.getOrCreateKType(((ClassBasedDeclarationContainer) klass).getJClass(), arguments, isMarkedNullable);
+        }
         return KClassifiers.createType(klass, arguments, isMarkedNullable, Collections.<Annotation>emptyList());
     }
 
@@ -155,7 +166,7 @@ public class ReflectionFactoryImpl extends ReflectionFactory {
     // Misc
 
     public static void clearCaches() {
-        ClassCachesKt.clearCaches();
+        CachesKt.clearCaches();
         ModuleByClassLoaderKt.clearModuleByClassLoaderCache();
     }
 }

core/reflection.jvm/src/kotlin/reflect/jvm/internal/caches.kt

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package kotlin.reflect.jvm.internal
+
+import java.util.concurrent.ConcurrentHashMap
+import kotlin.reflect.KType
+import kotlin.reflect.KTypeProjection
+import kotlin.reflect.full.createType
+
+import kotlin.reflect.KDeclarationContainer
+
+
+private val K_CLASS_CACHE = createCache { KClassImpl(it) }
+private val K_PACKAGE_CACHE = createCache { KPackageImpl(it) }
+
+// This function is invoked on each reflection access to Java classes, properties, etc. Performance is critical here.
+@Suppress("UNCHECKED_CAST")
+internal fun <T : Any> getOrCreateKotlinClass(jClass: Class<T>): KClassImpl<T> = K_CLASS_CACHE.get(jClass) as KClassImpl<T>
+
+internal fun <T : Any> getOrCreateKotlinPackage(jClass: Class<T>): KDeclarationContainer = K_PACKAGE_CACHE.get(jClass)
+
+internal fun clearCaches() {
+    K_CLASS_CACHE.clear()
+    K_PACKAGE_CACHE.clear()
+    CACHE_FOR_BASE_CLASSIFIERS.clear()
+    CACHE_FOR_NULLABLE_BASE_CLASSIFIERS.clear()
+    CACHE_FOR_GENERIC_CLASSIFIERS.clear()
+}
+
+// typeOf-related caches
+
+// Without type arguments and nullability
+private val CACHE_FOR_BASE_CLASSIFIERS = createCache {
+    getOrCreateKotlinClass(it).createType(emptyList(), false, emptyList())
+}
+
+private val CACHE_FOR_NULLABLE_BASE_CLASSIFIERS = createCache {
+    getOrCreateKotlinClass(it).createType(emptyList(), true, emptyList())
+}
+
+private typealias Key = Pair<List<KTypeProjection>, Boolean>
+// Class -> ((type arguments, is nullable) -> type)
+private val CACHE_FOR_GENERIC_CLASSIFIERS = createCache<ConcurrentHashMap<Key, KType>> {
+    ConcurrentHashMap()
+}
+
+internal fun <T : Any> getOrCreateKType(jClass: Class<T>, arguments: List<KTypeProjection>, isMarkedNullable: Boolean): KType {
+    return if (arguments.isEmpty()) {
+        if (isMarkedNullable) {
+            CACHE_FOR_NULLABLE_BASE_CLASSIFIERS.get(jClass)
+        } else {
+            CACHE_FOR_BASE_CLASSIFIERS.get(jClass)
+        }
+    } else {
+        getOrCreateKTypeWithTypeArguments(jClass, arguments, isMarkedNullable)
+    }
+}
+
+private fun <T : Any> getOrCreateKTypeWithTypeArguments(
+    jClass: Class<T>,
+    arguments: List<KTypeProjection>,
+    isMarkedNullable: Boolean
+): KType {
+    val cache = CACHE_FOR_GENERIC_CLASSIFIERS.get(jClass)
+    return cache.getOrPut(arguments to isMarkedNullable) {
+        getOrCreateKotlinClass(jClass).createType(arguments, isMarkedNullable, emptyList())
+    }
+}
+
+