[java] PSI augment provider cleanup

java/java-psi-api/src/com/intellij/psi/augment/PsiAugmentProvider.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2010 JetBrains s.r.o.
+ * Copyright 2000-2016 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,75 +19,109 @@ import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.extensions.ExtensionPointName;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.progress.ProcessCanceledException;
+import com.intellij.openapi.project.DumbAware;
 import com.intellij.openapi.project.DumbService;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiType;
 import com.intellij.psi.PsiTypeElement;
+import com.intellij.util.Processor;
+import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
 /**
- * An extension that enables one to add children to some PSI elements, e.g. methods to Java classes. The class code remains the same, but its
- * method accessors also include the results returned from {@link PsiAugmentProvider}s.
- * 
- * During indexing, only {@link com.intellij.openapi.project.DumbAware} augment providers are run.
+ * An extension that enables plugins to alter a behavior of Java PSI elements.
+ * During indexing, only {@link DumbAware} augment providers are run.
  */
 public abstract class PsiAugmentProvider {
-  private static final Logger LOG = Logger.getInstance("#" + PsiAugmentProvider.class.getName());
   public static final ExtensionPointName<PsiAugmentProvider> EP_NAME = ExtensionPointName.create("com.intellij.lang.psiAugmentProvider");
 
-  @NotNull
-  public abstract <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement element, @NotNull Class<Psi> type);
+  //<editor-fold desc="Methods to override in implementations.">
 
+  /**
+   * An extension that enables one to add children to some PSI elements, e.g. methods to Java classes.
+   * The class code remains the same, but its method accessors also include the results returned from {@link PsiAugmentProvider}s.
+   */
   @NotNull
-  public static <Psi extends PsiElement> List<Psi> collectAugments(@NotNull final PsiElement element, @NotNull final Class<Psi> type) {
-    List<Psi> result = Collections.emptyList();
-    for (PsiAugmentProvider provider : DumbService.getInstance(element.getProject()).filterByDumbAwareness(Extensions.getExtensions(EP_NAME))) {
-      List<Psi> augments = provider.getAugments(element, type);
-      if (!augments.isEmpty()) {
-        if (result.isEmpty()) result = new ArrayList<Psi>(augments.size());
-        result.addAll(augments);
-      }
-    }
-
-    return result;
+  protected <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement element, @NotNull Class<Psi> type) {
+    return Collections.emptyList();
   }
 
   /**
    * Extends {@link PsiTypeElement#getType()} so type could be retrieved from external place
    * e.g. from variable initializer in lombok case (http://projectlombok.org/features/val.html)
-   * 
-   * @param typeElement place where inference takes place, 
+   *
+   * @param typeElement place where inference takes place,
    *                    also nested PsiTypeElement-s (e.g. for List<String> PsiTypeElements corresponding to both List and String would be suggested)
    * @return inferred type or null, if inference is not applicable
-   * 
    * @since 14.1
    */
   @Nullable
-  protected PsiType inferType(PsiTypeElement typeElement) {
+  protected PsiType inferType(@NotNull PsiTypeElement typeElement) {
     return null;
   }
 
+  //</editor-fold>
+
+  //<editor-fold desc="API and the inner kitchen.">
+
+  @NotNull
+  public static <Psi extends PsiElement> List<Psi> collectAugments(@NotNull final PsiElement element, @NotNull final Class<Psi> type) {
+    final List<Psi> result = ContainerUtil.newSmartList();
+
+    forEach(element.getProject(), new Processor<PsiAugmentProvider>() {
+      @Override
+      public boolean process(PsiAugmentProvider provider) {
+        result.addAll(provider.getAugments(element, type));
+        return true;
+      }
+    });
+
+    return result;
+  }
+
   @Nullable
-  public static PsiType getInferredType(PsiTypeElement typeElement) {
-    for (PsiAugmentProvider provider : Extensions.getExtensions(EP_NAME)) {
-      try {
-        final PsiType type = provider.inferType(typeElement);
+  public static PsiType getInferredType(@NotNull final PsiTypeElement typeElement) {
+    final Ref<PsiType> result = Ref.create();
+
+    forEach(typeElement.getProject(), new Processor<PsiAugmentProvider>() {
+      @Override
+      public boolean process(PsiAugmentProvider provider) {
+        PsiType type = provider.inferType(typeElement);
         if (type != null) {
-          return type;
+          result.set(type);
+          return false;
+        }
+        else {
+          return true;
         }
       }
-      catch (ProcessCanceledException e) {
-        throw e;
-      }
-      catch (Exception e) {
-        LOG.error("provider: " + provider, e);
+    });
+
+    return result.get();
+  }
+
+  private static void forEach(Project project, Processor<PsiAugmentProvider> processor) {
+    for (PsiAugmentProvider provider : Extensions.getExtensions(EP_NAME)) {
+      if (!DumbService.isDumb(project) || DumbService.isDumbAware(provider)) {
+        try {
+          boolean goOn = processor.process(provider);
+          if (!goOn) break;
+        }
+        catch (ProcessCanceledException e) {
+          throw e;
+        }
+        catch (Exception e) {
+          Logger.getInstance(PsiAugmentProvider.class).error("provider: " + provider, e);
+        }
       }
     }
-    return null;
   }
-}
+
+  //</editor-fold>
+}

java/java-tests/testSrc/com/intellij/codeInsight/daemon/PsiAugmentProviderTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2015 JetBrains s.r.o.
+ * Copyright 2000-2016 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,9 +26,6 @@ import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.Collections;
-import java.util.List;
-
 public class PsiAugmentProviderTest extends LightCodeInsightFixtureTestCase {
   @Override
   protected String getTestDataPath() {
@@ -68,15 +65,9 @@ public class PsiAugmentProviderTest extends LightCodeInsightFixtureTestCase {
     private static final String LOMBOK_VAL_FQN = "lombok.val";
     private static final String LOMBOK_VAL_SHORT_NAME = "val";
 
-    @NotNull
-    @Override
-    public <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement element, @NotNull Class<Psi> type) {
-      return Collections.emptyList();
-    }
-
     @Nullable
     @Override
-    protected PsiType inferType(PsiTypeElement typeElement) {
+    protected PsiType inferType(@NotNull PsiTypeElement typeElement) {
       PsiElement parent = typeElement.getParent();
       if (parent instanceof PsiLocalVariable && ((PsiLocalVariable)parent).getInitializer() != null ||
           parent instanceof PsiParameter && ((PsiParameter)parent).getDeclarationScope() instanceof PsiForeachStatement) {