#30: Refactor file reference source roots lookup logic

Author: kshchepanovskyi

src/main/java/io/protostuff/jetbrains/plugin/psi/FileReferenceNode.java

@@ -79,7 +79,7 @@ public ProtoPsiFileRoot getTarget() {
     }
 
     private ProtoPsiFileRoot getTarget(@NotNull String filename, @NotNull Module module) {
-        Collection<PsiFileSystemItem> roots = FilePathReferenceProvider.getRoots(module);
+        Collection<PsiFileSystemItem> roots = new FilePathReferenceProvider().getRoots(module);
         for (PsiFileSystemItem root : roots) {
             VirtualFile file = root.getVirtualFile().findFileByRelativePath(getFilename());
             if (file != null) {

src/main/java/io/protostuff/jetbrains/plugin/reference/file/AllSourceRootsProvider.java

@@ -0,0 +1,19 @@
+package io.protostuff.jetbrains.plugin.reference.file;
+
+import com.intellij.openapi.module.Module;
+import com.intellij.openapi.roots.ModuleRootManager;
+import com.intellij.openapi.vfs.VirtualFile;
+
+/**
+ * Returns "all source roots". In the IntelliJ IDEA it include source & resource roots,
+ * for both main and test source sets.
+ *
+ * @author Kostiantyn Shchepanovskyi
+ */
+class AllSourceRootsProvider implements FilePathReferenceProvider.SourceRootsProvider {
+    @Override
+    public VirtualFile[] getSourceRoots(Module module) {
+        ModuleRootManager moduleRootManager = ModuleRootManager.getInstance(module);
+        return moduleRootManager.orderEntries().getAllSourceRoots();
+    }
+}

src/main/java/io/protostuff/jetbrains/plugin/reference/file/CustomIncludePathRootsProvider.java

@@ -0,0 +1,31 @@
+package io.protostuff.jetbrains.plugin.reference.file;
+
+import com.intellij.openapi.module.Module;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.vfs.LocalFileSystem;
+import com.intellij.openapi.vfs.VirtualFile;
+import io.protostuff.jetbrains.plugin.settings.ProtobufSettings;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Returns source roots for custom include paths.
+ *
+ * @author Kostiantyn Shchepanovskyi
+ */
+class CustomIncludePathRootsProvider implements FilePathReferenceProvider.SourceRootsProvider {
+    @Override
+    public VirtualFile[] getSourceRoots(Module module) {
+        List<VirtualFile> result = new ArrayList<>();
+        Project project = module.getProject();
+        ProtobufSettings settings = project.getComponent(ProtobufSettings.class);
+        List<String> includePaths = settings.getIncludePaths();
+        for (String includePath : includePaths) {
+            VirtualFile path = LocalFileSystem.getInstance().findFileByPath(includePath);
+            if (path != null && path.isDirectory()) {
+                result.add(path);
+            }
+        }
+        return result.toArray(new VirtualFile[0]);
+    }
+}

src/main/java/io/protostuff/jetbrains/plugin/reference/file/FilePathReferenceProvider.java

@@ -18,11 +18,8 @@
 
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.module.ModuleUtilCore;
-import com.intellij.openapi.project.Project;
-import com.intellij.openapi.roots.ModuleRootManager;
 import com.intellij.openapi.util.Condition;
 import com.intellij.openapi.util.TextRange;
-import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiCompiledElement;
 import com.intellij.psi.PsiDirectory;
@@ -35,7 +32,7 @@
 import com.intellij.psi.impl.source.resolve.reference.impl.providers.FileReferenceSet;
 import com.intellij.util.ProcessingContext;
 import com.intellij.util.containers.ContainerUtil;
-import io.protostuff.jetbrains.plugin.settings.ProtobufSettings;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -53,48 +50,39 @@ public class FilePathReferenceProvider extends PsiReferenceProvider {
 
     private final boolean myEndingSlashNotAllowed;
 
+    interface SourceRootsProvider {
+        VirtualFile[] getSourceRoots(Module module);
+    }
+
+    private List<SourceRootsProvider> sourceRootsProviders = new ArrayList<>();
+
     public FilePathReferenceProvider() {
         this(true);
     }
 
     public FilePathReferenceProvider(boolean endingSlashNotAllowed) {
         myEndingSlashNotAllowed = endingSlashNotAllowed;
+        sourceRootsProviders.add(new AllSourceRootsProvider());
+        sourceRootsProviders.add(new CustomIncludePathRootsProvider());
+        sourceRootsProviders.add(new LibrariesAndSdkClassesRootsProvider());
     }
 
     @NotNull
-    public static Collection<PsiFileSystemItem> getRoots(@Nullable final Module thisModule) {
-        if (thisModule == null) {
+    public Collection<PsiFileSystemItem> getRoots(@Nullable final Module module) {
+        if (module == null) {
             return Collections.emptyList();
         }
 
-        ModuleRootManager moduleRootManager = ModuleRootManager.getInstance(thisModule);
         Set<PsiFileSystemItem> result = ContainerUtil.newLinkedHashSet();
-        final PsiManager psiManager = PsiManager.getInstance(thisModule.getProject());
-
-        VirtualFile[] libraryUrls = moduleRootManager.orderEntries().getAllLibrariesAndSdkClassesRoots();
-        for (VirtualFile file : libraryUrls) {
-            PsiDirectory directory = psiManager.findDirectory(file);
-            if (directory != null) {
-                result.add(directory);
-            }
-        }
-
-        VirtualFile[] sourceRoots = moduleRootManager.orderEntries().getAllSourceRoots();
-        for (VirtualFile root : sourceRoots) {
-            final PsiDirectory directory = psiManager.findDirectory(root);
-            if (directory != null) {
-                result.add(directory);
-            }
-        }
-
-        Project project = thisModule.getProject();
-        ProtobufSettings settings = project.getComponent(ProtobufSettings.class);
-        List<String> includePaths = settings.getIncludePaths();
-        for (String includePath : includePaths) {
-            VirtualFile path = LocalFileSystem.getInstance().findFileByPath(includePath);
-            if (path != null && path.isDirectory()) {
-                PsiDirectory psiDirectory = psiManager.findDirectory(path);
-                result.add(psiDirectory);
+        PsiManager psiManager = PsiManager.getInstance(module.getProject());
+
+        for (SourceRootsProvider sourceRootsProvider : sourceRootsProviders) {
+            VirtualFile[] sourceRoots = sourceRootsProvider.getSourceRoots(module);
+            for (VirtualFile root : sourceRoots) {
+                final PsiDirectory directory = psiManager.findDirectory(root);
+                if (directory != null) {
+                    result.add(directory);
+                }
             }
         }
         return result;
@@ -188,4 +176,5 @@ protected boolean isPsiElementAccepted(PsiElement element) {
     protected FileReference createFileReference(FileReferenceSet referenceSet, final TextRange range, final int index, final String text) {
         return new FileReference(referenceSet, range, index, text);
     }
+
 }

src/main/java/io/protostuff/jetbrains/plugin/reference/file/LibrariesAndSdkClassesRootsProvider.java

@@ -0,0 +1,18 @@
+package io.protostuff.jetbrains.plugin.reference.file;
+
+import com.intellij.openapi.module.Module;
+import com.intellij.openapi.roots.ModuleRootManager;
+import com.intellij.openapi.vfs.VirtualFile;
+
+/**
+ * Returns source roots for libraries and SDK classes.
+ *
+ * @author Kostiantyn Shchepanovskyi
+ */
+class LibrariesAndSdkClassesRootsProvider implements FilePathReferenceProvider.SourceRootsProvider {
+    @Override
+    public VirtualFile[] getSourceRoots(Module module) {
+        ModuleRootManager moduleRootManager = ModuleRootManager.getInstance(module);
+        return moduleRootManager.orderEntries().getAllLibrariesAndSdkClassesRoots();
+    }
+}