package ru.artyushov.jmhPlugin.configuration;

import com.intellij.codeInspection.reference.EntryPoint;
import com.intellij.codeInspection.reference.RefElement;
import com.intellij.openapi.util.InvalidDataException;
import com.intellij.openapi.util.WriteExternalException;
import com.intellij.psi.PsiClass;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiMethod;
import org.jdom.Element;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import static ru.artyushov.jmhPlugin.configuration.ConfigurationUtils.isBenchmarkClass;
import static ru.artyushov.jmhPlugin.configuration.ConfigurationUtils.hasSetupOrTearDownAnnotation;
import static ru.artyushov.jmhPlugin.configuration.ConfigurationUtils.isBenchmarkMethod;

/**
 * User: nikart
 * Date: 07/08/14
 * Time: 23:33
 */
public class JmhEntryPoint extends EntryPoint {

  private boolean isSelected = true;

  @NotNull
  @Override
  public String getDisplayName() {
    return "jmhEntryPoint";
  }

  @Override
  public boolean isEntryPoint(@NotNull RefElement refElement, @NotNull PsiElement psiElement) {
    return isEntryPoint(psiElement);
  }

  @Override
  public boolean isEntryPoint(@NotNull PsiElement psiElement) {
    if (isSelected) {
      if (psiElement instanceof PsiClass) {
        final PsiClass aClass = (PsiClass)psiElement;
        if (isBenchmarkClass(aClass)) {
          return true;
        }
      } else if (psiElement instanceof PsiMethod) {
        final PsiMethod method = (PsiMethod)psiElement;
        if (isBenchmarkMethod(method) || hasSetupOrTearDownAnnotation(method)) {
          return true;
        }
      }
    }
    return false;
  }

  @Override
  public boolean isSelected() {
    return isSelected;
  }

  @Override
  public void setSelected(boolean isSelected) {
    this.isSelected = isSelected;
  }

  @Override
  public void readExternal(Element element) throws InvalidDataException {
    isSelected = Boolean.parseBoolean(element.getAttributeValue("isSelected"));
  }

  @Override
  public void writeExternal(Element element) throws WriteExternalException {
    element.setAttribute("isSelected", isSelected + "");
  }
}

@@ -270,42 +270,8 @@
 on;%0A
-import com.intellij.psi.PsiClass;%0A
 impo
@@ -408,139 +408,8 @@
 ll;%0A
-import org.jetbrains.annotations.Nullable;%0A%0Aimport static ru.artyushov.jmhPlugin.configuration.ConfigurationUtils.isBenchmarkClass;
 %0Aimp
@@ -586,22 +586,28 @@
 enchmark
-Method
+EntryElement
 ;%0A%0A/**%0A 
@@ -1094,162 +1094,50 @@
     
-if (psiElement instanceof PsiClass) %7B%0A        final PsiClass aClass = (PsiClass)psiElement;%0A        if (isBenchmarkClass(aClass)) %7B%0A          return true;
+return isBenchmarkEntryElement(psiElement)
 %0A   
@@ -1145,26 +1145,18 @@
     
-%7D%0A
       
-%7D else if 
+%7C%7C (
 (psi
@@ -1189,106 +1189,10 @@
 od) 
-%7B%0A        final PsiMethod method = (PsiMethod)psiElement;%0A        if (isBenchmarkMethod(method) %7C%7C
+&&
  has
@@ -1221,59 +1221,33 @@
 ion(
-m
+(PsiM
 ethod)
-) %7B%0A          return true;%0A        %7D%0A      %7D
+ psiElement));
 %0A   

Use isBenchmarkEntryElement()