[devkit] extract DevKitJvmInspection.ForClass

GitOrigin-RevId: ee8ba81a55f4e86b44d1642e0604c710c73851a0

plugins/devkit/devkit-core/src/inspections/ComponentNotRegisteredInspection.java

@@ -3,9 +3,6 @@ package org.jetbrains.idea.devkit.inspections;
 
 import com.intellij.codeInspection.LocalQuickFix;
 import com.intellij.codeInspection.options.OptPane;
-import com.intellij.lang.jvm.DefaultJvmElementVisitor;
-import com.intellij.lang.jvm.JvmClass;
-import com.intellij.lang.jvm.JvmElementVisitor;
 import com.intellij.openapi.actionSystem.AnAction;
 import com.intellij.openapi.components.BaseComponent;
 import com.intellij.openapi.diagnostic.Logger;
@@ -22,7 +19,6 @@ import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.Query;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.idea.devkit.DevKitBundle;
 import org.jetbrains.idea.devkit.inspections.quickfix.RegisterActionFix;
 import org.jetbrains.idea.devkit.inspections.quickfix.RegisterComponentFix;
@@ -35,7 +31,7 @@ import java.util.Set;
 import static com.intellij.codeInspection.options.OptPane.checkbox;
 import static com.intellij.codeInspection.options.OptPane.pane;
 
