Load all resources in RuntimePackagePartProvider; optimize code

The issue was reproducible when the same package is present in different
modules with the same -module-name (which is a popular case of src/test
roots in simple IDEA projects). The problem was in the fact that several
resource files containing package name mapping with the same name were
present in the classpath, but RuntimePackagePartProvider only considered
the first one. The fix is to use getResources instead of
getResourceAsStream and handle each returned resource.

Also, optimize internal representation to store the mapping in the form
which is the most convenient for findPackageParts, which should be
faster than registerModule because in theory, it's called more often.

 #KT-21973 Fixed
 #KT-24651

compiler/backend/src/org/jetbrains/kotlin/codegen/GeneratedClassLoader.java

@@ -16,17 +16,42 @@
 
 package org.jetbrains.kotlin.codegen;
 
+import kotlin.collections.CollectionsKt;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.backend.common.output.OutputFile;
+import org.jetbrains.kotlin.utils.ExceptionUtilsKt;
 
 import java.io.ByteArrayInputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.net.URLConnection;
+import java.net.URLStreamHandler;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.Enumeration;
 import java.util.List;
 import java.util.jar.Manifest;
 
 public class GeneratedClassLoader extends URLClassLoader {
+    private static final URLStreamHandler FAKE_BASE64_URL_HANDLER = new URLStreamHandler() {
+        @Override
+        protected URLConnection openConnection(URL url) {
+            return new URLConnection(url) {
+                @Override
+                public void connect() {
+                }
+
+                @Override
+                public InputStream getInputStream() {
+                    return new ByteArrayInputStream(Base64.getDecoder().decode(url.getPath()));
+                }
+            };
+        }
+    };
+
     private ClassFileFactory factory;
 
     public GeneratedClassLoader(@NotNull ClassFileFactory factory, ClassLoader parentClassLoader, URL... urls) {
@@ -43,6 +68,39 @@ public class GeneratedClassLoader extends URLClassLoader {
         return super.getResourceAsStream(name);
     }
 
+    @Override
+    public Enumeration<URL> findResources(String name) throws IOException {
+        Enumeration<URL> fromParent = super.findResources(name);
+
+        URL url = createFakeURLForResource(name);
+        if (url == null) return fromParent;
+
+        List<URL> fromMe = Collections.singletonList(url);
+        List<URL> result = fromParent.hasMoreElements()
+                           ? CollectionsKt.plus(fromMe, Collections.list(fromParent))
+                           : fromMe;
+        return Collections.enumeration(result);
+    }
+
+    @Override
+    public URL findResource(String name) {
+        URL url = createFakeURLForResource(name);
+        return url != null ? url : super.findResource(name);
+    }
+
+    @Nullable
+    private URL createFakeURLForResource(@NotNull String name) {
+        try {
+            OutputFile outputFile = factory.get(name);
+            // Encode the byte array in the URL path to prevent creating unneeded temporary files
+            return outputFile == null
+                   ? null
+                   : new URL(null, "bytes:" + Base64.getEncoder().encodeToString(outputFile.asByteArray()), FAKE_BASE64_URL_HANDLER);
+        } catch (IOException e) {
+            throw ExceptionUtilsKt.rethrow(e);
+        }
+    }
+
     @NotNull
     @Override
     protected Class<?> findClass(@NotNull String name) throws ClassNotFoundException {

compiler/testData/compileKotlinAgainstKotlin/reflectTopLevelFunctionOtherFile.kt

@@ -0,0 +1,20 @@
+// IGNORE_BACKEND: NATIVE
+
+// WITH_REFLECT
+// FILE: A.kt
+
+fun test() {
+}
+
+// FILE: B.kt
+
+import kotlin.reflect.jvm.javaMethod
+import kotlin.reflect.jvm.kotlinFunction
+
+fun box(): String {
+    val r = Class.forName("AKt").methods.single { it.name == "test" }.kotlinFunction
+    if (r?.toString() != "fun test(): kotlin.Unit")
+        return "Fail: $r"
+
+    return "OK"
+}

compiler/tests/org/jetbrains/kotlin/codegen/CompileKotlinAgainstKotlinTestGenerated.java

@@ -233,6 +233,11 @@ public class CompileKotlinAgainstKotlinTestGenerated extends AbstractCompileKotl
         runTest("compiler/testData/compileKotlinAgainstKotlin/recursiveGeneric.kt");
     }
 
+    @TestMetadata("reflectTopLevelFunctionOtherFile.kt")
+    public void testReflectTopLevelFunctionOtherFile() throws Exception {
+        runTest("compiler/testData/compileKotlinAgainstKotlin/reflectTopLevelFunctionOtherFile.kt");
+    }
+
     @TestMetadata("sealedClass.kt")
     public void testSealedClass() throws Exception {
         runTest("compiler/testData/compileKotlinAgainstKotlin/sealedClass.kt");

core/descriptors.runtime/src/kotlin/reflect/jvm/internal/components/RuntimePackagePartProvider.kt

@@ -21,27 +21,44 @@ import org.jetbrains.kotlin.load.kotlin.loadModuleMapping
 import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.serialization.deserialization.DeserializationConfiguration
-import java.util.concurrent.ConcurrentHashMap
+import java.io.IOException
+import java.util.*
 
 class RuntimePackagePartProvider(private val classLoader: ClassLoader) : PackagePartProvider {
-    private val module2Mapping = ConcurrentHashMap<String, ModuleMapping>()
+    // Names of modules which were registered with registerModule
+    private val visitedModules = hashSetOf<String>()
 
+    // Package FQ name -> list of JVM internal names of package parts in that package across all registered modules
+    private val packageParts = hashMapOf<String, LinkedHashSet<String>>()
+
+    @Synchronized
     fun registerModule(moduleName: String) {
-        val mapping = try {
-            val resourcePath = "META-INF/$moduleName.${ModuleMapping.MAPPING_FILE_EXT}"
-            classLoader.getResourceAsStream(resourcePath)?.use { stream ->
-                ModuleMapping.loadModuleMapping(stream.readBytes(), resourcePath, DeserializationConfiguration.Default)
-            }
-        } catch (e: Exception) {
-            // TODO: do not swallow this exception?
-            null
+        if (!visitedModules.add(moduleName)) return
+
+        val resourcePath = "META-INF/$moduleName.${ModuleMapping.MAPPING_FILE_EXT}"
+        val resources = try {
+            classLoader.getResources(resourcePath)
+        } catch (e: IOException) {
+            EmptyEnumeration
+        }
+
+        for (resource in resources) {
+            try {
+                resource.openStream()?.use { stream ->
+                    val mapping = ModuleMapping.loadModuleMapping(stream.readBytes(), resourcePath, DeserializationConfiguration.Default)
+                    for ((packageFqName, parts) in mapping.packageFqName2Parts) {
+                        packageParts.getOrPut(packageFqName) { linkedSetOf() }.addAll(parts.parts)
+                    }
+                }
+            } catch (e: Exception) {
+                // TODO: do not swallow this exception?
+            }
         }
-        module2Mapping.putIfAbsent(moduleName, mapping ?: ModuleMapping.EMPTY)
     }
 
-    override fun findPackageParts(packageFqName: String): List<String> {
-        return module2Mapping.values.mapNotNull { it.findPackageParts(packageFqName) }.flatMap { it.parts }.distinct()
-    }
+    @Synchronized
+    override fun findPackageParts(packageFqName: String): List<String> =
+        packageParts[packageFqName]?.toList().orEmpty()
 
     // TODO
     override fun findMetadataPackageParts(packageFqName: String): List<String> = TODO()
@@ -50,4 +67,9 @@ class RuntimePackagePartProvider(private val classLoader: ClassLoader) : Package
         // TODO: load annotations from resource files
         return emptyList()
     }
+
+    private object EmptyEnumeration : Enumeration<Nothing> {
+        override fun hasMoreElements(): Boolean = false
+        override fun nextElement(): Nothing = throw NoSuchElementException()
+    }
 }