#37: Resolve custom option references

Resolve custom options declared in imported proto files.

Author: kshchepanovskyi

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

@@ -6,18 +6,14 @@
 import com.intellij.psi.FileViewProvider;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiNamedElement;
-import com.intellij.psi.PsiReference;
 import io.protostuff.jetbrains.plugin.Icons;
 import io.protostuff.jetbrains.plugin.ProtoFileType;
 import io.protostuff.jetbrains.plugin.ProtoLanguage;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Deque;
 import java.util.List;
 import java.util.Objects;
-import java.util.stream.Collectors;
 import javax.swing.Icon;
 import org.antlr.jetbrains.adapter.psi.ScopeNode;
 import org.jetbrains.annotations.NotNull;
@@ -84,46 +80,6 @@ private ProtoRootNode getProtoRoot() {
         return findChildByClass(ProtoRootNode.class);
     }
 
-    /**
-     * Returns extensions defined in this proto file.
-     */
-    public Collection<ExtendNode> getLocalExtensions() {
-        List<ExtendNode> result = new ArrayList<>();
-        Deque<DataTypeContainer> queue = new ArrayDeque<>();
-        queue.push(getProtoRoot());
-        while (!queue.isEmpty()) {
-            DataTypeContainer container = queue.poll();
-            Collection<ExtendNode> extensions = container.getDeclaredExtensions();
-            result.addAll(extensions);
-            for (DataType dataType : container.getDeclaredDataTypes()) {
-                if (dataType instanceof DataTypeContainer) {
-                    DataTypeContainer childContainer = (DataTypeContainer) dataType;
-                    queue.push(childContainer);
-                }
-            }
-        }
-        return result;
-    }
-
-    /**
-     * Returns all extension for given target message.
-     */
-    public Collection<ExtendNode> getExtenstions(MessageNode target) {
-        return getLocalExtensions().stream()
-                .filter(node -> {
-                    TypeReferenceNode extendTarget = node.getTarget();
-                    if (extendTarget == null) {
-                        return false;
-                    }
-                    PsiReference extendTargetReference = extendTarget.getReference();
-                    if (extendTargetReference == null) {
-                        return false;
-                    }
-                    return extendTargetReference.isReferenceTo(target);
-                })
-                .collect(Collectors.toList());
-    }
-
     private void addChildrenRecursively(Collection<? extends ProtoType> elements, List<ProtoType> result) {
         for (ProtoType element : elements) {
             if (element instanceof DataTypeContainer) {

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

@@ -1,10 +1,18 @@
 package io.protostuff.jetbrains.plugin.psi;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiReference;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Deque;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Objects;
+import java.util.Queue;
+import java.util.Set;
+import java.util.stream.Collectors;
 import org.antlr.jetbrains.adapter.psi.AntlrPsiNode;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -65,7 +73,7 @@ public DataType resolve(String typeName, Deque<String> scopeLookupList, boolean
         if (!resolveInImports) {
             return null;
         }
-        ImportNode[] importNodes = findChildrenByClass(ImportNode.class);
+        List<ImportNode> importNodes = getImports();
         for (ImportNode importNode : importNodes) {
             ProtoRootNode targetProto = importNode.getTargetProto();
             if (targetProto != null) {
@@ -80,6 +88,25 @@ public DataType resolve(String typeName, Deque<String> scopeLookupList, boolean
         return null;
     }
 
+    /**
+     * Get imports declared in this proto file.
+     */
+    @NotNull
+    public List<ImportNode> getImports() {
+        return Arrays.asList(findChildrenByClass(ImportNode.class));
+    }
+
+    /**
+     * Get public imports declared in this proto file.
+     */
+    @NotNull
+    public List<ImportNode> getPublicImports() {
+        return Arrays.stream(findChildrenByClass(ImportNode.class))
+                .filter(ImportNode::isPublic)
+                .collect(Collectors.toList());
+    }
+
+
     private DataType resolveByQualifiedName(String qualifiedName) {
         String prefix = getCurrentProtoPrefix();
         if (qualifiedName.startsWith(prefix)) {
@@ -139,4 +166,66 @@ public Collection<ProtoType> getDeclaredTypes() {
         return Arrays.asList(findChildrenByClass(ProtoType.class));
     }
 
+    /**
+     * Returns all extension for given target message.
+     */
+    public Collection<ExtendNode> getExtenstions(MessageNode target) {
+        return getExtensions().stream()
+                .filter(node -> {
+                    TypeReferenceNode extendTarget = node.getTarget();
+                    if (extendTarget == null) {
+                        return false;
+                    }
+                    PsiReference extendTargetReference = extendTarget.getReference();
+                    return extendTargetReference.isReferenceTo(target);
+                })
+                .collect(Collectors.toList());
+    }
+
+    /**
+     * Returns all extensions visible inside of this proto file.
+     */
+    public Collection<ExtendNode> getExtensions() {
+        List<ExtendNode> result = new ArrayList<>();
+        result.addAll(getLocalExtensions());
+        Queue<ImportNode> queue = new ArrayDeque<>();
+        queue.addAll(getImports());
+        Set<ProtoRootNode> processedProtos = new HashSet<>();
+        while (!queue.isEmpty()) {
+            ImportNode importNode = queue.poll();
+            ProtoRootNode targetProto = importNode.getTargetProto();
+            if (processedProtos.contains(targetProto)) {
+                // do not enter into endless loop
+                // if proto files refer to each other
+                continue;
+            }
+            processedProtos.add(targetProto);
+            if (targetProto != null) {
+                result.addAll(targetProto.getLocalExtensions());
+                queue.addAll(targetProto.getPublicImports());
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Returns local extensions declared in this proto file.
+     */
+    public Collection<ExtendNode> getLocalExtensions() {
+        List<ExtendNode> result = new ArrayList<>();
+        Deque<DataTypeContainer> queue = new ArrayDeque<>();
+        queue.push(this);
+        while (!queue.isEmpty()) {
+            DataTypeContainer container = queue.poll();
+            Collection<ExtendNode> extensions = container.getDeclaredExtensions();
+            result.addAll(extensions);
+            for (DataType dataType : container.getDeclaredDataTypes()) {
+                if (dataType instanceof DataTypeContainer) {
+                    DataTypeContainer childContainer = (DataTypeContainer) dataType;
+                    queue.push(childContainer);
+                }
+            }
+        }
+        return result;
+    }
 }

src/main/java/io/protostuff/jetbrains/plugin/reference/FieldReferenceProviderImpl.java

@@ -136,12 +136,12 @@ private String getTarget(PsiElement element) {
         return null;
     }
 
-    private ProtoPsiFileRoot getProtoRoot(PsiElement element) {
+    private ProtoRootNode getProtoRoot(PsiElement element) {
         PsiElement parent = element.getParent();
-        while (!(parent instanceof ProtoPsiFileRoot)) {
+        while (!(parent instanceof ProtoRootNode)) {
             parent = parent.getParent();
         }
-        return (ProtoPsiFileRoot) parent;
+        return (ProtoRootNode) parent;
     }
 
     private MessageField resolveStandardOptionReference(PsiElement sourceElement, MessageNode target, String key) {
@@ -178,7 +178,7 @@ private MessageField resolveDefaultOptionReference(PsiElement element) {
 
     @Nullable
     private FieldNode resolveCustomOptionReference(PsiElement element, MessageNode target, String key) {
-        ProtoPsiFileRoot protoRoot = getProtoRoot(element);
+        ProtoRootNode protoRoot = getProtoRoot(element);
         DataTypeContainer container = getContainer(element);
         Deque<String> scopeLookupList = TypeReference.createScopeLookupList(container);
         // case 1: (.package.field)

src/test/java/io/protostuff/jetbrains/plugin/IconsTest.java

@@ -32,6 +32,35 @@ public void testCustomOptionReference() {
         Assert.assertTrue(field.getParent() instanceof ExtendEntryNode);
     }
 
+    public void testImportedCustomOptionReference() {
+        myFixture.configureByFile("reference/options/custom/SimpleExtensionField.proto");
+        PsiReference reference = myFixture.getReferenceAtCaretPositionWithAssertion(
+                "reference/options/custom/ImportedExtension.proto"
+        );
+        PsiElement target = reference.resolve();
+        Assert.assertNotNull(target);
+        Assert.assertTrue(target instanceof FieldNode);
+        FieldNode field = (FieldNode) target;
+        Assert.assertEquals("bar", field.getFieldName());
+        Assert.assertTrue(field.getParent() instanceof ExtendEntryNode);
+    }
+
+    public void testImportedImportedCustomOptionReference() {
+        myFixture.configureByFiles(
+                "reference/options/custom/SimpleExtensionField.proto",
+                "reference/options/custom/ImportedExtension.proto"
+        );
+        PsiReference reference = myFixture.getReferenceAtCaretPositionWithAssertion(
+                "reference/options/custom/ImportedImportedExtension.proto"
+        );
+        PsiElement target = reference.resolve();
+        Assert.assertNotNull(target);
+        Assert.assertTrue(target instanceof FieldNode);
+        FieldNode field = (FieldNode) target;
+        Assert.assertEquals("bar", field.getFieldName());
+        Assert.assertTrue(field.getParent() instanceof ExtendEntryNode);
+    }
+
     public void testExtensionIsMessage_PointToMessage() {
         PsiReference reference = myFixture.getReferenceAtCaretPositionWithAssertion(
                 "reference/options/custom/ExtensionIsMessage.proto"

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

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package reference.options;
+package reference.options.custom;
 
 import "google/protobuf/descriptor.proto";
 

src/test/resources/reference/options/custom/ExtensionIsMessage.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package reference.options;
+package reference.options.custom;
 
 import "google/protobuf/descriptor.proto";
 

src/test/resources/reference/options/custom/ExtensionOfExtension.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package reference.options;
+package reference.options.custom;
 
 import "google/protobuf/descriptor.proto";
 

src/test/resources/reference/options/custom/FieldOfExtension.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package reference.options.custom;
+
+import public "reference/options/custom/SimpleExtensionField.proto";
+
+message A {
+    option (ba<caret>r) = 10;
+}

src/test/resources/reference/options/custom/ImportedExtension.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package reference.options.custom;
+
+import public "reference/options/custom/ImportedExtension.proto";
+
+message A {
+    option (ba<caret>r) = 10;
+}