From 7ba8e7d1d3b12a1866ee79ea5135c6d94cca6940 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Wed, 21 Jul 2021 17:09:50 +0200
Subject: [PATCH] LL API: fix KtPackageProvider package existent checks for
 tests

Before it was not able to check nested packages for existence

Also, implement unimplemented getKotlinSubPackageFqNames
---
 .../test/base/KtPackageProviderTestImpl.kt    | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/test/base/KtPackageProviderTestImpl.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/test/base/KtPackageProviderTestImpl.kt
index 59d9bab569d..48caefb97b5 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/test/base/KtPackageProviderTestImpl.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/test/base/KtPackageProviderTestImpl.kt
@@ -15,13 +15,26 @@ internal class KtPackageProviderTestImpl(
     scope: GlobalSearchScope,
     files: Collection<KtFile>
 ) : KtPackageProvider() {
-    private val filesInScope = files.filter { scope.contains(it.virtualFile) }
+
+    private val packageToSubPackageNames: Map<FqName, Set<Name>> = run {
+        val filesInScope = files.filter { scope.contains(it.virtualFile) }
+        val packages = mutableMapOf<FqName, MutableSet<Name>>()
+        filesInScope.forEach { file ->
+            var currentPackage = FqName.ROOT
+            for (subPackage in file.packageFqName.pathSegments()) {
+                packages.getOrPut(currentPackage) { mutableSetOf() } += subPackage
+                currentPackage = currentPackage.child(subPackage)
+            }
+            packages.computeIfAbsent(currentPackage) { mutableSetOf() }
+        }
+        packages
+    }
 
     override fun isPackageExists(packageFqName: FqName): Boolean {
-        return filesInScope.any { it.packageFqName == packageFqName }
+        return packageFqName in packageToSubPackageNames
     }
 
     override fun getKotlinSubPackageFqNames(packageFqName: FqName): Set<Name> {
-        TODO("Not yet implemented")
+        return packageToSubPackageNames[packageFqName] ?: emptySet()
     }
 }
\ No newline at end of file