From e0b93b48f7c837c993ef976796dc05f015f4ec4d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Yann=20C=C3=A9bron?= <yann.cebron@jetbrains.com>
Date: Fri, 27 Apr 2018 14:59:26 +0200
Subject: [PATCH] [devkit] PluginDescriptorXIncludeReferenceContributor: filter
 and cache roots (IDEA-166677)

---
 ...DescriptorXIncludeFileReferenceHelper.java | 51 +++++++++++++------
 1 file changed, 35 insertions(+), 16 deletions(-)

diff --git a/plugins/devkit/devkit-core/src/references/PluginDescriptorXIncludeFileReferenceHelper.java b/plugins/devkit/devkit-core/src/references/PluginDescriptorXIncludeFileReferenceHelper.java
index aea510569fbb..282517e615a4 100644
--- a/plugins/devkit/devkit-core/src/references/PluginDescriptorXIncludeFileReferenceHelper.java
+++ b/plugins/devkit/devkit-core/src/references/PluginDescriptorXIncludeFileReferenceHelper.java
@@ -1,26 +1,31 @@
-// Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package org.jetbrains.idea.devkit.references;
 
 import com.intellij.ide.highlighter.XmlFileType;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.project.Project;
-import com.intellij.openapi.roots.ProjectRootManager;
+import com.intellij.openapi.roots.ProjectFileIndex;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiFileSystemItem;
 import com.intellij.psi.PsiManager;
 import com.intellij.psi.impl.source.resolve.reference.impl.providers.FileReferenceHelper;
+import com.intellij.psi.search.GlobalSearchScopesCore;
+import com.intellij.psi.util.CachedValueProvider;
+import com.intellij.psi.util.CachedValuesManager;
+import com.intellij.psi.util.PsiModificationTracker;
+import com.intellij.util.containers.ContainerUtil;
+import com.intellij.util.xml.DomService;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.idea.devkit.dom.IdeaPlugin;
 import org.jetbrains.idea.devkit.util.DescriptorUtil;
 import org.jetbrains.idea.devkit.util.PsiUtil;
-import org.jetbrains.jps.model.java.JavaModuleSourceRootTypes;
 
-import java.util.Collection;
-import java.util.List;
-import java.util.Objects;
+import java.util.*;
 import java.util.stream.Collectors;
 
 /**
  * Provides resource roots to be used as contexts in FileReferenceSet during the XIncludes resolving.
+ *
  * @see PluginDescriptorXIncludeReferenceContributor
  */
 public class PluginDescriptorXIncludeFileReferenceHelper extends FileReferenceHelper {
@@ -34,22 +39,36 @@ public class PluginDescriptorXIncludeFileReferenceHelper extends FileReferenceHe
   @NotNull
   @Override
   public Collection<PsiFileSystemItem> getContexts(Project project, @NotNull VirtualFile file) {
-   return getRoots(project);
+    return getRootsContainingPluginXmlFiles(project);
   }
 
   @NotNull
   @Override
   public Collection<PsiFileSystemItem> getRoots(@NotNull Module module) {
-    return getRoots(module.getProject());
+    return getRootsContainingPluginXmlFiles(module.getProject());
   }
 
-  private static Collection<PsiFileSystemItem> getRoots(@NotNull Project project) {
-    List<VirtualFile> roots = ProjectRootManager.getInstance(project)
-      .getModuleSourceRoots(JavaModuleSourceRootTypes.PRODUCTION);
-    PsiManager psiManager = PsiManager.getInstance(project);
-    return roots.stream()
-      .map(virtualFile -> psiManager.findDirectory(virtualFile))
-      .filter(Objects::nonNull)
-      .collect(Collectors.toList());
+  private static Collection<PsiFileSystemItem> getRootsContainingPluginXmlFiles(@NotNull Project project) {
+    return CachedValuesManager.getManager(project).getCachedValue(project, () -> {
+
+      Collection<VirtualFile> pluginXmlFilesInProductionScope =
+        DomService.getInstance().getDomFileCandidates(IdeaPlugin.class, project,
+                                                      GlobalSearchScopesCore.projectProductionScope(project));
+
+      ProjectFileIndex projectFileIndex = ProjectFileIndex.getInstance(project);
+      Set<VirtualFile> pluginXmlSourceRoots = new HashSet<>();
+      for (VirtualFile pluginXml : pluginXmlFilesInProductionScope) {
+        VirtualFile sourceRoot = projectFileIndex.getSourceRootForFile(pluginXml);
+        ContainerUtil.addIfNotNull(pluginXmlSourceRoots, sourceRoot);
+      }
+
+      PsiManager psiManager = PsiManager.getInstance(project);
+      List<PsiFileSystemItem> sourceRootsPsiFileSystemItems =
+        pluginXmlSourceRoots.stream()
+                            .map(virtualFile -> psiManager.findDirectory(virtualFile))
+                            .filter(Objects::nonNull)
+                            .collect(Collectors.toList());
+      return CachedValueProvider.Result.create(sourceRootsPsiFileSystemItems, PsiModificationTracker.MODIFICATION_COUNT);
+    });
   }
 }