#66: Resolve imports relative to file's location

Many projects are storing proto files in a locations that are not marked as source/resource root.
As a result, imports are highlighted with error "file does not exist", and all imported types are not resolvable.

Plugin can try to resolve imports relative to source file's location using following rules:

1. If `package` is not set, try to look in the same folder where source file is.
2. If package is set, try to look from parent folder.
3. Parent folder should be computed as source file's folders, with removed corresponding parts of the package.

For example:

`/home/user/project/foo/bar/hello.proto`:

```proto
package foo.bar;
import "baz/import.proto";
```

Import should be resolved to `/home/user/project/baz/import.proto`.

Author: kshchepanovskyi

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

@@ -5,6 +5,7 @@
 import com.intellij.openapi.module.ModuleManager;
 import com.intellij.openapi.module.ModuleUtilCore;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiFileSystemItem;
 import com.intellij.psi.PsiManager;
@@ -79,7 +80,7 @@ public ProtoPsiFileRoot getTarget() {
     }
 
     private ProtoPsiFileRoot getTarget(@NotNull String filename, @NotNull Module module) {
-        Collection<PsiFileSystemItem> roots = new FilePathReferenceProvider().getRoots(module);
+        Collection<PsiFileSystemItem> roots = new FilePathReferenceProvider().getRoots(module, getRoot());
         for (PsiFileSystemItem root : roots) {
             VirtualFile file = root.getVirtualFile().findFileByRelativePath(getFilename());
             if (file != null) {
@@ -101,6 +102,15 @@ private ProtoPsiFileRoot getTarget(@NotNull String filename, @NotNull Module mod
         return null;
     }
 
+    @Nullable
+    private ProtoPsiFileRoot getRoot() {
+        PsiElement node = this;
+        while (!(node instanceof ProtoPsiFileRoot || node == null)) {
+            node = node.getParent();
+        }
+        return (ProtoPsiFileRoot) node;
+    }
+
     @Nullable
     private String getFilename() {
         String text = getText();

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

@@ -3,6 +3,7 @@
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.roots.ModuleRootManager;
 import com.intellij.openapi.vfs.VirtualFile;
+import io.protostuff.jetbrains.plugin.psi.ProtoPsiFileRoot;
 
 /**
  * Returns "all source roots". In the IntelliJ IDEA it include source & resource roots,
@@ -12,7 +13,7 @@
  */
 class AllSourceRootsProvider implements FilePathReferenceProvider.SourceRootsProvider {
     @Override
-    public VirtualFile[] getSourceRoots(Module module) {
+    public VirtualFile[] getSourceRoots(Module module, ProtoPsiFileRoot psiFileRoot) {
         ModuleRootManager moduleRootManager = ModuleRootManager.getInstance(module);
         return moduleRootManager.orderEntries().getAllSourceRoots();
     }

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

@@ -4,6 +4,7 @@
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VirtualFile;
+import io.protostuff.jetbrains.plugin.psi.ProtoPsiFileRoot;
 import io.protostuff.jetbrains.plugin.settings.ProtobufSettings;
 import java.util.ArrayList;
 import java.util.List;
@@ -15,7 +16,7 @@
  */
 class CustomIncludePathRootsProvider implements FilePathReferenceProvider.SourceRootsProvider {
     @Override
-    public VirtualFile[] getSourceRoots(Module module) {
+    public VirtualFile[] getSourceRoots(Module module, ProtoPsiFileRoot psiFileRoot) {
         List<VirtualFile> result = new ArrayList<>();
         Project project = module.getProject();
         ProtobufSettings settings = project.getComponent(ProtobufSettings.class);

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

@@ -32,6 +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.psi.ProtoPsiFileRoot;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -51,7 +52,7 @@ public class FilePathReferenceProvider extends PsiReferenceProvider {
     private final boolean myEndingSlashNotAllowed;
 
     interface SourceRootsProvider {
-        VirtualFile[] getSourceRoots(Module module);
+        VirtualFile[] getSourceRoots(Module module, ProtoPsiFileRoot psiFileRoot);
     }
 
     private List<SourceRootsProvider> sourceRootsProviders = new ArrayList<>();
@@ -66,10 +67,11 @@ public FilePathReferenceProvider(boolean endingSlashNotAllowed) {
         sourceRootsProviders.add(new WebCoreResourcePathRootsProvider());
         sourceRootsProviders.add(new CustomIncludePathRootsProvider());
         sourceRootsProviders.add(new LibrariesAndSdkClassesRootsProvider());
+        sourceRootsProviders.add(new ProtoFileRelativePathRootsProvider());
     }
 
     @NotNull
-    public Collection<PsiFileSystemItem> getRoots(@Nullable final Module module) {
+    public Collection<PsiFileSystemItem> getRoots(@Nullable final Module module, ProtoPsiFileRoot psiFileRoot) {
         if (module == null) {
             return Collections.emptyList();
         }
@@ -78,7 +80,7 @@ public Collection<PsiFileSystemItem> getRoots(@Nullable final Module module) {
         PsiManager psiManager = PsiManager.getInstance(module.getProject());
 
         for (SourceRootsProvider sourceRootsProvider : sourceRootsProviders) {
-            VirtualFile[] sourceRoots = sourceRootsProvider.getSourceRoots(module);
+            VirtualFile[] sourceRoots = sourceRootsProvider.getSourceRoots(module, psiFileRoot);
             for (VirtualFile root : sourceRoots) {
                 final PsiDirectory directory = psiManager.findDirectory(root);
                 if (directory != null) {
@@ -130,12 +132,12 @@ public Collection<PsiFileSystemItem> computeDefaultContexts() {
                 if (forModules.length > 0) {
                     Set<PsiFileSystemItem> rootsForModules = ContainerUtil.newLinkedHashSet();
                     for (Module forModule : forModules) {
-                        rootsForModules.addAll(getRoots(forModule));
+                        rootsForModules.addAll(getRoots(forModule, null));
                     }
                     return rootsForModules;
                 }
 
-                return getRoots(ModuleUtilCore.findModuleForPsiElement(getElement()));
+                return getRoots(ModuleUtilCore.findModuleForPsiElement(getElement()), null);
             }
 
             @Override

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

@@ -3,6 +3,7 @@
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.roots.ModuleRootManager;
 import com.intellij.openapi.vfs.VirtualFile;
+import io.protostuff.jetbrains.plugin.psi.ProtoPsiFileRoot;
 
 /**
  * Returns source roots for libraries and SDK classes.
@@ -11,7 +12,7 @@
  */
 class LibrariesAndSdkClassesRootsProvider implements FilePathReferenceProvider.SourceRootsProvider {
     @Override
-    public VirtualFile[] getSourceRoots(Module module) {
+    public VirtualFile[] getSourceRoots(Module module, ProtoPsiFileRoot psiFileRoot) {
         ModuleRootManager moduleRootManager = ModuleRootManager.getInstance(module);
         return moduleRootManager.orderEntries().getAllLibrariesAndSdkClassesRoots();
     }

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

@@ -0,0 +1,82 @@
+package io.protostuff.jetbrains.plugin.reference.file;
+
+import com.intellij.openapi.module.Module;
+import com.intellij.openapi.vfs.VirtualFile;
+import io.protostuff.jetbrains.plugin.psi.ProtoPsiFileRoot;
+import java.util.Arrays;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Many projects are storing proto files in a locations that are not marked as source/resource root.
+ * As a result, imports are highlighted with error "file does not exist", and all imported types are not resolvable.
+ *
+ * Plugin can try to resolve imports relative to source file's location using following rules:
+ *
+ * 1. If `package` is not set, try to look in the same folder where source file is.
+ * 2. If package is set, try to look from parent folder.
+ * 3. Parent folder should be computed as source file's folders, with removed corresponding parts of the package.
+ *
+ * For example:
+ *
+ * `/home/user/project/foo/bar/hello.proto`:
+ *
+ * ```proto
+ * package foo.bar;
+ * import "baz/import.proto";
+ * ```
+ *
+ * Import should be resolved to `/home/user/project/baz/import.proto`.
+ *
+ * @author Kostiantyn Shchepanovskyi
+ */
+class ProtoFileRelativePathRootsProvider implements FilePathReferenceProvider.SourceRootsProvider {
+
+    private static final VirtualFile[] NONE = new VirtualFile[0];
+
+    @Override
+    public VirtualFile[] getSourceRoots(Module module, @Nullable ProtoPsiFileRoot psiFileRoot) {
+        if (psiFileRoot != null) {
+            VirtualFile file = psiFileRoot.getVirtualFile();
+            if (file != null) {
+                String packageName = psiFileRoot.getPackageName();
+                // no package - return parent dir
+                VirtualFile dir = file.getParent();
+                int nestingLevel = computeNestingLevel(packageName);
+                dir = goUp(dir, nestingLevel);
+                if (dir != null) {
+                    return toArray(dir);
+                }
+            }
+        }
+        return NONE;
+    }
+
+    @Nullable
+    private VirtualFile goUp(VirtualFile dir, int times) {
+        while (times > 0 && dir != null) {
+            dir = dir.getParent();
+            times--;
+        }
+        return dir;
+    }
+
+    private int computeNestingLevel(String packageName) {
+        // compute nesting level
+        int nestingLevel = 0;
+        if (!packageName.isEmpty()) {
+            nestingLevel++;
+            for (char c : packageName.toCharArray()) {
+                if (c == '.') {
+                    nestingLevel++;
+                }
+            }
+        }
+        return nestingLevel;
+    }
+
+    @NotNull
+    private VirtualFile[] toArray(VirtualFile dir) {
+        return new VirtualFile[]{dir};
+    }
+}

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

@@ -7,6 +7,7 @@
 import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.util.containers.MultiMap;
+import io.protostuff.jetbrains.plugin.psi.ProtoPsiFileRoot;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -47,7 +48,7 @@ class WebCoreResourcePathRootsProvider implements FilePathReferenceProvider.Sour
 
     @SuppressWarnings("unchecked")
     @Override
-    public VirtualFile[] getSourceRoots(Module module) {
+    public VirtualFile[] getSourceRoots(Module module, ProtoPsiFileRoot psiFileRoot) {
         try {
             if (GET_INSTANCE != null && GET_RESOURCE_ROOTS != null) {
                 Object configurationInstance = GET_INSTANCE.invoke(null, module.getProject());

src/test/java/io/protostuff/jetbrains/plugin/reference/ReferenceTest.java

@@ -31,7 +31,27 @@ public void testEnumReference() {
     }
 
     public void testImportedMessageReference() {
-        checkReferenceToDataType(".reference.ImportedMessage", "reference/ImportedMessageReferenceTestData.proto", "reference/ImportedTestData.proto");
+        checkReferenceToDataType(".reference.ImportedMessage",
+                "reference/ImportedMessageReferenceTestData.proto",
+                "reference/ImportedTestData.proto");
+    }
+
+    public void testImportedRelativelyMessageReference() {
+        checkReferenceToDataType(".import.ImportedMessage1",
+                "reference/relative/import/ImportedRelativelyMessageReferenceTestData.proto",
+                "reference/relative/import/ImportedRelativelyTestData.proto");
+    }
+
+    public void testImportedRelativelyNoPackageMessageReference() {
+        checkReferenceToDataType(".import.ImportedMessage1",
+                "reference/relative/import/ImportedRelativelyNoPackageMessageReferenceTestData.proto",
+                "reference/relative/import/ImportedRelativelyTestData.proto");
+    }
+
+    public void testImportedRelativelyTwoLevelsMessageReference() {
+        checkReferenceToDataType(".import.ImportedMessage1",
+                "reference/relative/import/ImportedRelativelyTwoLevelsMessageReferenceTestData.proto",
+                "reference/relative/import/ImportedRelativelyTestData.proto");
     }
 
     /**

src/test/resources/reference/relative/import/ImportedRelativelyMessageReferenceTestData.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package import;
+
+import "import/ImportedRelativelyTestData.proto";
+
+message TestMessage {
+    <caret>ImportedMessage1 message = 1;
+}

src/test/resources/reference/relative/import/ImportedRelativelyNoPackageMessageReferenceTestData.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+import "ImportedRelativelyTestData.proto";
+
+message TestMessage {
+    import.<caret>ImportedMessage1 message = 1;
+}