Dispose IntelliJ platform components after annotation processing complete.

Annotation processors may cache ProcessingEnvironment.

(cherry picked from commit bd7a9c6)

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/AnnotationProcessingExtension.kt

@@ -130,6 +130,7 @@ abstract class AbstractAnnotationProcessingExtension(
                 project, psiManager, javaPsiFacade, projectScope, bindingTrace.bindingContext, appendJavaSourceRootsHandler)
 
         val processingResult = processingEnvironment.doAnnotationProcessing(files)
+        processingEnvironment.dispose()
 
         annotationProcessingComplete = true
         log {
@@ -161,19 +162,19 @@ abstract class AbstractAnnotationProcessingExtension(
     }
     
     private fun KotlinProcessingEnvironment.createTypeMapper(): KotlinTypeMapper {
-        return KotlinTypeMapper(bindingContext, ClassBuilderMode.full(false), NoResolveFileClassesProvider,
+        return KotlinTypeMapper(bindingContext(), ClassBuilderMode.full(false), NoResolveFileClassesProvider,
                                 IncompatibleClassTracker.DoNothing, JvmAbi.DEFAULT_MODULE_NAME, false)
     }
 
     private fun KotlinProcessingEnvironment.doAnnotationProcessing(files: Collection<KtFile>): ProcessingResult {
         run initializeProcessors@ {
-            processors.forEach { it.init(this) }
+            processors().forEach { it.init(this) }
-            log { "Initialized processors: " + processors.joinToString { it.javaClass.name } }
+            log { "Initialized processors: " + processors().joinToString { it.javaClass.name } }
         }
 
         val firstRoundAnnotations = RoundAnnotations(
                 sourceRetentionAnnotationHandler,
-                bindingContext,
+                bindingContext(),
                 createTypeMapper())
         
         run analyzeFilesForFirstRound@ {
@@ -185,7 +186,7 @@ abstract class AbstractAnnotationProcessingExtension(
                 javaSourceRoot.walk().filter { it.isFile && it.extension == "java" }.forEach {
                     val vFile = StandardFileSystems.local().findFileByPath(it.absolutePath)
                     if (vFile != null) {
-                        val javaFile = psiManager.findFile(vFile) as? PsiJavaFile
+                        val javaFile = psiManager().findFile(vFile) as? PsiJavaFile
                         if (javaFile != null) {
                             firstRoundAnnotations.analyzeFile(javaFile)
                         }
@@ -212,7 +213,7 @@ abstract class AbstractAnnotationProcessingExtension(
                     for (line in incrementalData.lines()) {
                         if (line.length < 3 || !line.startsWith("i ")) continue
                         val fqName = line.drop(2) 
-                        val psiClass = javaPsiFacade.findClass(fqName, projectScope) ?: continue
+                        val psiClass = javaPsiFacade().findClass(fqName, projectScope()) ?: continue
                         if (firstRoundAnnotations.analyzeDeclaration(psiClass)) {
                             analyzedClasses += fqName
                         }
@@ -237,14 +238,15 @@ abstract class AbstractAnnotationProcessingExtension(
         
         val finalRoundNumber = run annotationProcessing@ {
             val firstRoundEnvironment = KotlinRoundEnvironment(firstRoundAnnotations, false, 1)
-            process(firstRoundEnvironment)
+            process(firstRoundEnvironment) // Dispose for firstRoundEnvironment is called inside process
         } + 1
         
         log { "Starting round $finalRoundNumber (final)" }
         val finalRoundEnvironment = KotlinRoundEnvironment(firstRoundAnnotations.copy(), true, finalRoundNumber)
-        for (processor in processors) {
+        for (processor in processors()) {
             processor.process(emptySet(), finalRoundEnvironment)
         }
+        finalRoundEnvironment.dispose()
 
         return ProcessingResult(messager.errorCount, messager.warningCount, filer.wasAnythingGenerated)
     }
@@ -258,16 +260,16 @@ abstract class AbstractAnnotationProcessingExtension(
         
         // Add new Java source roots after the first round
         if (roundEnvironment.roundNumber == 1) {
-            appendJavaSourceRootsHandler(listOf(generatedSourcesOutputDir))
+            appendJavaSourceRootsHandler()(listOf(generatedSourcesOutputDir))
         }
         
         // Update the platform caches
-        (PsiManager.getInstance(project).modificationTracker as? PsiModificationTrackerImpl)?.incCounter()
+        (psiManager().modificationTracker as? PsiModificationTrackerImpl)?.incCounter()
         
         // Find generated files
         val localFileSystem = StandardFileSystems.local()
         val psiFiles = newJavaFiles
-                .map { localFileSystem.findFileByPath(it.absolutePath)?.let { psiManager.findFile(it) } }
+                .map { localFileSystem.findFileByPath(it.absolutePath)?.let { psiManager().findFile(it) } }
                 .filterIsInstance<PsiJavaFile>()
 
         if (psiFiles.isEmpty()) {
@@ -276,8 +278,9 @@ abstract class AbstractAnnotationProcessingExtension(
         }
         
         // Start the next round
-        val nextRoundAnnotations = roundEnvironment.roundAnnotations.copy().apply { analyzeFiles(psiFiles) }
+        val nextRoundAnnotations = roundEnvironment.roundAnnotations().copy().apply { analyzeFiles(psiFiles) }
         val nextRoundEnvironment = KotlinRoundEnvironment(nextRoundAnnotations, false, roundEnvironment.roundNumber + 1)
+        roundEnvironment.dispose()
         return process(nextRoundEnvironment)
     }
     
@@ -287,13 +290,13 @@ abstract class AbstractAnnotationProcessingExtension(
         val newFiles = mutableListOf<File>()
         filer.onFileCreatedHandler = { newFiles += it }
 
-        for (processor in processors) {
+        for (processor in processors()) {
             val supportedAnnotationNames = processor.supportedAnnotationTypes
             val acceptsAnyAnnotation = supportedAnnotationNames.contains("*")
 
             val applicableAnnotationNames = when (acceptsAnyAnnotation) {
-                true -> roundEnvironment.roundAnnotations.annotationsMap.keys
+                true -> roundEnvironment.roundAnnotations().annotationsMap.keys
-                false -> processor.supportedAnnotationTypes.filter { it in roundEnvironment.roundAnnotations.annotationsMap }
+                false -> processor.supportedAnnotationTypes.filter { it in roundEnvironment.roundAnnotations().annotationsMap }
             }
 
             if (applicableAnnotationNames.isEmpty()) {
@@ -302,7 +305,7 @@ abstract class AbstractAnnotationProcessingExtension(
             }
 
             val applicableAnnotations = applicableAnnotationNames
-                    .map { javaPsiFacade.findClass(it, projectScope)?.let { JeTypeElement(it) } }
+                    .map { javaPsiFacade().findClass(it, projectScope())?.let { JeTypeElement(it) } }
                     .filterNotNullTo(hashSetOf())
 
             log {

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/impl/DisposableRef.kt

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.annotation.processing.impl
+
+import com.intellij.openapi.Disposable
+
+class DisposableRef<out T : Any>(initialValue: T): Disposable {
+    @Volatile
+    private var value: T? = initialValue
+    
+    operator fun invoke() = value ?: throw IllegalStateException("Reference is disposed")
+    
+    override fun dispose() {
+        value = null
+    }
+}
+
+fun <T : Any> T.toDisposable() = DisposableRef(this)
+
+fun dispose(vararg refs: DisposableRef<*>) {
+    refs.forEach { it.dispose() }
+}

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/impl/KotlinElements.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.annotation.processing.impl
 
+import com.intellij.openapi.Disposable
 import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.*
@@ -31,7 +32,15 @@ import java.io.Writer
 import javax.lang.model.element.*
 import javax.lang.model.util.Elements
 
-class KotlinElements(val javaPsiFacade: JavaPsiFacade, val scope: GlobalSearchScope) : Elements {
+class KotlinElements(
+        javaPsiFacade: JavaPsiFacade, 
+        scope: GlobalSearchScope
+) : Elements, Disposable {
+    internal val javaPsiFacade = javaPsiFacade.toDisposable()
+    internal val scope = scope.toDisposable()
+
+    override fun dispose() = dispose(javaPsiFacade, scope)
+
     override fun hides(hider: Element, hidden: Element): Boolean {
         val hiderMethod = (hider as? JeMethodExecutableElement)?.psi ?: return false
         val hiddenMethod = (hidden as? JeMethodExecutableElement)?.psi ?: return false
@@ -87,12 +96,12 @@ class KotlinElements(val javaPsiFacade: JavaPsiFacade, val scope: GlobalSearchSc
     }
 
     override fun getPackageElement(name: CharSequence): PackageElement? {
-        val psiPackage = javaPsiFacade.findPackage(name.toString()) ?: return null
+        val psiPackage = javaPsiFacade().findPackage(name.toString()) ?: return null
         return JePackageElement(psiPackage)
     }
 
     override fun getTypeElement(name: CharSequence): TypeElement? {
-        val psiClass = javaPsiFacade.findClass(name.toString(), scope) ?: return null
+        val psiClass = javaPsiFacade().findClass(name.toString(), scope()) ?: return null
         return JeTypeElement(psiClass)
     }
 

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/impl/KotlinProcessingEnvironment.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.annotation.processing.impl
 
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiManager
@@ -26,32 +27,51 @@ import java.util.*
 import javax.annotation.processing.ProcessingEnvironment
 import javax.annotation.processing.Processor
 import javax.lang.model.SourceVersion
-import javax.lang.model.util.Elements
-import javax.lang.model.util.Types
 
 class KotlinProcessingEnvironment(
-        private val elements: Elements,
+        elements: KotlinElements,
-        private val types: Types,
+        types: KotlinTypes,
-        private val messager: KotlinMessager,
+        messager: KotlinMessager,
         options: Map<String, String>,
-        private val filer: KotlinFiler,
+        filer: KotlinFiler,
         
-        internal val processors: List<Processor>,
+        processors: List<Processor>,
         
-        internal val project: Project,
+        project: Project,
-        internal val psiManager: PsiManager,
+        psiManager: PsiManager,
-        internal val javaPsiFacade: JavaPsiFacade,
+        javaPsiFacade: JavaPsiFacade,
-        internal val projectScope: GlobalSearchScope,
+        projectScope: GlobalSearchScope,
-        internal val bindingContext: BindingContext,
+        bindingContext: BindingContext,
-        internal val appendJavaSourceRootsHandler: (List<File>) -> Unit
+        appendJavaSourceRootsHandler: (List<File>) -> Unit
-) : ProcessingEnvironment {
+) : ProcessingEnvironment, Disposable {
-    private val options = Collections.unmodifiableMap(options) 
+    private val elements = elements.toDisposable()
+    private val types = types.toDisposable()
+    private val messager = messager.toDisposable()
+    private val filer = filer.toDisposable()
+    internal val processors = processors.toDisposable()
     
-    override fun getElementUtils() = elements
+    internal val project = project.toDisposable()
-    override fun getTypeUtils() = types
+    internal val psiManager = psiManager.toDisposable()
-    override fun getMessager() = messager
+    internal val javaPsiFacade = javaPsiFacade.toDisposable()
+    internal val projectScope = projectScope.toDisposable()
+    internal val bindingContext = bindingContext.toDisposable()
+    
+    internal val appendJavaSourceRootsHandler = appendJavaSourceRootsHandler.toDisposable()
+    private val options = Collections.unmodifiableMap(options).toDisposable()
+
+    override fun dispose() {
+        types.dispose()
+        elements.dispose()
+        dispose(elements, types, messager, filer, processors,
+                project, psiManager, javaPsiFacade, projectScope, bindingContext,
+                appendJavaSourceRootsHandler, options)
+    }
+
+    override fun getElementUtils() = elements()
+    override fun getTypeUtils() = types()
+    override fun getMessager() = messager()
     override fun getLocale() = Locale.getDefault()
     override fun getSourceVersion() = SourceVersion.RELEASE_8
-    override fun getOptions() = options
+    override fun getOptions() = options()
-    override fun getFiler() = filer
+    override fun getFiler() = filer()
 }

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/impl/KotlinRoundEnvironment.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.annotation.processing.impl
 
+import com.intellij.openapi.Disposable
 import org.jetbrains.kotlin.annotation.processing.RoundAnnotations
 import org.jetbrains.kotlin.java.model.toJeElement
 import javax.annotation.processing.RoundEnvironment
@@ -23,10 +24,14 @@ import javax.lang.model.element.Element
 import javax.lang.model.element.TypeElement
 
 internal class KotlinRoundEnvironment(
-        val roundAnnotations: RoundAnnotations,
+        roundAnnotations: RoundAnnotations,
         private val isProcessingOver: Boolean,
         internal val roundNumber: Int
-) : RoundEnvironment {
+) : RoundEnvironment, Disposable {
+    val roundAnnotations = roundAnnotations.toDisposable()
+    
+    override fun dispose() = dispose(roundAnnotations)
+    
     private var isError = false
     
     override fun getRootElements() = emptySet<Element>()
@@ -34,7 +39,7 @@ internal class KotlinRoundEnvironment(
     override fun processingOver() = isProcessingOver
 
     private fun getElementsAnnotatedWith(fqName: String): Set<Element> {
-        val declarations = roundAnnotations.annotationsMap[fqName] ?: return emptySet()
+        val declarations = roundAnnotations().annotationsMap[fqName] ?: return emptySet()
         return hashSetOf<Element>().apply {
             for (declaration in declarations) {
                 declaration.toJeElement()?.let { add(it) }

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/impl/KotlinTypes.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.annotation.processing.impl
 
+import com.intellij.openapi.Disposable
 import com.intellij.psi.*
 import com.intellij.psi.impl.source.PsiImmediateClassType
 import com.intellij.psi.search.GlobalSearchScope
@@ -31,7 +32,17 @@ import javax.lang.model.element.TypeElement
 import javax.lang.model.type.*
 import javax.lang.model.util.Types
 
-class KotlinTypes(val javaPsiFacade: JavaPsiFacade, val psiManager: PsiManager, val scope: GlobalSearchScope) : Types {
+class KotlinTypes(
+        javaPsiFacade: JavaPsiFacade,
+        psiManager: PsiManager,
+        scope: GlobalSearchScope
+) : Types, Disposable {
+    val javaPsiFacade = javaPsiFacade.toDisposable()
+    val psiManager = psiManager.toDisposable()
+    val scope = scope.toDisposable()
+
+    override fun dispose() = dispose(javaPsiFacade, psiManager, scope)
+
     override fun contains(containing: TypeMirror, contained: TypeMirror): Boolean {
         assertKindNot(containing, TypeKind.PACKAGE, TypeKind.EXECUTABLE)
         assertKindNot(contained, TypeKind.PACKAGE, TypeKind.EXECUTABLE)
@@ -62,7 +73,7 @@ class KotlinTypes(val javaPsiFacade: JavaPsiFacade, val psiManager: PsiManager,
         if (componentType is ExecutableType || componentType is NoType) error(componentType)
         assertJeType(componentType); componentType as JePsiType
         
-        return JeArrayType(PsiArrayType(componentType.psiType), psiManager, isRaw = false)
+        return JeArrayType(PsiArrayType(componentType.psiType), psiManager(), isRaw = false)
     }
 
     override fun isAssignable(t1: TypeMirror, t2: TypeMirror): Boolean {
@@ -85,11 +96,11 @@ class KotlinTypes(val javaPsiFacade: JavaPsiFacade, val psiManager: PsiManager,
         if (superBound != null && superBound !is JePsiType) illegalArg("superBound should have PsiType")
         
         return JeWildcardType(if (extendsBound != null) {
-            PsiWildcardType.createExtends(psiManager, (extendsBound as JePsiType).psiType)
+            PsiWildcardType.createExtends(psiManager(), (extendsBound as JePsiType).psiType)
         } else if (superBound != null) {
-            PsiWildcardType.createSuper(psiManager, (superBound as JePsiType).psiType)
+            PsiWildcardType.createSuper(psiManager(), (superBound as JePsiType).psiType)
         } else {
-            PsiWildcardType.createUnbounded(psiManager)
+            PsiWildcardType.createUnbounded(psiManager())
         }, isRaw = false)
     }
 
@@ -106,7 +117,7 @@ class KotlinTypes(val javaPsiFacade: JavaPsiFacade, val psiManager: PsiManager,
         if (t.kind == TypeKind.PACKAGE) throw IllegalArgumentException("Invalid type: $t")
         return when (t) {
             is JeTypeVariableType -> TypeConversionUtil.typeParameterErasure(t.parameter).toJeType(t.psiManager, isRaw = true)
-            is JePsiType -> TypeConversionUtil.erasure(t.psiType).toJeType(psiManager, isRaw = true)
+            is JePsiType -> TypeConversionUtil.erasure(t.psiType).toJeType(psiManager(), isRaw = true)
             is JeMethodExecutableTypeMirror -> {
                 val oldSignature = t.signature
                 val parameterTypes = oldSignature?.parameterTypes?.toList() ?: t.psi.parameterList.parameters.map { it.type }
@@ -128,17 +139,17 @@ class KotlinTypes(val javaPsiFacade: JavaPsiFacade, val psiManager: PsiManager,
         if (t is NoType || t is ExecutableType) throw IllegalArgumentException("Invalid type: $t")
 
         if (t is JeDeclaredType && t.psiType is PsiImmediateClassType) {
-            return t.psiClass.superTypes.map { it.toJeType(psiManager) }
+            return t.psiClass.superTypes.map { it.toJeType(psiManager()) }
         }
 
         val psiType = (t as? JePsiType)?.psiType as? PsiClassType ?: return emptyList()
-        return psiType.superTypes.map { it.toJeType(psiManager) }
+        return psiType.superTypes.map { it.toJeType(psiManager()) }
     }
 
     override fun boxedClass(p: PrimitiveType): TypeElement? {
         p as? JePrimitiveType ?: throw IllegalArgumentException("Unknown type: $p")
         val boxedTypeName = p.psiType.boxedTypeName
-        val boxedClass = javaPsiFacade.findClass(boxedTypeName, scope) 
+        val boxedClass = javaPsiFacade().findClass(boxedTypeName, scope()) 
                          ?: throw IllegalStateException("Can't find boxed class $boxedTypeName")
         return JeTypeElement(boxedClass)
     }
@@ -253,7 +264,7 @@ class KotlinTypes(val javaPsiFacade: JavaPsiFacade, val psiManager: PsiManager,
                 val returnType = substitutor.substitute(element.psi.returnType)
                 JeMethodExecutableTypeMirror(method, signature, returnType)
             }
-            is JeVariableElement -> substitutor.substitute(element.psi.type).toJeType(psiManager)
+            is JeVariableElement -> substitutor.substitute(element.psi.type).toJeType(psiManager())
             else -> throw IllegalArgumentException("Invalid element type: $element")
         }
     }

plugins/plugins-tests/tests/org/jetbrains/kotlin/annotation/processing/test/processor/ProcessorTests.kt

@@ -17,12 +17,15 @@
 package org.jetbrains.kotlin.annotation.processing.test.processor
 
 import org.intellij.lang.annotations.Language
+import org.jetbrains.kotlin.annotation.processing.impl.DisposableRef
+import org.jetbrains.kotlin.annotation.processing.impl.KotlinProcessingEnvironment
 import org.jetbrains.kotlin.incremental.SourceRetentionAnnotationHandlerImpl
 import org.jetbrains.kotlin.java.model.elements.*
 import org.jetbrains.kotlin.java.model.types.JeDeclaredType
 import org.jetbrains.kotlin.java.model.types.JeMethodExecutableTypeMirror
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
+import javax.annotation.processing.ProcessingEnvironment
 import javax.lang.model.element.AnnotationMirror
 import javax.lang.model.element.Element
 import javax.lang.model.type.DeclaredType
@@ -273,4 +276,40 @@ class ProcessorTests : AbstractProcessorTest() {
         assertEquals("T", baseF.asType().toString())
         check(baseF, "java.lang.String")
     }
+    
+    fun testDispose() {
+        var savedEnv: ProcessingEnvironment? = null
+        
+        test("AsMemberOf", "*") { set, roundEnv, env ->
+            savedEnv = env
+        }
+        
+        fun <T> T.test(f: T.() -> Unit) {
+            var exceptionCaught = false
+            try {
+                f()
+            } 
+            catch (e: IllegalStateException) { 
+                exceptionCaught = true
+            }
+            
+            assertEquals("Exception was not caught", true, exceptionCaught)
+        }
+
+        with (savedEnv as KotlinProcessingEnvironment) {
+            test { elementUtils }
+            test { typeUtils }
+            test { messager }
+            test { options }
+            test { filer }
+
+            fun testDisposable(name: String) {
+                test { (javaClass.methods.first { it.name.startsWith("$name$") }.invoke(this) as DisposableRef<*>).invoke() }
+            }
+
+            testDisposable("getProject")
+            testDisposable("getJavaPsiFacade")
+            testDisposable("getBindingContext")
+        }
+    }
 }