-final class ComponentNotRegisteredInspection extends DevKitJvmInspection {
+final class ComponentNotRegisteredInspection extends DevKitJvmInspection.ForClass {
   private static final Logger LOG = Logger.getInstance(ComponentNotRegisteredInspection.class);
 
   public boolean CHECK_ACTIONS = true;
@@ -55,23 +51,8 @@ final class ComponentNotRegisteredInspection extends DevKitJvmInspection {
     );
   }
 
-  @Nullable
   @Override
-  protected JvmElementVisitor<Boolean> buildVisitor(@NotNull Project project, @NotNull HighlightSink sink, boolean isOnTheFly) {
-    return new DefaultJvmElementVisitor<>() {
-      @Override
-      public Boolean visitClass(@NotNull JvmClass clazz) {
-        PsiElement sourceElement = clazz.getSourceElement();
-        if (!(sourceElement instanceof PsiClass)) {
-          return null;
-        }
-        checkClass(project, (PsiClass)sourceElement, sink);
-        return false;
-      }
-    };
-  }
-
-  private void checkClass(@NotNull Project project, @NotNull PsiClass checkedClass, @NotNull HighlightSink sink) {
+  protected void checkClass(@NotNull Project project, @NotNull PsiClass checkedClass, @NotNull HighlightSink sink) {
     if (checkedClass.getQualifiedName() == null ||
         checkedClass.getContainingFile().getVirtualFile() == null ||
         checkedClass.hasModifierProperty(PsiModifier.ABSTRACT) ||
@@ -119,7 +100,8 @@ final class ComponentNotRegisteredInspection extends DevKitJvmInspection {
     }
 
     for (ComponentType componentType : ComponentType.values()) {
-      if (InheritanceUtil.isInheritor(checkedClass, componentType.myClassName) && checkComponentRegistration(checkedClass, sink, componentType)) {
+      if (InheritanceUtil.isInheritor(checkedClass, componentType.myClassName) &&
+          checkComponentRegistration(checkedClass, sink, componentType)) {
         return;
       }
     }

plugins/devkit/devkit-core/src/inspections/DescriptionNotFoundInspectionBase.java

@@ -2,9 +2,6 @@
 package org.jetbrains.idea.devkit.inspections;
 
 import com.intellij.codeInspection.util.InspectionMessage;
-import com.intellij.lang.jvm.DefaultJvmElementVisitor;
-import com.intellij.lang.jvm.JvmClass;
-import com.intellij.lang.jvm.JvmElementVisitor;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.module.ModuleUtilCore;
 import com.intellij.openapi.project.Project;
@@ -12,13 +9,12 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiDirectory;
-import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.idea.devkit.inspections.quickfix.CreateHtmlDescriptionFix;
 
-abstract class DescriptionNotFoundInspectionBase extends DevKitJvmInspection {
+abstract class DescriptionNotFoundInspectionBase extends DevKitJvmInspection.ForClass {
 
   private final DescriptionType myDescriptionType;
 
@@ -26,21 +22,8 @@ abstract class DescriptionNotFoundInspectionBase extends DevKitJvmInspection {
     myDescriptionType = descriptionType;
   }
 
-  protected JvmElementVisitor<Boolean> buildVisitor(@NotNull Project project, @NotNull HighlightSink sink, boolean isOnTheFly) {
-    return new DefaultJvmElementVisitor<>() {
-      @Override
-      public Boolean visitClass(@NotNull JvmClass clazz) {
-        PsiElement sourceElement = clazz.getSourceElement();
-        if (!(sourceElement instanceof PsiClass)) {
-          return null;
-        }
-        checkClass((PsiClass)sourceElement, sink);
-        return false;
-      }
-    };
-  }
-
-  private void checkClass(@NotNull PsiClass psiClass, @NotNull HighlightSink sink) {
+  @Override
+  protected void checkClass(@NotNull Project project, @NotNull PsiClass psiClass, @NotNull HighlightSink sink) {
     if (!ExtensionUtil.isExtensionPointImplementationCandidate(psiClass)) return;
     if (!myDescriptionType.matches(psiClass)) return;
 

plugins/devkit/devkit-core/src/inspections/DevKitJvmInspection.kt

@@ -2,11 +2,17 @@
 package org.jetbrains.idea.devkit.inspections
 
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.lang.jvm.DefaultJvmElementVisitor
+import com.intellij.lang.jvm.JvmClass
+import com.intellij.lang.jvm.JvmElementVisitor
 import com.intellij.lang.jvm.inspection.JvmLocalInspection
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElementVisitor
 
 /**
  * Base class for JVM API (declaration-based) DevKit inspections.
+ * Use [ForClass] when only inspecting classes.
  *
  * Override [isAllowed] to add additional constraints (e.g., required class in scope via [DevKitInspectionUtil.isClassAvailable])
  * to skip running inspection completely whenever possible.
@@ -23,4 +29,21 @@ abstract class DevKitJvmInspection : JvmLocalInspection() {
   }
 
   protected open fun isAllowed(holder: ProblemsHolder) = DevKitInspectionUtil.isAllowed(holder.file)
+
+
+  abstract class ForClass : DevKitJvmInspection() {
+
+    final override fun buildVisitor(project: Project, sink: HighlightSink, isOnTheFly: Boolean): JvmElementVisitor<Boolean?>? {
+      return object : DefaultJvmElementVisitor<Boolean> {
+        override fun visitClass(clazz: JvmClass): Boolean {
+          val sourceElement = clazz.sourceElement
+          if (sourceElement !is PsiClass) return true
+          checkClass(project, sourceElement, sink)
+          return true
+        }
+      }
+    }
+
+    protected abstract fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink)
+  }
 }

plugins/devkit/devkit-core/src/inspections/ExtensionClassShouldBeFinalAndNonPublicInspection.kt

@@ -2,9 +2,7 @@
 package org.jetbrains.idea.devkit.inspections
 
 import com.intellij.codeInspection.IntentionWrapper
-import com.intellij.lang.jvm.DefaultJvmElementVisitor
 import com.intellij.lang.jvm.JvmClass
-import com.intellij.lang.jvm.JvmElementVisitor
 import com.intellij.lang.jvm.JvmModifier
 import com.intellij.lang.jvm.actions.createModifierActions
 import com.intellij.lang.jvm.actions.modifierRequest
@@ -25,45 +23,40 @@ private inline val visibleForTestingAnnotations
     "org.jetbrains.annotations.VisibleForTesting"
   )
 
-internal class ExtensionClassShouldBeFinalAndNonPublicInspection : DevKitJvmInspection() {
+internal class ExtensionClassShouldBeFinalAndNonPublicInspection : DevKitJvmInspection.ForClass() {
 
-  override fun buildVisitor(project: Project, sink: HighlightSink, isOnTheFly: Boolean): JvmElementVisitor<Boolean> {
-    return object : DefaultJvmElementVisitor<Boolean> {
-      override fun visitClass(clazz: JvmClass): Boolean {
-        if (clazz !is PsiClass) return true
-        if (!ExtensionUtil.isExtensionPointImplementationCandidate(clazz)) {
-          return true
-        }
-        val sourceElement = clazz.sourceElement
-        val language = sourceElement?.language ?: return true
-        val file = clazz.containingFile ?: return true
+  override fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink) {
+    if (!ExtensionUtil.isExtensionPointImplementationCandidate(psiClass)) {
+      return
+    }
+    val sourceElement = psiClass.sourceElement
+    val language = sourceElement?.language ?: return
+    val file = psiClass.containingFile ?: return
 
-        val isFinal = clazz.hasModifier(JvmModifier.FINAL)
-        val extensionClassShouldNotBePublicProvider = getProvider(ExtensionClassShouldNotBePublicProviders, language) ?: return true
-        val isPublic = extensionClassShouldNotBePublicProvider.isPublic(clazz)
-        if (isFinal && !isPublic) return true
+    val isFinal = psiClass.hasModifier(JvmModifier.FINAL)
+    val extensionClassShouldNotBePublicProvider = getProvider(ExtensionClassShouldNotBePublicProviders, language) ?: return
+    val isPublic = extensionClassShouldNotBePublicProvider.isPublic(psiClass)
+    if (isFinal && !isPublic) return
 
-        if (!ExtensionUtil.isInstantiatedExtension(clazz) { false }) return true
+    if (!ExtensionUtil.isInstantiatedExtension(psiClass) { false }) return
 
-        if (!isFinal && !hasInheritors(clazz)) {
-          val actions = createModifierActions(clazz, modifierRequest(JvmModifier.FINAL, true))
-          val errorMessageProvider = getProvider(ExtensionClassShouldBeFinalErrorMessageProviders, language) ?: return true
-          val message = errorMessageProvider.provideErrorMessage()
-          val fixes = IntentionWrapper.wrapToQuickFixes(actions.toTypedArray(), file)
-          sink.highlight(message, *fixes)
-        }
-        if (isPublic && !isAnnotatedAsVisibleForTesting(clazz)) {
-          val message = when {
-            isServiceImplementationRegisteredInPluginXml(clazz) -> DevKitBundle.message("inspection.extension.class.should.not.be.public.service")
-            else -> DevKitBundle.message("inspection.extension.class.should.not.be.public.text")
-          }
-          val fixes = extensionClassShouldNotBePublicProvider.provideQuickFix(clazz, file)
-          sink.highlight(message, *fixes)
-        }
-        return true
+    if (!isFinal && !hasInheritors(psiClass)) {
+      val actions = createModifierActions(psiClass, modifierRequest(JvmModifier.FINAL, true))
+      val errorMessageProvider = getProvider(ExtensionClassShouldBeFinalErrorMessageProviders, language) ?: return
+      val message = errorMessageProvider.provideErrorMessage()
+      val fixes = IntentionWrapper.wrapToQuickFixes(actions.toTypedArray(), file)
+      sink.highlight(message, *fixes)
+    }
+    if (isPublic && !isAnnotatedAsVisibleForTesting(psiClass)) {
+      val message = when {
+        isServiceImplementationRegisteredInPluginXml(psiClass) -> DevKitBundle.message("inspection.extension.class.should.not.be.public.service")
+        else -> DevKitBundle.message("inspection.extension.class.should.not.be.public.text")
       }
+      val fixes = extensionClassShouldNotBePublicProvider.provideQuickFix(psiClass, file)
+      sink.highlight(message, *fixes)
     }
   }
+
 }
 
 private fun isServiceImplementationRegisteredInPluginXml(aClass: PsiClass): Boolean {
@@ -84,6 +77,6 @@ private fun hasInheritors(aClass: PsiClass): Boolean {
   return DirectClassInheritorsSearch.search(aClass).findFirst() != null
 }
 
-private fun isAnnotatedAsVisibleForTesting(clazz: JvmClass): Boolean {
-  return clazz.annotations.any { visibleForTestingAnnotations.contains(it.qualifiedName) }
+private fun isAnnotatedAsVisibleForTesting(psiClass: JvmClass): Boolean {
+  return psiClass.annotations.any { visibleForTestingAnnotations.contains(it.qualifiedName) }
 }

plugins/devkit/devkit-core/src/inspections/ExtensionRegisteredAsServiceOrComponentInspection.kt

@@ -1,9 +1,6 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.idea.devkit.inspections
 
-import com.intellij.lang.jvm.DefaultJvmElementVisitor
-import com.intellij.lang.jvm.JvmClass
-import com.intellij.lang.jvm.JvmElementVisitor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiClass
 import com.intellij.psi.util.InheritanceUtil
@@ -13,22 +10,11 @@ import org.jetbrains.idea.devkit.DevKitBundle
 import org.jetbrains.idea.devkit.dom.Extension
 import org.jetbrains.idea.devkit.util.locateExtensionsByPsiClass
 
-internal class ExtensionRegisteredAsServiceOrComponentInspection : DevKitJvmInspection() {
+internal class ExtensionRegisteredAsServiceOrComponentInspection : DevKitJvmInspection.ForClass() {
 
   private val serviceAttributeNames = setOf("service")
 
-  override fun buildVisitor(project: Project, sink: HighlightSink, isOnTheFly: Boolean): JvmElementVisitor<Boolean?>? {
-    return object : DefaultJvmElementVisitor<Boolean> {
-      override fun visitClass(clazz: JvmClass): Boolean {
-        val sourceElement = clazz.sourceElement
-        if (sourceElement !is PsiClass) return true
-        checkClass(project, sourceElement, sink)
-        return true
-      }
-    }
-  }
-
-  private fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink) {
+  override fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink) {
     if (!ExtensionUtil.isExtensionPointImplementationCandidate(psiClass)) {
       return
     }

plugins/devkit/devkit-core/src/inspections/JComponentDataProviderInspection.kt

@@ -2,10 +2,7 @@
 package org.jetbrains.idea.devkit.inspections
 
 import com.intellij.codeInspection.ProblemsHolder
-import com.intellij.lang.jvm.DefaultJvmElementVisitor
-import com.intellij.lang.jvm.JvmClass
 import com.intellij.lang.jvm.JvmClassKind
-import com.intellij.lang.jvm.JvmElementVisitor
 import com.intellij.openapi.actionSystem.DataProvider
 import com.intellij.openapi.actionSystem.UiCompatibleDataProvider
 import com.intellij.openapi.actionSystem.UiDataProvider
@@ -15,25 +12,14 @@ import com.intellij.psi.util.InheritanceUtil
 import org.jetbrains.idea.devkit.DevKitBundle
 import javax.swing.JComponent
 
-internal class JComponentDataProviderInspection : DevKitJvmInspection() {
+internal class JComponentDataProviderInspection : DevKitJvmInspection.ForClass() {
 
   override fun isAllowed(holder: ProblemsHolder): Boolean {
     return super.isAllowed(holder) &&
            DevKitInspectionUtil.isClassAvailable(holder, UiDataProvider::class.java.name)
   }
 
-  override fun buildVisitor(project: Project, sink: HighlightSink, isOnTheFly: Boolean): JvmElementVisitor<Boolean?>? {
-    return object : DefaultJvmElementVisitor<Boolean> {
-      override fun visitClass(clazz: JvmClass): Boolean {
-        val sourceElement = clazz.sourceElement
-        if (sourceElement !is PsiClass) return true
-        checkClass(sourceElement, sink)
-        return true
-      }
-    }
-  }
-
-  private fun checkClass(psiClass: PsiClass, sink: HighlightSink) {
+  override fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink) {
     if (psiClass.classKind != JvmClassKind.CLASS) return
 
     @Suppress("UsagesOfObsoleteApi")

plugins/devkit/devkit-core/src/inspections/LightServiceMustBeFinalInspection.kt

@@ -3,9 +3,6 @@ package org.jetbrains.idea.devkit.inspections
 
 import com.intellij.codeInspection.IntentionWrapper
 import com.intellij.codeInspection.ProblemHighlightType
-import com.intellij.lang.jvm.DefaultJvmElementVisitor
-import com.intellij.lang.jvm.JvmClass
-import com.intellij.lang.jvm.JvmElementVisitor
 import com.intellij.lang.jvm.JvmModifier
 import com.intellij.lang.jvm.actions.annotationRequest
 import com.intellij.lang.jvm.actions.createModifierActions
@@ -16,33 +13,26 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiClass
 import org.jetbrains.idea.devkit.DevKitBundle
 
-internal class LightServiceMustBeFinalInspection : DevKitJvmInspection() {
+internal class LightServiceMustBeFinalInspection : DevKitJvmInspection.ForClass() {
 
-  override fun buildVisitor(project: Project, sink: HighlightSink, isOnTheFly: Boolean): JvmElementVisitor<Boolean> {
-    return object : DefaultJvmElementVisitor<Boolean> {
-      override fun visitClass(clazz: JvmClass): Boolean {
-        val sourceElement = clazz.sourceElement
-        if (sourceElement !is PsiClass) return true
-        if (sourceElement.isAnnotationType || sourceElement.isEnum || sourceElement.hasModifier(JvmModifier.FINAL)) return true
-        val file = sourceElement.containingFile ?: return true
-        val serviceAnnotation = sourceElement.getAnnotation(Service::class.java.canonicalName) ?: return true
-        val elementToReport = serviceAnnotation.nameReferenceElement ?: return true
-        if (sourceElement.isInterface || sourceElement.hasModifier(JvmModifier.ABSTRACT)) {
-          val actions = createRemoveAnnotationActions(sourceElement, annotationRequest(Service::class.java.canonicalName))
-          val fixes = IntentionWrapper.wrapToQuickFixes(actions.toTypedArray(), file)
-          val message = DevKitBundle.message("inspection.light.service.must.be.concrete.class.message")
-          val holder = (sink as HighlightSinkImpl).holder
-          holder.registerProblem(elementToReport, message, ProblemHighlightType.GENERIC_ERROR, *fixes)
-        }
-        else {
-          val errorMessageProvider = getProvider(LightServiceMustBeFinalErrorMessageProviders, sourceElement.language) ?: return true
-          val message = errorMessageProvider.provideErrorMessage()
-          val actions = createModifierActions(sourceElement, modifierRequest(JvmModifier.FINAL, true))
-          val fixes = IntentionWrapper.wrapToQuickFixes(actions.toTypedArray(), file)
-          sink.highlight(message, ProblemHighlightType.GENERIC_ERROR, *fixes)
-        }
-        return true
-      }
+  override fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink) {
+    if (psiClass.isAnnotationType || psiClass.isEnum || psiClass.hasModifier(JvmModifier.FINAL)) return
+    val file = psiClass.containingFile ?: return
+    val serviceAnnotation = psiClass.getAnnotation(Service::class.java.canonicalName) ?: return
+    val elementToReport = serviceAnnotation.nameReferenceElement ?: return
+    if (psiClass.isInterface || psiClass.hasModifier(JvmModifier.ABSTRACT)) {
+      val actions = createRemoveAnnotationActions(psiClass, annotationRequest(Service::class.java.canonicalName))
+      val fixes = IntentionWrapper.wrapToQuickFixes(actions.toTypedArray(), file)
+      val message = DevKitBundle.message("inspection.light.service.must.be.concrete.class.message")
+      val holder = (sink as HighlightSinkImpl).holder
+      holder.registerProblem(elementToReport, message, ProblemHighlightType.GENERIC_ERROR, *fixes)
+      return
     }
+
+    val errorMessageProvider = getProvider(LightServiceMustBeFinalErrorMessageProviders, psiClass.language) ?: return
+    val message = errorMessageProvider.provideErrorMessage()
+    val actions = createModifierActions(psiClass, modifierRequest(JvmModifier.FINAL, true))
+    val fixes = IntentionWrapper.wrapToQuickFixes(actions.toTypedArray(), file)
+    sink.highlight(message, ProblemHighlightType.GENERIC_ERROR, *fixes)
   }
 }

plugins/devkit/devkit-core/src/inspections/ListenerImplementationMustNotBeDisposableInspection.kt

@@ -1,9 +1,6 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.idea.devkit.inspections
 
-import com.intellij.lang.jvm.DefaultJvmElementVisitor
-import com.intellij.lang.jvm.JvmClass
-import com.intellij.lang.jvm.JvmElementVisitor
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiClass
@@ -12,20 +9,9 @@ import org.jetbrains.idea.devkit.DevKitBundle
 import org.jetbrains.idea.devkit.dom.index.IdeaPluginRegistrationIndex
 import org.jetbrains.idea.devkit.util.PluginRelatedLocatorsUtils
 
-internal class ListenerImplementationMustNotBeDisposableInspection : DevKitJvmInspection() {
+internal class ListenerImplementationMustNotBeDisposableInspection : DevKitJvmInspection.ForClass() {
 
-  override fun buildVisitor(project: Project, sink: HighlightSink, isOnTheFly: Boolean): JvmElementVisitor<Boolean?>? {
-    return object : DefaultJvmElementVisitor<Boolean> {
-      override fun visitClass(clazz: JvmClass): Boolean {
-        val sourceElement = clazz.sourceElement
-        if (sourceElement !is PsiClass) return true
-        checkClass(project, sourceElement, sink)
-        return true
-      }
-    }
-  }
-
-  private fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink) {
+  override fun checkClass(project: Project, psiClass: PsiClass, sink: HighlightSink) {
     if (!ExtensionUtil.isExtensionPointImplementationCandidate(psiClass)) return
 
     if (!InheritanceUtil.isInheritor(psiClass, Disposable::class.java.canonicalName)) return

plugins/devkit/devkit-core/src/inspections/MissingActionUpdateThread.java

@@ -11,19 +11,18 @@ import com.intellij.openapi.actionSystem.ActionUpdateThreadAware;
 import com.intellij.openapi.actionSystem.AnAction;
 import com.intellij.openapi.actionSystem.AnActionEvent;
 import com.intellij.openapi.project.Project;
-import com.intellij.openapi.util.IntellijInternalApi;
 import com.intellij.psi.CommonClassNames;
+import com.intellij.psi.PsiClass;
 import com.intellij.util.containers.JBIterable;
 import com.intellij.util.containers.JBTreeTraverser;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.VisibleForTesting;
 import org.jetbrains.idea.devkit.DevKitBundle;
 
 @VisibleForTesting
 @ApiStatus.Internal
-public final class MissingActionUpdateThread extends DevKitJvmInspection {
+public final class MissingActionUpdateThread extends DevKitJvmInspection.ForClass {
 
   @Override
   protected boolean isAllowed(@NotNull ProblemsHolder holder) {
@@ -31,61 +30,55 @@ public final class MissingActionUpdateThread extends DevKitJvmInspection {
            DevKitInspectionUtil.isClassAvailable(holder, ActionUpdateThreadAware.class.getName());
   }
 
-  @Nullable
   @Override
-  protected JvmElementVisitor<Boolean> buildVisitor(@NotNull Project project, @NotNull HighlightSink sink, boolean isOnTheFly) {
-    return new DefaultJvmElementVisitor<>() {
-      @Override
-      public Boolean visitClass(@NotNull JvmClass clazz) {
-        if (clazz.getClassKind() != JvmClassKind.CLASS ||
-            clazz.hasModifier(JvmModifier.ABSTRACT) ||
-            !JvmInheritanceUtil.isInheritor(clazz, ActionUpdateThreadAware.class.getName())) {
-          return false;
-        }
-        boolean isAnAction = false;
-        boolean hasUpdateMethod = false;
-        JBIterable<JvmReferenceType> superInterfaces = JBIterable.empty();
-        for (JvmClass c = clazz; c != null; c = JvmUtil.resolveClass(c.getSuperClassType())) {
-          String className = c.getQualifiedName();
-          if (CommonClassNames.JAVA_LANG_OBJECT.equals(className) ||
-              (isAnAction = AnAction.class.getName().equals(className))) {
-            break;
-          }
-          for (JvmMethod method : c.getMethods()) {
-            String name = method.getName();
-            if ("getActionUpdateThread".equals(name) && method.getParameters().length == 0) {
-              return null;
-            }
-            else if ("update".equals(name) && !hasUpdateMethod) {
-              JvmParameter[] parameters = method.getParameters();
-              JvmType pt = parameters.length == 1 ? parameters[0].getType() : null;
-              JvmClass pc = pt instanceof JvmReferenceType ? JvmUtil.resolveClass((JvmReferenceType)pt) : null;
-              if (pc != null && AnActionEvent.class.getName().equals(pc.getQualifiedName())) {
-                hasUpdateMethod = true;
-              }
-            }
-          }
-          superInterfaces = superInterfaces.append(JBIterable.of(c.getInterfaceTypes()));
-        }
-        if (!isAnAction) {
-          // Check super-interfaces for default methods - no need to check if the method is default.
-          // Default override is good, non-default override without the implementation is a compiler error.
-          JBTreeTraverser<JvmClass> traverser = JBTreeTraverser.from(o -> JBIterable.of(o.getSuperClassType())
-            .filterMap(JvmUtil::resolveClass));
-          for (JvmClass c : traverser.unique().withRoots(superInterfaces.filterMap(JvmUtil::resolveClass))) {
-            if (ActionUpdateThreadAware.class.getName().equals(c.getQualifiedName())) continue;
-            for (JvmMethod method : c.getMethods()) {
-              if ("getActionUpdateThread".equals(method.getName()) && method.getParameters().length == 0) {
-                return null;
-              }
-            }
-          }
-        }
-        if (!isAnAction || hasUpdateMethod) {
-          sink.highlight(DevKitBundle.message("inspections.action.update.thread.message"));
-        }
-        return false;
+  protected void checkClass(@NotNull Project project, @NotNull PsiClass psiClass, @NotNull HighlightSink sink) {
+
+    if (psiClass.getClassKind() != JvmClassKind.CLASS ||
+        psiClass.hasModifier(JvmModifier.ABSTRACT) ||
+        !JvmInheritanceUtil.isInheritor(psiClass, ActionUpdateThreadAware.class.getName())) {
+      return;
+    }
+    boolean isAnAction = false;
+    boolean hasUpdateMethod = false;
+    JBIterable<JvmReferenceType> superInterfaces = JBIterable.empty();
+    for (JvmClass c = psiClass; c != null; c = JvmUtil.resolveClass(c.getSuperClassType())) {
+      String className = c.getQualifiedName();
+      if (CommonClassNames.JAVA_LANG_OBJECT.equals(className) ||
+          (isAnAction = AnAction.class.getName().equals(className))) {
+        break;
       }
-    };
+      for (JvmMethod method : c.getMethods()) {
+        String name = method.getName();
+        if ("getActionUpdateThread".equals(name) && method.getParameters().length == 0) {
+          return;
+        }
+        else if ("update".equals(name) && !hasUpdateMethod) {
+          JvmParameter[] parameters = method.getParameters();
+          JvmType pt = parameters.length == 1 ? parameters[0].getType() : null;
+          JvmClass pc = pt instanceof JvmReferenceType ? JvmUtil.resolveClass((JvmReferenceType)pt) : null;
+          if (pc != null && AnActionEvent.class.getName().equals(pc.getQualifiedName())) {
+            hasUpdateMethod = true;
+          }
+        }
+      }
+      superInterfaces = superInterfaces.append(JBIterable.of(c.getInterfaceTypes()));
+    }
+    if (!isAnAction) {
+      // Check super-interfaces for default methods - no need to check if the method is default.
+      // Default override is good, non-default override without the implementation is a compiler error.
+      JBTreeTraverser<JvmClass> traverser = JBTreeTraverser.from(o -> JBIterable.of(o.getSuperClassType())
+        .filterMap(JvmUtil::resolveClass));
+      for (JvmClass c : traverser.unique().withRoots(superInterfaces.filterMap(JvmUtil::resolveClass))) {
+        if (ActionUpdateThreadAware.class.getName().equals(c.getQualifiedName())) continue;
+        for (JvmMethod method : c.getMethods()) {
+          if ("getActionUpdateThread".equals(method.getName()) && method.getParameters().length == 0) {
+            return;
+          }
+        }
+      }
+    }
+    if (!isAnAction || hasUpdateMethod) {
+      sink.highlight(DevKitBundle.message("inspections.action.update.thread.message"));
+    }
   }
 }

plugins/devkit/devkit-core/src/inspections/internal/StatisticsCollectorNotRegisteredInspection.java

@@ -1,22 +1,14 @@
 // Copyright 2000-2024 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package org.jetbrains.idea.devkit.inspections.internal;
 
-import com.intellij.lang.jvm.DefaultJvmElementVisitor;
-import com.intellij.lang.jvm.JvmClass;
-import com.intellij.lang.jvm.JvmElementVisitor;
 import com.intellij.openapi.project.Project;
-import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiClass;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiModifier;
-import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.util.ClassUtil;
 import com.intellij.psi.util.InheritanceUtil;
 import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.VisibleForTesting;
 import org.jetbrains.idea.devkit.DevKitBundle;
 import org.jetbrains.idea.devkit.dom.ExtensionPoint;
@@ -34,25 +26,11 @@ import static org.jetbrains.idea.devkit.util.ExtensionLocatorKt.processExtension
 
 @VisibleForTesting
 @ApiStatus.Internal
-public final class StatisticsCollectorNotRegisteredInspection extends DevKitJvmInspection {
+public final class StatisticsCollectorNotRegisteredInspection extends DevKitJvmInspection.ForClass {
   private static final String FEATURE_USAGES_COLLECTOR = "com.intellij.internal.statistic.service.fus.collectors.FeatureUsagesCollector";
 
   @Override
-  protected JvmElementVisitor<Boolean> buildVisitor(@NotNull Project project, @NotNull HighlightSink sink, boolean isOnTheFly) {
-    return new DefaultJvmElementVisitor<>() {
-      @Override
-      public Boolean visitClass(@NotNull JvmClass clazz) {
-        PsiElement sourceElement = clazz.getSourceElement();
-        if (!(sourceElement instanceof PsiClass)) {
-          return null;
-        }
-        checkClass(project, (PsiClass)sourceElement, sink);
-        return false;
-      }
-    };
-  }
-
-  private static void checkClass(@NotNull Project project, @NotNull PsiClass checkedClass, @NotNull HighlightSink sink) {
+  protected void checkClass(@NotNull Project project, @NotNull PsiClass checkedClass, @NotNull HighlightSink sink) {
     if (!ExtensionUtil.isExtensionPointImplementationCandidate(checkedClass)) {
       return;
     }