Highlight exit points of a function when cursor is on func

Author: ilinum

src/com/goide/highlighting/GoHighlightExitPointsHandlerFactory.java

@@ -63,6 +63,9 @@ 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);
+      }
       new GoRecursiveVisitor() {
         @Override
         public void visitFunctionLit(@NotNull GoFunctionLit literal) {
@@ -85,12 +88,22 @@ public void visitCallExpr(@NotNull GoCallExpr o) {
     public static MyHandler createForElement(@NotNull Editor editor, PsiFile file, PsiElement element) {
       GoTypeOwner function = PsiTreeUtil.getParentOfType(element, GoFunctionLit.class, GoFunctionOrMethodDeclaration.class);
       if (function == null) return null;
-      if (element instanceof LeafPsiElement && ((LeafPsiElement)element).getElementType() == GoTypes.RETURN || isPanicCall(element)) {
+      if (shouldCreateMyHandler(element)) {
         return new MyHandler(editor, file, element, function);
       }
       return null;
     }
 
+    private static boolean shouldCreateMyHandler(PsiElement element) {
+      if (element instanceof LeafPsiElement) {
+        LeafPsiElement leaf = (LeafPsiElement)element;
+        return leaf.getElementType() == GoTypes.RETURN || leaf.getElementType() == GoTypes.FUNC || isPanicCall(leaf);
+      }
+      else {
+        return false;
+      }
+    }
+
     private static boolean isPanicCall(@NotNull PsiElement e) {
       PsiElement parent = e.getParent();
       if (parent instanceof GoReferenceExpression) {