Fix memory leak in KotlinBuiltins related to unsigned types

Having a map from ModuleDescriptor leads to modules leakage
Especially, it's critical for DefaultBuiltins which is used in JS

ModuleDescriptor's instances were leaked there forever until the daemon dies
(up to 350M while compiling Kotlin project itself)

core/descriptors/src/org/jetbrains/kotlin/builtins/KotlinBuiltIns.java

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.FqNameUnsafe;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
+import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
 import org.jetbrains.kotlin.resolve.scopes.ChainedMemberScope;
 import org.jetbrains.kotlin.resolve.scopes.MemberScope;
 import org.jetbrains.kotlin.storage.MemoizedFunctionToNotNull;
@@ -64,7 +65,6 @@ public abstract class KotlinBuiltIns {
     private ModuleDescriptorImpl builtInsModule;
 
     private final NotNullLazyValue<Primitives> primitives;
-    private final MemoizedFunctionToNotNull<ModuleDescriptor, UnsignedPrimitives> unsignedPrimitives;
     private final NotNullLazyValue<PackageFragments> packageFragments;
 
     private final MemoizedFunctionToNotNull<Name, ClassDescriptor> builtInClassesByName;
@@ -114,29 +114,6 @@ public abstract class KotlinBuiltIns {
             }
         });
 
