IJPL-339 support several contexts for one virtual file in PsiDirectory

GitOrigin-RevId: ba422e49fd17778dc290188b29e1bf3003e476c4

platform/core-api/src/com/intellij/psi/search/CodeInsightContextAwareSearchScopes.kt

@@ -57,6 +57,16 @@ fun SearchScope.getAnyCorrespondingContext(file: VirtualFile): CodeInsightContex
   }
 }
 
+//todo ijpl-339 mark experimental
+@Internal
+fun SearchScope.getCorrespondingContexts(file: VirtualFile): Collection<CodeInsightContext> {
+  return when (val contextInfo = getFileContextInfo(file)) {
+    is ActualContextFileInfo -> contextInfo.contexts
+    is NoContextFileInfo -> listOf(anyContext())
+    else -> emptyList()
+  }
+}
+
 val SearchScope.codeInsightContextInfo: CodeInsightContextInfo
   @Internal
   get() = if (this is CodeInsightContextAwareSearchScope) this.codeInsightContextInfo else NoContextInformation()

platform/core-impl/src/com/intellij/psi/impl/file/PsiDirectoryImpl.java

@@ -30,7 +30,10 @@ import com.intellij.psi.impl.CheckUtil;
 import com.intellij.psi.impl.PsiElementBase;
 import com.intellij.psi.impl.PsiManagerImpl;
 import com.intellij.psi.impl.source.PsiFileImpl;
-import com.intellij.psi.search.*;
+import com.intellij.psi.search.CodeInsightContextAwareSearchScopes;
+import com.intellij.psi.search.GlobalSearchScope;
+import com.intellij.psi.search.PsiElementProcessor;
+import com.intellij.psi.search.PsiFileSystemItemProcessor;
 import com.intellij.psi.util.PsiUtilCore;
 import com.intellij.testFramework.LightVirtualFile;
 import com.intellij.util.ArrayUtilRt;
@@ -42,6 +45,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
 
@@ -159,39 +163,20 @@ public class PsiDirectoryImpl extends PsiElementBase implements PsiDirectory, Qu
     for (VirtualFile file : files) {
       // The scope allows us to pre-filter the virtual files and avoid creating unnecessary PSI files.
 
-      PsiFile psiFile;
-      if (sharedSourceSupportEnabled) {
-        if (scope != null) {
-          CodeInsightContextFileInfo contextInfo = CodeInsightContextAwareSearchScopes.getFileContextInfo(scope, file);
-          if (contextInfo instanceof NoContextFileInfo) {
-            // file is in scope, but context is not specified
-            psiFile = myManager.findFile(file);
+      if (sharedSourceSupportEnabled && scope != null) {
+        Collection<CodeInsightContext> contexts = CodeInsightContextAwareSearchScopes.getCorrespondingContexts(scope, file);
+        for (CodeInsightContext context : contexts) {
+          PsiFile psiFile = myManager.findFile(file, context);
+          if (psiFile != null) {
+            psiFiles.add(psiFile);
           }
-          else if (contextInfo instanceof ActualContextFileInfo) {
-            // file is in scope, context is specified
-            CodeInsightContext context = ((ActualContextFileInfo)contextInfo).getContexts().iterator().next();
-            psiFile = myManager.findFile(file, context);
-          }
-          else {
-            // DoesNotContainFileInfo, file is not in scope
-            psiFile = null;
-          }
-        }
-        else {
-          psiFile = myManager.findFile(file);
         }
       }
-      else {
-        if (scope == null || scope.contains(file)) {
-          psiFile = myManager.findFile(file);
+      else if (scope == null || scope.contains(file)) {
+        PsiFile psiFile = myManager.findFile(file);
+        if (psiFile != null) {
+          psiFiles.add(psiFile);
         }
-        else {
-          psiFile = null;
-        }
-      }
-
-      if (psiFile != null) {
-        psiFiles.add(psiFile);
       }
     }
     return PsiUtilCore.toPsiFileArray(psiFiles);

platform/platform-tests/intellij.platform.tests.iml

@@ -120,5 +120,6 @@
     <orderEntry type="module" module-name="intellij.platform.bookmarks" scope="TEST" />
     <orderEntry type="module" module-name="intellij.platform.statistics.testFramework" scope="TEST" />
     <orderEntry type="module" module-name="intellij.platform.ide.bootstrap.eel" scope="TEST" />
+    <orderEntry type="module" module-name="intellij.platform.testFramework.junit5.projectStructure" scope="TEST" />
   </component>
 </module>

platform/platform-tests/testSrc/com/intellij/multiverse/PsiDirectoryTest.kt

@@ -0,0 +1,57 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.multiverse
+
+import com.intellij.openapi.application.readAction
+import com.intellij.openapi.module.ModuleManager
+import com.intellij.openapi.module.impl.scopes.ModulesScope
+import com.intellij.openapi.vfs.VfsUtil
+import com.intellij.platform.testFramework.junit5.projectStructure.fixture.multiverseProjectFixture
+import com.intellij.psi.PsiManager
+import com.intellij.testFramework.common.timeoutRunBlocking
+import com.intellij.testFramework.junit5.TestApplication
+import org.junit.jupiter.api.Test
+import kotlin.io.path.Path
+
+
+@TestApplication
+internal class PsiDirectoryTest {
+  companion object {
+    private val projectFixture = multiverseProjectFixture {
+      module("foo") {
+        contentRoot("root") {
+          sourceRoot("src", SHARED_ID) {
+            file("A.java", "")
+          }
+        }
+      }
+      module("bar") {
+        sharedSourceRoot(SHARED_ID)
+      }
+    }
+    private const val SHARED_ID = "shared"
+  }
+
+  @Test
+  fun `PsiDirectory#getFiles should return all PsiFiles matching scope`() = timeoutRunBlocking {
+    readAction {
+      val project = projectFixture.get()
+      val dirPath = Path(project.basePath!!, "foo/root/src")
+      val vDir = VfsUtil.findFile(dirPath, false) ?: error("Can't find virtual directory $dirPath")
+      val modules = ModuleManager.getInstance(project).modules
+      assert(modules.size == 2) { modules.contentToString() }
+
+      val modulesScope = ModulesScope(modules.toSet(), project)
+      val psiDir = PsiManager.getInstance(project).findDirectory(vDir) ?: error("Can't find PsiDirectory for $vDir")
+
+      // we have `psiDir` that matches `modulesScope` 2 times (once for the module `foo` and once for module `bar`)
+      // expected result: psiDir.getFiles(modulesScope) returns two PsiFiles corresponding to A.java, one PsiFile in the context of `foo`, and another one in the scope of `bar`
+
+      val files = psiDir.getFiles(modulesScope)
+      assert(files.size == 2) { files.contentToString() }
+      assert(files[0].name == "A.java")
+      assert(files[1].name == "A.java")
+
+      assert(files[0] != files[1])
+    }
+  }
+}