From 5cf5024d0837e1c1abf523a71626fbfb15b13cdd Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Tue, 22 Mar 2016 19:45:55 +0300
Subject: [PATCH] Lint: fix Lint additional checkers with the custom handlers

---
 .../tools/klint/client/api/LintDriver.java    |  9 ++--
 .../org/jetbrains/uast/check/UastChecker.kt   | 51 ++++++++++++-------
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/plugins/lint/lint-api/src/com/android/tools/klint/client/api/LintDriver.java b/plugins/lint/lint-api/src/com/android/tools/klint/client/api/LintDriver.java
index 83edd30eae9..6ee62e11931 100644
--- a/plugins/lint/lint-api/src/com/android/tools/klint/client/api/LintDriver.java
+++ b/plugins/lint/lint-api/src/com/android/tools/klint/client/api/LintDriver.java
@@ -1445,7 +1445,6 @@ public class LintDriver {
             if (ideaProject == null) {
                 return;
             }
-            List<UastLanguagePlugin> plugins = mClient.getLanguagePlugins();
 
             for (JavaContext context : contexts) {
                 fireEvent(LintListener.EventType.SCANNING_FILE, context);
@@ -1456,10 +1455,10 @@ public class LintDriver {
                         UastVisitor customVisitor = scanner.createUastVisitor(context);
                         if (customVisitor != null) {
                             UastChecker.INSTANCE.checkWithCustomHandler(
-                                    ideaProject, context.file, plugins, customVisitor);
+                                    ideaProject, context.file, context, customVisitor);
                         } else {
                             UastChecker.INSTANCE.check(
-                                    ideaProject, context.file, (UastScanner)check, plugins, context);
+                                    ideaProject, context.file, (UastScanner)check, context);
                         }
                     }
                 }
@@ -1516,9 +1515,9 @@ public class LintDriver {
             for (UastScanner detector : detectors) {
                 UastVisitor customHandler = detector.createUastVisitor(context);
                 if (customHandler != null) {
-                    checker.checkWithCustomHandler(intellijProject, file, plugins, customHandler);
+                    checker.checkWithCustomHandler(intellijProject, file, context, customHandler);
                 } else {
-                    checker.check(intellijProject, file, detector, plugins, context);
+                    checker.check(intellijProject, file, detector, context);
                 }
             }
         }
diff --git a/plugins/lint/uast-android/src/org/jetbrains/uast/check/UastChecker.kt b/plugins/lint/uast-android/src/org/jetbrains/uast/check/UastChecker.kt
index 5fe613fc862..af50994863f 100644
--- a/plugins/lint/uast-android/src/org/jetbrains/uast/check/UastChecker.kt
+++ b/plugins/lint/uast-android/src/org/jetbrains/uast/check/UastChecker.kt
@@ -41,13 +41,24 @@ object UastChecker {
     fun checkWithCustomHandler(
             project: Project,
             file: File,
-            converters: List<UastLanguagePlugin>,
+            context: UastAndroidContext,
             visitor: UastVisitor) {
-        check(project, file, converters, UastHandler { visitor.handle(it) })
+        check(project, file, context, UastHandler { visitor.handle(it) })
     }
 
-    fun check(project: Project, file: File, plugins: List<UastLanguagePlugin>, handler: UastHandler) {
+    fun check(project: Project, file: File, context: UastAndroidContext, handler: UastHandler) {
         val vfile = VirtualFileManager.getInstance().findFileByUrl("file://" + file.absolutePath) ?: return
+
+        val plugins = context.languagePlugins
+        val additionalCheckers = plugins.fold(mutableListOf<UastAdditionalChecker>()) { list, plugin ->
+            for (checker in plugin.additionalCheckers) {
+                if (checker is AndroidUastAdditionalChecker) list += checker
+            }
+            list
+        }
+
+        val handlerWrapper = HandlerWrapper(handler, additionalCheckers, context)
+
         ApplicationManager.getApplication().runReadAction {
             val psiFile = PsiManager.getInstance(project).findFile(vfile)
 
@@ -55,12 +66,12 @@ object UastChecker {
                 when (psiFile) {
                     is PsiJavaFile -> {
                         val ufile = JavaUastLanguagePlugin.converter.convertWithParent(psiFile)
-                        ufile?.handleTraverse(handler)
+                        ufile?.handleTraverse(handlerWrapper)
                     }
                     else -> for (plugin in plugins) {
                         val ufile = plugin.converter.convertWithParent(psiFile)
                         if (ufile != null) {
-                            ufile.handleTraverse(handler)
+                            ufile.handleTraverse(handlerWrapper)
                             break
                         }
                     }
@@ -69,21 +80,27 @@ object UastChecker {
         }
     }
 
-    fun check(project: Project, file: File, scanner: UastScanner, plugins: List<UastLanguagePlugin>, context: UastAndroidContext) {
+    private class HandlerWrapper(
+            val original: UastHandler,
+            val additionalCheckers: List<UastAdditionalChecker>,
+            val context: UastAndroidContext
+    ) : UastHandler {
+        override fun invoke(element: UElement) {
+            original(element)
+            for (checker in additionalCheckers) {
+                checker(element, this, context)
+            }
+        }
+    }
+
+    fun check(project: Project, file: File, scanner: UastScanner, context: UastAndroidContext) {
         val applicableFunctionNames = scanner.applicableFunctionNames ?: emptyList()
         val applicableSuperClasses = scanner.applicableSuperClasses ?: emptyList()
         val applicableConstructorTypes = scanner.applicableConstructorTypes ?: emptyList()
 
         val appliesToResourcesRefs = scanner.appliesToResourceRefs()
 
-        val additionalCheckers = plugins.fold(mutableListOf<UastAdditionalChecker>()) { list, plugin ->
-            for (checker in plugin.additionalCheckers) {
-                if (checker is AndroidUastAdditionalChecker) list += checker
-            }
-            list
-        }
-
-        var handler: UastHandler? = null
+        var handler: UastHandler?
         handler = UastHandler { element ->
             when (element) {
                 is UCallExpression -> {
@@ -131,13 +148,9 @@ object UastChecker {
                     }
                 }
             }
-
-            for (checker in additionalCheckers) {
-                checker(element, handler!!, context)
-            }
         }
 
-        check(project, file, plugins, handler)
+        check(project, file, context, handler)
     }
 
 }
\ No newline at end of file