-        this.unsignedPrimitives = storageManager.createMemoizedFunction(new Function1<ModuleDescriptor, UnsignedPrimitives>() {
-            @Override
-            public UnsignedPrimitives invoke(ModuleDescriptor module) {
-                Map<KotlinType, SimpleType> unsignedKotlinTypeToKotlinArrayType = new HashMap<KotlinType, SimpleType>();
-                Map<SimpleType, SimpleType> kotlinArrayTypeToUnsignedKotlinType = new HashMap<SimpleType, SimpleType>();
-                for (UnsignedType unsigned : UnsignedType.values()) {
-                    ClassDescriptor descriptor = FindClassInModuleKt.findClassAcrossModuleDependencies(module, unsigned.getClassId());
-                    if (descriptor == null) continue;
-
-                    ClassDescriptor arrayDescriptor =
-                            FindClassInModuleKt.findClassAcrossModuleDependencies(module, unsigned.getArrayClassId());
-                    if (arrayDescriptor == null) continue;
-
-                    SimpleType type = descriptor.getDefaultType();
-                    SimpleType arrayType = arrayDescriptor.getDefaultType();
-
-                    unsignedKotlinTypeToKotlinArrayType.put(type, arrayType);
-                    kotlinArrayTypeToUnsignedKotlinType.put(arrayType, type);
-                }
-                return new UnsignedPrimitives(unsignedKotlinTypeToKotlinArrayType, kotlinArrayTypeToUnsignedKotlinType);
-            }
-        });
-
         this.builtInClassesByName = storageManager.createMemoizedFunction(new Function1<Name, ClassDescriptor>() {
             @Override
             public ClassDescriptor invoke(Name name) {
@@ -824,8 +801,7 @@ public abstract class KotlinBuiltIns {
 
         ModuleDescriptor module = DescriptorUtils.getContainingModuleOrNull(notNullArrayType);
         if (module != null) {
-            //noinspection SuspiciousMethodCalls
-            SimpleType unsignedType = unsignedPrimitives.invoke(module).kotlinArrayTypeToUnsignedKotlinType.get(notNullArrayType);
+            KotlinType unsignedType = getElementTypeForUnsignedArray(notNullArrayType, module);
             if (unsignedType != null) return unsignedType;
         }
 
@@ -833,6 +809,24 @@ public abstract class KotlinBuiltIns {
         throw new IllegalStateException("not array: " + arrayType);
     }
 
+    @Nullable
+    private static KotlinType getElementTypeForUnsignedArray(@NotNull KotlinType notNullArrayType, @NotNull ModuleDescriptor module) {
+        ClassifierDescriptor descriptor = notNullArrayType.getConstructor().getDeclarationDescriptor();
+        if (descriptor == null) return null;
+        if (!UnsignedTypes.INSTANCE.isShortNameOfUnsignedArray(descriptor.getName())) return null;
+
+        ClassId arrayClassId = DescriptorUtilsKt.getClassId(descriptor);
+        if (arrayClassId == null) return null;
+
+        ClassId elementClassId = UnsignedTypes.INSTANCE.getUnsignedClassIdByArrayClassId(arrayClassId);
+        if (elementClassId == null) return null;
+
+        ClassDescriptor elementClassDescriptor = FindClassInModuleKt.findClassAcrossModuleDependencies(module, elementClassId);
+        if (elementClassDescriptor == null) return null;
+
+        return elementClassDescriptor.getDefaultType();
+    }
+
     @NotNull
     public SimpleType getPrimitiveArrayKotlinType(@NotNull PrimitiveType primitiveType) {
         return primitives.invoke().primitiveTypeToArrayKotlinType.get(primitiveType);
@@ -847,10 +841,21 @@ public abstract class KotlinBuiltIns {
         if (primitiveArray != null) return primitiveArray;
 
         if (UnsignedTypes.INSTANCE.isUnsignedType(kotlinType)) {
+            if (TypeUtils.isNullableType(kotlinType)) return null;
+
             ModuleDescriptor module = DescriptorUtils.getContainingModuleOrNull(kotlinType);
             if (module == null) return null;
 
-            return unsignedPrimitives.invoke(module).unsignedKotlinTypeToKotlinArrayType.get(kotlinType);
+            ClassId unsignedClassId = DescriptorUtilsKt.getClassId(kotlinType.getConstructor().getDeclarationDescriptor());
+            assert unsignedClassId != null : "unsignedClassId should not be null for unsigned type " + kotlinType;
+
+            ClassId arrayClassId = UnsignedTypes.INSTANCE.getUnsignedArrayClassIdByUnsignedClassId(unsignedClassId);
+            assert arrayClassId != null : "arrayClassId should not be null for unsigned type " + unsignedClassId;
+
+            ClassDescriptor arrayClassDescriptor = FindClassInModuleKt.findClassAcrossModuleDependencies(module, arrayClassId);
+            if (arrayClassDescriptor == null) return null;
+
+            return arrayClassDescriptor.getDefaultType();
         }
 
         return null;

core/descriptors/src/org/jetbrains/kotlin/builtins/UnsignedType.kt

@@ -20,12 +20,27 @@ enum class UnsignedType(val classId: ClassId) {
     ULONG(ClassId.fromString("kotlin/ULong"));
 
     val typeName = classId.shortClassName
-    val arrayTypeName = Name.identifier(typeName.asString() + "Array")
-    val arrayClassId = ClassId(classId.packageFqName, arrayTypeName)
+    val arrayClassId = ClassId(classId.packageFqName, Name.identifier(typeName.asString() + "Array"))
 }
 
 object UnsignedTypes {
-    val unsignedTypeNames = enumValues<UnsignedType>().map { it.typeName }.toSet()
+    private val unsignedTypeNames = enumValues<UnsignedType>().map { it.typeName }.toSet()
+    private val arrayClassIdToUnsignedClassId = hashMapOf<ClassId, ClassId>()
+    private val unsignedClassIdToArrayClassId = hashMapOf<ClassId, ClassId>()
+
+    private val arrayClassesShortNames: Set<Name> = UnsignedType.values().mapTo(mutableSetOf()) { it.arrayClassId.shortClassName }
+
+    init {
+        for (unsignedType in UnsignedType.values()) {
+            arrayClassIdToUnsignedClassId[unsignedType.arrayClassId] = unsignedType.classId
+            unsignedClassIdToArrayClassId[unsignedType.classId] = unsignedType.arrayClassId
+        }
+    }
+
+    fun isShortNameOfUnsignedArray(name: Name) = name in arrayClassesShortNames
+
+    fun getUnsignedClassIdByArrayClassId(arrayClassId: ClassId): ClassId? = arrayClassIdToUnsignedClassId[arrayClassId]
+    fun getUnsignedArrayClassIdByUnsignedClassId(arrayClassId: ClassId): ClassId? = unsignedClassIdToArrayClassId[arrayClassId]
 
     fun isUnsignedType(type: KotlinType): Boolean {
         if (TypeUtils.noExpectedType(type)) return false
@@ -40,4 +55,4 @@ object UnsignedTypes {
                 container.fqName == KotlinBuiltIns.BUILT_INS_PACKAGE_FQ_NAME &&
                 descriptor.name in UnsignedTypes.unsignedTypeNames
     }
-}
+}

core/descriptors/src/org/jetbrains/kotlin/resolve/DescriptorUtils.kt

@@ -72,8 +72,8 @@ val ClassifierDescriptorWithTypeParameters.denotedClassDescriptor: ClassDescript
         else -> throw UnsupportedOperationException("Unexpected descriptor kind: $this")
     }
 
-val ClassifierDescriptorWithTypeParameters.classId: ClassId?
-    get() = containingDeclaration.let { owner ->
+val ClassifierDescriptor?.classId: ClassId?
+    get() = this?.containingDeclaration?.let { owner ->
         when (owner) {
             is PackageFragmentDescriptor -> ClassId(owner.fqName, name)
             is ClassifierDescriptorWithTypeParameters -> owner.classId?.createNestedClassId(name)