From c970aac1a2236f35441b3ebeee68d7722287272b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Yann=20C=C3=A9bron?= <yann.cebron@jetbrains.com>
Date: Thu, 6 Apr 2017 18:14:38 +0200
Subject: [PATCH] [devkit] UndesirableClassUsageInspection -> UAST

---
 plugins/devkit/resources/META-INF/plugin.xml  |  2 +-
 .../UndesirableClassUsageInspection.java      | 49 +++++++++++--------
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/plugins/devkit/resources/META-INF/plugin.xml b/plugins/devkit/resources/META-INF/plugin.xml
index 0800326f69b3..a2f1e9a96acb 100644
--- a/plugins/devkit/resources/META-INF/plugin.xml
+++ b/plugins/devkit/resources/META-INF/plugin.xml
@@ -86,7 +86,7 @@
                      level="WARNING"
                      implementationClass="org.jetbrains.idea.devkit.inspections.InspectionMappingConsistencyInspection"/>
 
-    <localInspection language="JAVA" shortName="UndesirableClassUsage" displayName="Undesirable class usage"
+    <localInspection language="UAST" shortName="UndesirableClassUsage" displayName="Undesirable class usage"
                      groupKey="inspections.group.name"
                      enabledByDefault="true" level="WARNING"
                      implementationClass="org.jetbrains.idea.devkit.inspections.internal.UndesirableClassUsageInspection"/>
diff --git a/plugins/devkit/src/inspections/internal/UndesirableClassUsageInspection.java b/plugins/devkit/src/inspections/internal/UndesirableClassUsageInspection.java
index b2f339d40480..42c0d9b9fbb3 100644
--- a/plugins/devkit/src/inspections/internal/UndesirableClassUsageInspection.java
+++ b/plugins/devkit/src/inspections/internal/UndesirableClassUsageInspection.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2014 JetBrains s.r.o.
+ * Copyright 2000-2017 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,22 +19,29 @@ import com.intellij.codeInspection.ProblemHighlightType;
 import com.intellij.codeInspection.ProblemsHolder;
 import com.intellij.openapi.application.QueryExecutorBase;
 import com.intellij.openapi.ui.ComboBox;
-import com.intellij.psi.*;
+import com.intellij.psi.PsiClass;
+import com.intellij.psi.PsiElementVisitor;
+import com.intellij.psi.util.PsiTypesUtil;
+import com.intellij.uast.UastVisitorAdapter;
 import com.intellij.ui.components.JBList;
 import com.intellij.ui.components.JBScrollPane;
 import com.intellij.ui.components.JBTabbedPane;
 import com.intellij.ui.table.JBTable;
 import com.intellij.ui.treeStructure.Tree;
+import com.intellij.util.ObjectUtils;
 import com.intellij.util.QueryExecutor;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.idea.devkit.inspections.DevKitInspectionBase;
+import org.jetbrains.idea.devkit.inspections.DevKitUastInspectionBase;
+import org.jetbrains.uast.UCallExpression;
+import org.jetbrains.uast.UastCallKind;
+import org.jetbrains.uast.visitor.AbstractUastVisitor;
 
 import javax.swing.*;
 import java.awt.image.BufferedImage;
 import java.util.Map;
 
-public class UndesirableClassUsageInspection extends DevKitInspectionBase {
+public class UndesirableClassUsageInspection extends DevKitUastInspectionBase {
 
   private static final Map<String, String> CLASSES = ContainerUtil.<String, String>immutableMapBuilder()
     .put(JList.class.getName(), JBList.class.getName())
@@ -50,23 +57,23 @@ public class UndesirableClassUsageInspection extends DevKitInspectionBase {
   @Override
   @NotNull
   public PsiElementVisitor buildInternalVisitor(@NotNull final ProblemsHolder holder, boolean isOnTheFly) {
-    return new JavaElementVisitor() {
+    return new UastVisitorAdapter(new AbstractUastVisitor() {
       @Override
-      public void visitNewExpression(PsiNewExpression expression) {
-        PsiJavaCodeReferenceElement ref = expression.getClassReference();
-        if (ref == null) return;
-
-        PsiElement res = ref.resolve();
-        if (res == null) return;
-
-        String name = ((PsiClass)res).getQualifiedName();
-        if (name == null) return;
-
-        String replacement = CLASSES.get(name);
-        if (replacement == null) return;
-
-        holder.registerProblem(expression, "Please use '" + replacement + "' instead", ProblemHighlightType.LIKE_DEPRECATED);
+      public boolean visitCallExpression(@NotNull UCallExpression expression) {
+        if (expression.getKind() == UastCallKind.CONSTRUCTOR_CALL) {
+          final PsiClass psiClass = PsiTypesUtil.getPsiClass(expression.getReturnType());
+          if (psiClass != null) {
+            final String name = psiClass.getQualifiedName();
+            String replacement = CLASSES.get(name);
+            if (replacement != null) {
+              holder.registerProblem(ObjectUtils.assertNotNull(expression.getPsi()),
+                                     "Please use '" + replacement + "' instead",
+                                     ProblemHighlightType.LIKE_DEPRECATED);
+            }
+          }
+        }
+        return true;
       }
-    };
+    });
   }
-}
+}
\ No newline at end of file