[inspections] IDEA-321711 Try to eliminate overhead of PsiElementVisitor calling checkCancelled in inspections

GitOrigin-RevId: a507f6e0a41aae12abe4257e5047eb4a210cbf44

platform/analysis-api/src/com/intellij/codeInspection/LocalInspectionTool.java

@@ -133,28 +133,39 @@ public abstract class LocalInspectionTool extends InspectionProfileEntry {
    * @see PsiRecursiveVisitor
    */
   public @NotNull PsiElementVisitor buildVisitor(@NotNull ProblemsHolder holder, boolean isOnTheFly) {
-    return new PsiElementVisitor() {
+    return new PsiFileElementVisitor(holder, isOnTheFly);
+  }
+
+  private class PsiFileElementVisitor extends PsiElementVisitor {
+    private final @NotNull ProblemsHolder myHolder;
+    private final boolean myIsOnTheFly;
+
+    private PsiFileElementVisitor(@NotNull ProblemsHolder holder, boolean fly) {
+      this.myHolder = holder;
+      this.myIsOnTheFly = fly;
+    }
+
     @Override
     public void visitFile(@NotNull PsiFile file) {
-        addDescriptors(checkFile(file, holder.getManager(), isOnTheFly));
+      addDescriptors(checkFile(file, myHolder.getManager(), myIsOnTheFly));
     }
 
     private void addDescriptors(ProblemDescriptor @Nullable [] descriptors) {
       if (descriptors != null) {
         for (ProblemDescriptor descriptor : descriptors) {
           if (descriptor != null) {
-              holder.registerProblem(descriptor);
+            myHolder.registerProblem(descriptor);
           }
           else {
             Class<?> inspectionToolClass = LocalInspectionTool.this.getClass();
-              LOG.error(PluginException.createByClass("Array returned from checkFile() method of " + inspectionToolClass + " contains null element: " +
+            LOG.error(PluginException.createByClass(
+              "Array returned from checkFile() method of " + inspectionToolClass + " contains null element: " +
               Arrays.toString(descriptors),
               null, inspectionToolClass));
           }
         }
       }
     }
-    };
   }
 
   /**

platform/analysis-impl/src/com/intellij/codeInspection/InspectionEngine.java

@@ -40,19 +40,24 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
+import static java.util.Collections.singletonList;
+
 public final class InspectionEngine {
   private static final Logger LOG = Logger.getInstance(InspectionEngine.class);
   private static boolean createVisitorAndAcceptElements(@NotNull LocalInspectionTool tool,
                                                         @NotNull ProblemsHolder holder,
                                                         boolean isOnTheFly,
                                                         @NotNull LocalInspectionToolSession session,
-                                                     @NotNull List<? extends PsiElement> elements) {
+                                                        @NotNull List<? extends PsiElement> elements,
+                                                        Map<Class<?>, Collection<Class<?>>> targetPsiClasses) {
     PsiElementVisitor visitor = createVisitor(tool, holder, isOnTheFly, session);
     // if inspection returned empty visitor then it should be skipped
     if (visitor == PsiElementVisitor.EMPTY_VISITOR) return false;
 
+    List<Class<?>> acceptingPsiTypes = InspectionVisitorsOptimizer.getAcceptingPsiTypes(visitor);
+
     tool.inspectionStarted(session, isOnTheFly);
-    acceptElements(elements, visitor);
+    acceptElements(elements, visitor, targetPsiClasses, acceptingPsiTypes);
     tool.inspectionFinished(session, holder);
     return true;
   }
@@ -73,14 +78,33 @@ public final class InspectionEngine {
     return visitor;
   }
 
-  private static void acceptElements(@NotNull List<? extends PsiElement> elements, @NotNull PsiElementVisitor elementVisitor) {
-    //noinspection ForLoopReplaceableByForEach
+  private static void acceptElements(@NotNull List<? extends PsiElement> elements,
+                                     @NotNull PsiElementVisitor elementVisitor,
+                                     Map<Class<?>, Collection<Class<?>>> targetPsiClasses,
+                                     List<Class<?>> acceptingPsiTypes) {
+    if (acceptingPsiTypes == InspectionVisitorsOptimizer.ALL_ELEMENTS_VISIT_LIST) {
       for (int i = 0, elementsSize = elements.size(); i < elementsSize; i++) {
         PsiElement element = elements.get(i);
         element.accept(elementVisitor);
         ProgressManager.checkCanceled();
       }
     }
+    else {
+      Set<Class<?>> accepts = InspectionVisitorsOptimizer.getVisitorAcceptClasses(targetPsiClasses, acceptingPsiTypes);
+      if (accepts == null || accepts.isEmpty()) {
+        return; // nothing to visit in this run
+      }
+
+      for (int i = 0, elementsSize = elements.size(); i < elementsSize; i++) {
+        PsiElement element = elements.get(i);
+
+        if (accepts.contains(element.getClass())) {
+          element.accept(elementVisitor);
+          ProgressManager.checkCanceled();
+        }
+      }
+    }
+  }
 
   /**
    * @deprecated use {@link #inspectEx(List, PsiFile, TextRange, TextRange, boolean, boolean, boolean, ProgressIndicator, PairProcessor)}
@@ -265,6 +289,9 @@ public final class InspectionEngine {
     Map<LocalInspectionToolWrapper, List<ProblemDescriptor>> resultDescriptors = new ConcurrentHashMap<>();
     withSession(psiFile, restrictRange, restrictRange, HighlightSeverity.INFORMATION, isOnTheFly, session -> {
       List<LocalInspectionToolWrapper> applicableTools = filterToolsApplicableByLanguage(toolWrappers, elementDialectIds);
+
+      Map<Class<?>, Collection<Class<?>>> targetPsiClasses = InspectionVisitorsOptimizer.getTargetPsiClasses(elements);
+
       Processor<LocalInspectionToolWrapper> processor = toolWrapper -> {
         ProblemsHolder holder = new ProblemsHolder(InspectionManager.getInstance(psiFile.getProject()), psiFile, isOnTheFly){
           @Override
@@ -287,7 +314,7 @@ public final class InspectionEngine {
         LocalInspectionTool tool = toolWrapper.getTool();
 
         long inspectionStartTime = System.nanoTime();
-        boolean inspectionWasRun = createVisitorAndAcceptElements(tool, holder, isOnTheFly, session, elements);
+        boolean inspectionWasRun = createVisitorAndAcceptElements(tool, holder, isOnTheFly, session, elements, targetPsiClasses);
         long inspectionDuration = TimeoutUtil.getDurationMillis(inspectionStartTime);
 
         reportToQodana(psiFile, isOnTheFly, toolWrapper, inspectionWasRun, inspectionDuration, holder.getResultCount());

platform/analysis-impl/src/com/intellij/codeInspection/InspectionVisitorsOptimizer.java

@@ -0,0 +1,201 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInspection;
+
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.util.UserDataHolderBase;
+import com.intellij.openapi.util.UserDataHolderEx;
+import com.intellij.openapi.util.registry.Registry;
+import com.intellij.psi.BasicInspectionVisitorBean;
+import com.intellij.psi.HintedPsiElementVisitor;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiElementVisitor;
+import com.intellij.psi.impl.ElementBase;
+import com.intellij.psi.impl.source.tree.CompositePsiElement;
+import com.intellij.util.containers.CollectionFactory;
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
+import java.util.*;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static java.util.Collections.*;
+
+/**
+ * Infers classes of elements for inspection visitors in order to skip some of PSI elements during inspection pass.
+ * <p>
+ * Declare `inspection.basicVisitor` in plugin.xml for your language to get speed up of inspection runs.
+ */
+@ApiStatus.Internal
+public final class InspectionVisitorsOptimizer {
+  private InspectionVisitorsOptimizer() {
+  }
+
+  private static final Logger LOG = Logger.getInstance(InspectionVisitorsOptimizer.class);
+
+  public static final List<Class<?>> ALL_ELEMENTS_VISIT_LIST = singletonList(PsiElement.class);
+
+  private static final boolean useOptimizedVisitors = Registry.is("ide.optimize.inspection.visitors");
+  private static final boolean inTests = ApplicationManager.getApplication().isUnitTestMode();
+
+  public static @NotNull List<Class<?>> getAcceptingPsiTypes(@NotNull PsiElementVisitor visitor) {
+    if (!useOptimizedVisitors) return ALL_ELEMENTS_VISIT_LIST;
+
+    List<Class<?>> acceptingPsiTypes;
+    if (visitor instanceof HintedPsiElementVisitor) {
+      acceptingPsiTypes = ((HintedPsiElementVisitor)visitor).getHintPsiElements();
+
+      if (inTests && !VISITOR_TYPES.get(visitor.getClass()).overridesVisitPsiElement) {
+        LOG.error("HintedPsiElementVisitor implementations must override PsiElementVisitor.visitElement",
+                  visitor.getClass().getName());
+      }
+
+      if (acceptingPsiTypes.contains(PsiElement.class) || acceptingPsiTypes.isEmpty()) {
+        acceptingPsiTypes = ALL_ELEMENTS_VISIT_LIST;
+      }
+    }
+    else {
+      acceptingPsiTypes = VISITOR_TYPES.get(visitor.getClass()).handlesElementTypes;
+    }
+
+    return acceptingPsiTypes;
+  }
+
+  public static @NotNull Map<Class<?>, Collection<Class<?>>> getTargetPsiClasses(@NotNull List<? extends PsiElement> elements) {
+    if (!useOptimizedVisitors) return emptyMap();
+
+    Set<Class<?>> uniqueElementClasses = CollectionFactory.createSmallMemoryFootprintSet();
+    for (int i = 0; i < elements.size(); i++) {
+      PsiElement element = elements.get(i);
+      uniqueElementClasses.add(element.getClass());
+    }
+
+    Map<Class<?>, Collection<Class<?>>> targetPsiClasses = new IdentityHashMap<>();
+    for (Class<?> elementClass : uniqueElementClasses) {
+      for (Class<?> aSuper : ELEMENT_TYPE_SUPERS.get(elementClass)) {
+        Collection<Class<?>> classes = targetPsiClasses.get(aSuper);
+        if (classes == null) {
+          classes = CollectionFactory.createSmallMemoryFootprintSet();
+          targetPsiClasses.put(aSuper, classes);
+          if (!aSuper.isInterface() && !Modifier.isAbstract(aSuper.getModifiers())) { // PSI elements in tree cannot be abstract
+            classes.add(aSuper);
+          }
+        }
+
+        classes.add(elementClass);
+      }
+    }
+    return targetPsiClasses;
+  }
+
+  public static @Nullable Set<Class<?>> getVisitorAcceptClasses(
+    @NotNull Map<Class<?>, Collection<Class<?>>> targetPsiClasses,
+    @NotNull List<Class<?>> acceptingPsiTypes
+  ) {
+    if (acceptingPsiTypes.size() == 1) {
+      return Set.copyOf(targetPsiClasses.getOrDefault(acceptingPsiTypes.get(0), emptyList()));
+    }
+
+    Set<Class<?>> accepts = null;
+    for (Class<?> psiType : acceptingPsiTypes) {
+      Collection<Class<?>> classes = targetPsiClasses.getOrDefault(psiType, emptyList());
+      if (!classes.isEmpty()) {
+        if (accepts == null) {
+          accepts = new HashSet<>(classes);
+        }
+        accepts.addAll(classes);
+      }
+    }
+
+    return accepts;
+  }
+
+  private static final ClassValue<Collection<Class<?>>> ELEMENT_TYPE_SUPERS = new ClassValue<>() {
+    @Override
+    protected Collection<Class<?>> computeValue(Class<?> type) {
+      return getAllSupers(type);
+    }
+
+    private static @NotNull Collection<Class<?>> getAllSupers(@NotNull Class<?> clazz) {
+      List<Class<?>> supers = new ArrayList<>();
+      supers.add(clazz);
+      addInterfaces(clazz, supers);
+
+      Class<?> superClass = clazz.getSuperclass();
+      while (superClass != null) {
+        if (superClass != Object.class) {
+          supers.add(superClass);
+          addInterfaces(superClass, supers);
+        }
+        superClass = superClass.getSuperclass();
+      }
+
+      supers.removeIf(aSuper -> {
+        return (aSuper == UserDataHolderBase.class
+                || aSuper == UserDataHolderEx.class
+                || aSuper == CompositePsiElement.class
+                || aSuper == ElementBase.class
+                || aSuper == Cloneable.class
+                || aSuper == AtomicReference.class);
+      });
+
+      return supers;
+    }
+
+    private static void addInterfaces(Class<?> clazz, List<Class<?>> supers) {
+      Class<?>[] interfaces = clazz.getInterfaces();
+      addAll(supers, interfaces);
+      for (Class<?> anInterface : interfaces) {
+        supers.addAll(getAllSupers(anInterface));
+      }
+    }
+  };
+
+  private record VisitorTypes(boolean hasBasicVisitor,
+                              List<Class<?>> handlesElementTypes,
+                              boolean overridesVisitPsiElement) {
+  }
+
+  private static final ClassValue<VisitorTypes> VISITOR_TYPES = new ClassValue<>() {
+    @Override
+    protected VisitorTypes computeValue(Class<?> type) {
+      List<Class<?>> visitClasses = new ArrayList<>();
+
+      Collection<String> visitorClasses = BasicInspectionVisitorBean.getVisitorClasses();
+
+      Class<?> superClass = type;
+      while (superClass != null) {
+        if (superClass == PsiElementVisitor.class) {
+          // no `inspection.basicVisitor` defined in hierarchy
+          return new VisitorTypes(false, ALL_ELEMENTS_VISIT_LIST, visitClasses.contains(PsiElement.class));
+        }
+
+        if (visitorClasses.contains(superClass.getName())) {
+          break;
+        }
+
+        for (Method declaredMethod : superClass.getDeclaredMethods()) {
+          if (declaredMethod.getParameterCount() == 1
+              && declaredMethod.getName().startsWith("visit")
+              && Modifier.isPublic(declaredMethod.getModifiers())
+              && !Modifier.isAbstract(declaredMethod.getModifiers())
+              && !Modifier.isStatic(declaredMethod.getModifiers())) {
+            Class<?> parameterType = declaredMethod.getParameterTypes()[0];
+            visitClasses.add(parameterType);
+          }
+        }
+
+        superClass = superClass.getSuperclass();
+      }
+
+      if (visitClasses.contains(PsiElement.class)) {
+        return new VisitorTypes(true, ALL_ELEMENTS_VISIT_LIST, true);
+      }
+
+      return new VisitorTypes(true, List.copyOf(visitClasses), false);
+    }
+  };
+}

platform/core-api/resources/META-INF/Core.xml

@@ -52,6 +52,7 @@
       <with attribute="implementationClass" implements="com.intellij.psi.LanguageSubstitutor"/>
     </extensionPoint>
     <extensionPoint name="iconProvider" interface="com.intellij.ide.IconProvider" dynamic="true"/>
+    <extensionPoint name="inspection.basicVisitor" beanClass="com.intellij.psi.BasicInspectionVisitorBean" dynamic="true"/>
   </extensionPoints>
   <extensions defaultExtensionNs="com.intellij">
     <applicationService serviceInterface="com.intellij.util.messages.MessageBusFactory"

platform/core-api/src/com/intellij/psi/BasicInspectionVisitorBean.java

@@ -0,0 +1,36 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.psi;
+
+import com.intellij.openapi.extensions.ExtensionPointName;
+import com.intellij.openapi.extensions.RequiredElement;
+import com.intellij.util.xmlb.annotations.Attribute;
+import org.jetbrains.annotations.ApiStatus;
+
+import java.util.Collection;
+import java.util.Set;
+
+import static java.util.stream.Collectors.toSet;
+
+@ApiStatus.Internal
+public final class BasicInspectionVisitorBean {
+  @Attribute("class")
+  @RequiredElement
+  public String clazz;
+
+  private static final ExtensionPointName<BasicInspectionVisitorBean> EP_NAME =
+    ExtensionPointName.create("com.intellij.inspection.basicVisitor");
+
+  private static volatile Set<String> ourClasses;
+
+  public static Collection<String> getVisitorClasses() {
+    Set<String> set = ourClasses;
+    if (set != null) return set;
+
+    set = EP_NAME.getExtensionList().stream()
+      .map(x -> x.clazz)
+      .collect(toSet());
+    ourClasses = set;
+
+    return set;
+  }
+}

platform/core-api/src/com/intellij/psi/HintedPsiElementVisitor.java

@@ -0,0 +1,25 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.psi;
+
+import org.jetbrains.annotations.ApiStatus;
+
+import java.util.List;
+
+/**
+ * {@link PsiElementVisitor} that exposes desired element classes to visit. Called once per each run of inspection tool.
+ * Inspection engine then may skip elements with all types not in this list.
+ * <p>
+ * Limitations:
+ * <ul>
+ * <li>It must always return the same set of classes</li>
+ * <li>The result may not depend on any configuration/settings</li>
+ * <li>Implementations must override {@link PsiElementVisitor#visitElement(PsiElement)}</li>
+ * </ul>
+ */
+@ApiStatus.Internal
+public interface HintedPsiElementVisitor {
+  /**
+   * @return PSI element classes to visit
+   */
+  List<Class<?>> getHintPsiElements();
+}

platform/core-api/src/com/intellij/psi/PsiLanguageInjectionHostVisitor.java

@@ -0,0 +1,13 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.psi;
+
+import java.util.List;
+
+import static java.util.Collections.singletonList;
+
+public abstract class PsiLanguageInjectionHostVisitor extends PsiElementVisitor implements HintedPsiElementVisitor {
+  @Override
+  public List<Class<?>> getHintPsiElements() {
+    return singletonList(PsiLanguageInjectionHost.class);
+  }
+}

platform/lang-impl/src/com/intellij/codeInsight/daemon/impl/InspectionRunner.java

@@ -44,6 +44,14 @@ import java.util.function.BiPredicate;
 import java.util.function.Consumer;
 
 class InspectionRunner {
+
+  private static final PsiElementVisitor TOMB_STONE_VISITOR = new PsiElementVisitor() {
+    @Override
+    public String toString() {
+      return "TOMB_STONE_VISITOR";
+    }
+  };
+
   private final PsiFile myPsiFile;
   private final TextRange myRestrictRange;
   private final TextRange myPriorityRange;
@@ -81,11 +89,13 @@ class InspectionRunner {
       this.tool = tool;
       this.holder = holder;
       this.visitor = visitor;
+      this.acceptingPsiTypes = InspectionVisitorsOptimizer.getAcceptingPsiTypes(visitor);
     }
 
     final @NotNull LocalInspectionToolWrapper tool;
     final @NotNull InspectionProblemHolder holder;
     final @NotNull PsiElementVisitor visitor;
+    final @NotNull List<Class<?>> acceptingPsiTypes;
     volatile PsiElement myFavoriteElement; // the element during visiting which, some diagnostics were generated in the previous run
   }
 
@@ -178,7 +188,7 @@ class InspectionRunner {
   private @NotNull InspectionContext createTombStone() {
     LocalInspectionToolWrapper tool = new LocalInspectionToolWrapper(new LocalInspectionEP());
     InspectionProblemHolder holder = new InspectionProblemHolder(myPsiFile, tool, false, myInspectionProfileWrapper, empty);
-    return new InspectionContext(tool, holder, new PsiElementVisitor() {});
+    return new InspectionContext(tool, holder, TOMB_STONE_VISITOR);
   }
 
   private static @NotNull <T> Map<PsiFile, T> createInjectedFileMap() {
@@ -308,6 +318,8 @@ class InspectionRunner {
     ApplicationEx application = ApplicationManagerEx.getApplicationEx();
     boolean shouldFailFastAcquiringReadAction = application.isInImpatientReader();
 
+    Map<Class<?>, Collection<Class<?>>> targetPsiClasses = InspectionVisitorsOptimizer.getTargetPsiClasses(elements);
+
     boolean processed =
       ((JobLauncherImpl)JobLauncher.getInstance()).processQueue(contexts, new LinkedBlockingQueue<>(), new SensitiveProgressWrapper(myProgress), TOMB_STONE, context ->
         AstLoadingFilter.disallowTreeLoading(() -> AstLoadingFilter.<Boolean, RuntimeException>forceAllowTreeLoading(myPsiFile, () -> {
@@ -328,17 +340,40 @@ class InspectionRunner {
                 }
               }
 
-              //noinspection ForLoopReplaceableByForEach
+              if (context.acceptingPsiTypes == InspectionVisitorsOptimizer.ALL_ELEMENTS_VISIT_LIST) {
                 for (int i = 0; i < elements.size(); i++) {
                   PsiElement element = elements.get(i);
-                ProgressManager.checkCanceled();
+
                   if (element == favoriteElement) continue; // already visited
+
+                  ProgressManager.checkCanceled();
                   element.accept(context.visitor);
                   if (resultCount != -1 && holder.getResultCount() != resultCount) {
                     context.myFavoriteElement = element;
                     resultCount = -1; // mark as "new favorite element is stored"
                   }
                 }
+              }
+              else {
+                Set<Class<?>> accepts = InspectionVisitorsOptimizer.getVisitorAcceptClasses(targetPsiClasses, context.acceptingPsiTypes);
+                if (accepts != null && !accepts.isEmpty()) {
+                  for (int i = 0; i < elements.size(); i++) {
+                    PsiElement element = elements.get(i);
+
+                    if (element == favoriteElement) continue; // already visited
+
+                    if (accepts.contains(element.getClass())) {
+                      ProgressManager.checkCanceled();
+                      element.accept(context.visitor);
+                      if (resultCount != -1 && holder.getResultCount() != resultCount) {
+                        context.myFavoriteElement = element;
+                        resultCount = -1; // mark as "new favorite element is stored"
+                      }
+                    }
+                  }
+                }
+              }
+
               afterProcessCallback.accept(context);
             };
             if (!application.tryRunReadAction(action)) {

platform/util/resources/misc/registry.properties

@@ -1003,6 +1003,9 @@ ide.goto.symbol.include.overrides.on.qualified.patterns.description=Disable dedu
 ide.structural.navigation.visit.fields=false
 ide.structural.navigation.visit.fields.description=Whether fields should be stopped at when navigating to the nex/previous structural member by Alt+Down/Up.
 
+ide.optimize.inspection.visitors=false
+ide.optimize.inspection.visitors.description=Whether to skip elements not suitable for visitor during inspections run
+
 ide.dfa.time.limit.online=1000
 ide.dfa.time.limit.online.description=Time limit (in milliseconds) that is allowed to analyze data flow for one method