Highlight func keyword when caret is on return

go-lang-plugin-org · Nov 26, 2015 · 070fbb6 · 070fbb6
1 parent 6dcdfa4
commit 070fbb6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/src/com/goide/highlighting/exitpoint/FunctionExitPointHandler.java b/src/com/goide/highlighting/exitpoint/FunctionExitPointHandler.java
@@ -16,14 +16,12 @@
 
 package com.goide.highlighting.exitpoint;
 
-import com.goide.GoTypes;
 import com.goide.psi.*;
 import com.goide.psi.impl.GoPsiImplUtil;
 import com.intellij.codeInsight.highlighting.HighlightUsagesHandlerBase;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
-import com.intellij.psi.impl.source.tree.LeafPsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.Consumer;
 import com.intellij.util.containers.ContainerUtil;
@@ -55,8 +53,8 @@ protected void selectTargets(List<PsiElement> targets, @NotNull Consumer<List<Ps
 
   @Override
   public void computeUsages(List<PsiElement> targets) {
-    if (myTarget instanceof LeafPsiElement && ((LeafPsiElement)myTarget).getElementType() == GoTypes.FUNC) {
-      addOccurrence(myTarget);
+    if (myFunction instanceof GoFunctionOrMethodDeclaration) {
+      addOccurrence(((GoFunctionOrMethodDeclaration)myFunction).getFunc());
     }
     new GoRecursiveVisitor() {
       @Override

diff --git a/tests/com/goide/editor/GoExitPointsHighlightingTest.java b/tests/com/goide/editor/GoExitPointsHighlightingTest.java
@@ -48,7 +48,7 @@ public void testBasicExitPoints() {
                   "  }" +
                   "  return x\n" +
                   "}";
-    doTest(text, "return -1", "return x");
+    doTest(text, "func", "return -1", "return x");
   }
 
   public void testCaretOnFuncWithReturnAndPanic() {