Kapt: Dispose every JeElement and JeTypeMirror because some annotation processors may store them in static fields. (KT-13917)

Hello, Data Binding, why you cache them? :(

(cherry picked from commit d86a30b)

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

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.CompilerConfigurationKey
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
+import org.jetbrains.kotlin.java.model.internal.JeElementRegistry
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import java.io.File
 
@@ -128,6 +129,8 @@ class AnnotationProcessingComponentRegistrar : ComponentRegistrar {
                 classpath, generatedOutputDirFile, classesOutputDir, javaRoots, verboseOutput, 
                 incrementalDataFile, sourceRetentionAnnotationHandler)
 
+        project.registerService(JeElementRegistry::class.java, JeElementRegistry())
+        
         AnalysisCompletedHandlerExtension.registerExtension(project, annotationProcessingExtension)
     }
 }

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

@@ -31,6 +31,6 @@ class DisposableRef<out T : Any>(initialValue: T): Disposable {
 
 fun <T : Any> T.toDisposable() = DisposableRef(this)
 
-fun dispose(vararg refs: DisposableRef<*>) {
+fun dispose(vararg refs: DisposableRef<*>?) {
-    refs.forEach { it.dispose() }
+    refs.forEach { it?.dispose() }
 }

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

@@ -17,10 +17,12 @@
 package org.jetbrains.kotlin.annotation.processing.impl
 
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiManager
 import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.java.model.internal.JeElementRegistry
 import org.jetbrains.kotlin.resolve.BindingContext
 import java.io.File
 import java.util.*
@@ -56,12 +58,14 @@ class KotlinProcessingEnvironment(
     internal val projectScope = projectScope.toDisposable()
     internal val bindingContext = bindingContext.toDisposable()
 
+    private val registry = ServiceManager.getService(project, JeElementRegistry::class.java).toDisposable()
     internal val appendJavaSourceRootsHandler = appendJavaSourceRootsHandler.toDisposable()
     private val options = Collections.unmodifiableMap(options).toDisposable()
 
     override fun dispose() {
-        types.dispose()
+        types().dispose()
-        elements.dispose()
+        elements().dispose()
+        registry().dispose()
         dispose(elements, types, messager, filer, processors,
                 project, psiManager, javaPsiFacade, projectScope, bindingContext,
                 appendJavaSourceRootsHandler, options)

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

@@ -272,8 +272,8 @@ class KotlinTypes(
     }
 
     override fun isSubsignature(t1: ExecutableType, t2: ExecutableType): Boolean {
-        val m1 = (t1 as JeClassInitializerExecutableTypeMirror).initializer
+        val m1 = (t1 as JeClassInitializerExecutableTypeMirror).psi
-        val m2 = (t2 as JeClassInitializerExecutableTypeMirror).initializer
+        val m2 = (t2 as JeClassInitializerExecutableTypeMirror).psi
         
         // No parameters
         if (m1 is PsiClassInitializer && m2 is PsiClassInitializer) return true

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/JeElement.kt

@@ -16,9 +16,37 @@
 
 package org.jetbrains.kotlin.java.model
 
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.ServiceManager
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiManager
+import org.jetbrains.kotlin.annotation.processing.impl.toDisposable
+import org.jetbrains.kotlin.java.model.internal.JeElementRegistry
 import javax.lang.model.element.Element
 
-interface JeElement : Element {
+interface JePsiElementOwner {
     val psi: PsiElement
 }
+
+interface JeElement : Element, JePsiElementOwner
+
+abstract class JeDisposablePsiElementOwner<out T : PsiElement>(element: T) : JePsiElementOwner, Disposable {
+    private val disposablePsi = element.toDisposable()
+
+    init {
+        @Suppress("LeakingThis")
+        ServiceManager.getService(element.project, JeElementRegistry::class.java).register(this)
+    }
+
+    override fun dispose() {
+        disposablePsi.dispose()
+    }
+
+    override val psi: T
+        get() = disposablePsi()
+
+    val psiManager: PsiManager
+        get() = disposablePsi().manager
+}
+
+abstract class JeAbstractElement<out T : PsiElement>(element: T) : JeDisposablePsiElementOwner<T>(element), JeElement, Disposable

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JeAnnotationMirror.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.java.model.elements
 
 import com.intellij.psi.*
+import org.jetbrains.kotlin.java.model.JeDisposablePsiElementOwner
 import org.jetbrains.kotlin.java.model.internal.getTypeWithTypeParameters
 import org.jetbrains.kotlin.java.model.types.JeDeclaredErrorType
 import org.jetbrains.kotlin.java.model.types.JeDeclaredType
@@ -25,7 +26,7 @@ import javax.lang.model.element.AnnotationValue
 import javax.lang.model.element.ExecutableElement
 import javax.lang.model.type.DeclaredType
 
-class JeAnnotationMirror(val psi: PsiAnnotation) : AnnotationMirror {
+class JeAnnotationMirror(psi: PsiAnnotation) : JeDisposablePsiElementOwner<PsiAnnotation>(psi), AnnotationMirror {
     override fun getAnnotationType(): DeclaredType? {
         val psiClass = resolveAnnotationClass() ?: return JeDeclaredErrorType
         return JeDeclaredType(psiClass.getTypeWithTypeParameters(), psiClass)

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JeAnnotationValue.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.java.model.elements
 
 import com.intellij.psi.*
+import org.jetbrains.kotlin.java.model.JeDisposablePsiElementOwner
 import org.jetbrains.kotlin.java.model.internal.calcConstantValue
 import org.jetbrains.kotlin.java.model.types.toJeType
 import java.util.*
@@ -48,22 +49,22 @@ fun JeAnnotationValue(psi: PsiAnnotationMemberValue): AnnotationValue {
     return annotationValue
 }
 
-class JeAnnotationAnnotationValue(val psi: PsiAnnotation) : AnnotationValue {
+class JeAnnotationAnnotationValue(psi: PsiAnnotation) : JeDisposablePsiElementOwner<PsiAnnotation>(psi), AnnotationValue {
     override fun getValue() = JeAnnotationMirror(psi)
     override fun <R : Any?, P : Any?> accept(v: AnnotationValueVisitor<R, P>, p: P) = v.visitAnnotation(value, p)
 }
 
-class JeEnumValueAnnotationValue(val psi: PsiEnumConstant) : AnnotationValue {
+class JeEnumValueAnnotationValue(psi: PsiEnumConstant) : JeDisposablePsiElementOwner<PsiEnumConstant>(psi), AnnotationValue {
     override fun getValue() = JeVariableElement(psi)
     override fun <R : Any?, P : Any?> accept(v: AnnotationValueVisitor<R, P>, p: P) = v.visitEnumConstant(value, p) 
 }
 
-class JeTypeAnnotationValue(val psi: PsiClassObjectAccessExpression) : AnnotationValue {
+class JeTypeAnnotationValue(psi: PsiClassObjectAccessExpression) : JeDisposablePsiElementOwner<PsiClassObjectAccessExpression>(psi), AnnotationValue {
     override fun getValue() = psi.operand.type.toJeType(psi.manager)
     override fun <R : Any?, P : Any?> accept(v: AnnotationValueVisitor<R, P>, p: P) = v.visitType(value, p)
 }
 
-abstract class JePrimitiveAnnotationValue : AnnotationValue {
+abstract class JePrimitiveAnnotationValue<out T : PsiElement>(psi: T) : JeDisposablePsiElementOwner<T>(psi), AnnotationValue {
     override fun <R : Any?, P : Any?> accept(v: AnnotationValueVisitor<R, P>, p: P): R {
         val value = this.value
         return when (value) {
@@ -81,25 +82,25 @@ abstract class JePrimitiveAnnotationValue : AnnotationValue {
     }
 }
 
-class JeLiteralAnnotationValue(val psi: PsiLiteral) : JePrimitiveAnnotationValue() {
+class JeLiteralAnnotationValue(psi: PsiLiteral) : JePrimitiveAnnotationValue<PsiLiteral>(psi) {
     override fun getValue() = psi.value
 }
 
-class JeExpressionAnnotationValue(val psi: PsiExpression) : JePrimitiveAnnotationValue() {
+class JeExpressionAnnotationValue(psi: PsiExpression) : JePrimitiveAnnotationValue<PsiExpression>(psi) {
     override fun getValue() = psi.calcConstantValue()
 }
 
-class JeArrayAnnotationValue(val psi: PsiArrayInitializerMemberValue) : AnnotationValue {
+class JeArrayAnnotationValue(psi: PsiArrayInitializerMemberValue) : JeDisposablePsiElementOwner<PsiArrayInitializerMemberValue>(psi), AnnotationValue {
     override fun getValue() = psi.initializers.map(::JeAnnotationValue)
     override fun <R : Any?, P : Any?> accept(v: AnnotationValueVisitor<R, P>, p: P) = v.visitArray(value, p)
 }
 
-class JeSingletonArrayAnnotationValue(val psi: PsiAnnotationMemberValue) : AnnotationValue {
+class JeSingletonArrayAnnotationValue(psi: PsiAnnotationMemberValue) : JeDisposablePsiElementOwner<PsiAnnotationMemberValue>(psi), AnnotationValue {
     override fun getValue(): List<AnnotationValue> = Collections.singletonList(JeAnnotationValue(psi))
     override fun <R : Any?, P : Any?> accept(v: AnnotationValueVisitor<R, P>, p: P) = v.visitArray(value, p)
 }
 
-class JeErrorAnnotationValue(val psi: PsiElement) : AnnotationValue {
+class JeErrorAnnotationValue(psi: PsiElement) : JeDisposablePsiElementOwner<PsiElement>(psi), AnnotationValue {
     override fun <R : Any?, P : Any?> accept(v: AnnotationValueVisitor<R, P>, p: P) = v.visitString(psi.text, p)
     override fun getValue() = null
 }

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JeClassInitializerExecutableElement.kt

@@ -24,9 +24,9 @@ import org.jetbrains.kotlin.java.model.types.JeNoneType
 import javax.lang.model.element.*
 import javax.lang.model.type.TypeMirror
 
-class JeClassInitializerExecutableElement(override val psi: PsiClassInitializer) : JeElement,
+class JeClassInitializerExecutableElement(
-        ExecutableElement, JeNoAnnotations, JeModifierListOwner
+        psi: PsiClassInitializer
-{
+) : JeAbstractElement<PsiClassInitializer>(psi), ExecutableElement, JeNoAnnotations, JeModifierListOwner {
     val isStaticInitializer = psi.isStatic
 
     override fun getEnclosingElement() = psi.containingClass?.let { JeTypeElement(it) }

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JeMethodExecutableElement.kt

@@ -26,7 +26,7 @@ import org.jetbrains.kotlin.java.model.types.toJeType
 import javax.lang.model.element.*
 import javax.lang.model.type.TypeMirror
 
-class JeMethodExecutableElement(override val psi: PsiMethod) : JeElement, ExecutableElement, JeModifierListOwner, JeAnnotationOwner {
+class JeMethodExecutableElement(psi: PsiMethod) : JeAbstractElement<PsiMethod>(psi), ExecutableElement, JeModifierListOwner, JeAnnotationOwner {
     override fun getEnclosingElement() = psi.containingClass?.let(::JeTypeElement)
 
     override fun getSimpleName(): JeName {

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JePackageElement.kt

@@ -17,16 +17,13 @@
 package org.jetbrains.kotlin.java.model.elements
 
 import com.intellij.psi.PsiPackage
-import org.jetbrains.kotlin.java.model.JeElement
+import org.jetbrains.kotlin.java.model.*
-import org.jetbrains.kotlin.java.model.JeModifierListOwner
-import org.jetbrains.kotlin.java.model.JeName
-import org.jetbrains.kotlin.java.model.JeNoAnnotations
 import org.jetbrains.kotlin.java.model.types.JePackageTypeMirror
 import javax.lang.model.element.ElementKind
 import javax.lang.model.element.ElementVisitor
 import javax.lang.model.element.PackageElement
 
-class JePackageElement(override val psi: PsiPackage) : JeElement, PackageElement, JeModifierListOwner, JeNoAnnotations {
+class JePackageElement(psi: PsiPackage) : JeAbstractElement<PsiPackage>(psi), PackageElement, JeModifierListOwner, JeNoAnnotations {
     override fun getEnclosingElement() = null
 
     override fun getSimpleName() = JeName(psi.name)

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JeTypeElement.kt

@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.java.model.types.toJeType
 import javax.lang.model.element.*
 import javax.lang.model.type.TypeMirror
 
-class JeTypeElement(override val psi: PsiClass) : JeElement, TypeElement, JeAnnotationOwner, JeModifierListOwner {
+class JeTypeElement(psi: PsiClass) : JeAbstractElement<PsiClass>(psi), TypeElement, JeAnnotationOwner, JeModifierListOwner {
     override fun getEnclosingElement(): Element? {
         psi.containingClass?.let { return JeTypeElement(it) }
         val javaFile = psi.containingFile as? PsiJavaFile ?: return null

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JeTypeParameterElement.kt

@@ -26,9 +26,9 @@ import javax.lang.model.element.ElementVisitor
 import javax.lang.model.element.TypeParameterElement
 
 class JeTypeParameterElement(
-        override val psi: PsiTypeParameter,
+        psi: PsiTypeParameter,
         val parent: JeElement?
-) : JeElement, TypeParameterElement, JeAnnotationOwner, JeModifierListOwner {
+) : JeAbstractElement<PsiTypeParameter>(psi), TypeParameterElement, JeAnnotationOwner, JeModifierListOwner {
     override fun getSimpleName() = JeName(psi.name)
 
     override fun getEnclosingElement() = parent

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/elements/JeVariableElement.kt

@@ -26,10 +26,12 @@ import javax.lang.model.element.ElementKind
 import javax.lang.model.element.ElementVisitor
 import javax.lang.model.element.VariableElement
 
-class JeVariableElement(override val psi: PsiVariable) : JeElement, VariableElement, JeModifierListOwner, JeAnnotationOwner {
+class JeVariableElement(psi: PsiVariable) : JeAbstractElement<PsiVariable>(psi), VariableElement, JeModifierListOwner, JeAnnotationOwner {
     override fun getSimpleName() = JeName(psi.name)
 
     override fun getEnclosingElement(): JeElement? {
+        val psi = psi
+
         if (psi is PsiParameter) {
             (psi.declarationScope as? PsiMethod)?.let { return JeMethodExecutableElement(it) }
         }

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/internal/JeElementRegistry.kt

@@ -0,0 +1,42 @@
+/*
+ * 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.java.model.internal
+
+import com.intellij.openapi.Disposable
+
+/*
+    Some annotation processors (I'm looking at you, Data Binding)
+      have a habit to cache Element or TypeMirror instances into static fields.
+    Because almost every JeElement or JeTypeMirror is dependent on some PsiElement, and PsiElements are dependent on PsiProject,
+      the entire project will not be gc'ed normally.
+    JeElementRegistry allows us to remove all Psi*-references from our wrapper Element and TypeMirror instances.
+
+    So this class basically remembers all instances created by kapt, and disposes all of them as the Annotation Processing is complete.
+ */
+class JeElementRegistry : Disposable {
+    private val list = mutableListOf<Disposable>()
+
+    fun register(obj: Disposable) {
+        list += obj
+    }
+
+    override fun dispose() {
+        val list = this.list
+        list.forEach { it.dispose() }
+        this.list.clear()
+    }
+}

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JeArrayType.kt

@@ -23,10 +23,10 @@ import javax.lang.model.type.TypeKind
 import javax.lang.model.type.TypeVisitor
 
 class JeArrayType(
-        override val psiType: PsiArrayType, 
+        psiType: PsiArrayType,
-        override val psiManager: PsiManager,
+        psiManager: PsiManager,
         private val isRaw: Boolean
-) : JePsiType(), JeTypeWithManager, ArrayType {
+) : JePsiTypeBase<PsiArrayType>(psiType, psiManager), ArrayType {
     override fun getKind() = TypeKind.ARRAY
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitArray(this, p)
     override fun getComponentType() = psiType.componentType.toJeType(psiManager, isRaw = isRaw)

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JeClassInitializerExecutableTypeMirror.kt

@@ -17,18 +17,17 @@
 package org.jetbrains.kotlin.java.model.types
 
 import com.intellij.psi.PsiClassInitializer
-import com.intellij.psi.PsiManager
+import org.jetbrains.kotlin.java.model.JeDisposablePsiElementOwner
 import org.jetbrains.kotlin.java.model.internal.isStatic
 import javax.lang.model.type.*
 
-class JeClassInitializerExecutableTypeMirror(val initializer: PsiClassInitializer) : JeTypeMirror, JeTypeWithManager, ExecutableType {
+class JeClassInitializerExecutableTypeMirror(
+        psi: PsiClassInitializer
+) : JeDisposablePsiElementOwner<PsiClassInitializer>(psi), JeTypeMirror, JeTypeWithManager, ExecutableType {
     override fun getKind() = TypeKind.EXECUTABLE
     
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitExecutable(this, p)
 
-    override val psiManager: PsiManager
-        get() = initializer.manager
-
     override fun getReturnType() = JeVoidType
 
     override fun getReceiverType() = JeNoneType
@@ -39,15 +38,15 @@ class JeClassInitializerExecutableTypeMirror(val initializer: PsiClassInitialize
 
     override fun getTypeVariables() = emptyList<TypeVariable>()
 
-    override fun toString() = (initializer.containingClass?.qualifiedName?.let { it + "." } ?: "") +
+    override fun toString() = (psi.containingClass?.qualifiedName?.let { it + "." } ?: "") +
-                              (if (initializer.isStatic) "<clinit>" else "<instinit>")
+                              (if (psi.isStatic) "<clinit>" else "<instinit>")
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (other?.javaClass != javaClass) return false
         other as? JeClassInitializerExecutableTypeMirror ?: return false
-        return initializer == other.initializer
+        return psi == other.psi
     }
 
-    override fun hashCode() = initializer.hashCode()
+    override fun hashCode() = psi.hashCode()
 }

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JeDeclaredType.kt

@@ -21,6 +21,7 @@ import com.intellij.psi.*
 import com.intellij.psi.impl.PsiSubstitutorImpl
 import com.intellij.psi.impl.source.PsiClassReferenceType
 import com.intellij.psi.impl.source.PsiImmediateClassType
+import org.jetbrains.kotlin.annotation.processing.impl.toDisposable
 import org.jetbrains.kotlin.java.model.elements.JeTypeElement
 import org.jetbrains.kotlin.java.model.internal.getTypeWithTypeParameters
 import org.jetbrains.kotlin.java.model.internal.isStatic
@@ -39,11 +40,22 @@ fun createImmediateClassType(psiClass: PsiClass, typeArgs: List<PsiType>): PsiIm
 }
 
 class JeDeclaredType(
-        override val psiType: PsiClassType,
+        psiType: PsiClassType,
-        val psiClass: PsiClass,
+        psiClass: PsiClass,
         val enclosingDeclaredType: DeclaredType? = null,
         val isRaw: Boolean = false
-) : JePsiType(), JeTypeWithManager, DeclaredType {
+) : JePsiTypeBase<PsiClassType>(psiType, psiClass.manager), DeclaredType {
+    private val disposablePsiClass = psiClass.toDisposable()
+
+    val psiClass: PsiClass
+        get() = disposablePsiClass()
+
+    // JeElementRegistry registration is done in JePsiType
+    override fun dispose() {
+        super.dispose()
+        disposablePsiClass.dispose()
+    }
+
     override fun getKind() = TypeKind.DECLARED
     
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitDeclared(this, p)

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JeIntersectionType.kt

@@ -23,10 +23,10 @@ import javax.lang.model.type.TypeKind
 import javax.lang.model.type.TypeVisitor
 
 class JeIntersectionType(
-        override val psiType: PsiIntersectionType, 
+        psiType: PsiIntersectionType,
-        override val psiManager: PsiManager,
+        psiManager: PsiManager,
         private val isRaw: Boolean
-) : JePsiType(), JeTypeWithManager, IntersectionType {
+) : JePsiTypeBase<PsiIntersectionType>(psiType, psiManager), IntersectionType {
     override fun getKind() = TypeKind.INTERSECTION
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitIntersection(this, p)
 

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JeMethodExecutableTypeMirror.kt

@@ -21,6 +21,9 @@ import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiType
 import com.intellij.psi.util.MethodSignature
 import com.intellij.psi.util.PsiTypesUtil
+import org.jetbrains.kotlin.annotation.processing.impl.toDisposable
+import org.jetbrains.kotlin.annotation.processing.impl.dispose
+import org.jetbrains.kotlin.java.model.JeDisposablePsiElementOwner
 import org.jetbrains.kotlin.java.model.elements.getReceiverTypeMirror
 import javax.lang.model.type.ExecutableType
 import javax.lang.model.type.TypeKind
@@ -28,11 +31,26 @@ import javax.lang.model.type.TypeMirror
 import javax.lang.model.type.TypeVisitor
 
 class JeMethodExecutableTypeMirror(
-        val psi: PsiMethod,
+        psi: PsiMethod,
-        val signature: MethodSignature? = null,
+        signature: MethodSignature? = null,
-        val returnType: PsiType? = null,
+        returnType: PsiType? = null,
         val isRaw: Boolean = false
-) : JeTypeMirror, JeTypeWithManager, ExecutableType {
+) : JeDisposablePsiElementOwner<PsiMethod>(psi), JeTypeMirror, JeTypeWithManager, ExecutableType {
+    private val disposableSignature = signature?.toDisposable()
+    private val disposableReturnType = returnType?.toDisposable()
+
+    val signature: MethodSignature?
+        get() = disposableSignature?.invoke()
+
+    val returnType: PsiType?
+        get() = disposableReturnType?.invoke()
+
+    // JeElementRegistry registration is done in JeDisposablePsiElementOwner
+    override fun dispose() {
+        super.dispose()
+        dispose(disposableSignature, disposableReturnType)
+    }
+
     override fun getKind() = TypeKind.EXECUTABLE
     
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitExecutable(this, p)
@@ -43,9 +61,6 @@ class JeMethodExecutableTypeMirror(
 
     override fun getThrownTypes() = psi.throwsList.referencedTypes.map { it.toJeType(psi.manager, isRaw = isRaw) }
 
-    override val psiManager: PsiManager
-        get() = psi.manager
-
     override fun getParameterTypes(): List<TypeMirror> {
         signature?.parameterTypes?.let { types -> return types.map { it.toJeType(psi.manager, isRaw = isRaw) } }
         return psi.parameterList.parameters.map { it.type.toJeType(psi.manager, isRaw = isRaw) }

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JePrimitiveType.kt

@@ -22,7 +22,7 @@ import javax.lang.model.type.PrimitiveType
 import javax.lang.model.type.TypeKind
 import javax.lang.model.type.TypeVisitor
 
-class JePrimitiveType(override val psiType: PsiPrimitiveType) : JePsiType(), PrimitiveType {
+class JePrimitiveType(override val psiType: PsiPrimitiveType) : JePsiType, PrimitiveType {
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitPrimitive(this, p)
     
     override fun getKind() = when (psiType) {

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JePsiType.kt

@@ -16,8 +16,34 @@
 
 package org.jetbrains.kotlin.java.model.types
 
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.ServiceManager
+import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiType
+import org.jetbrains.kotlin.annotation.processing.impl.toDisposable
+import org.jetbrains.kotlin.java.model.internal.JeElementRegistry
+import org.jetbrains.kotlin.annotation.processing.impl.dispose
 
-abstract class JePsiType : JeTypeMirror {
+interface JePsiType : JeTypeMirror {
-    abstract val psiType: PsiType
+    val psiType: PsiType
+}
+
+abstract class JePsiTypeBase<out T : PsiType>(
+        psiType: T,
+        manager: PsiManager
+) : JePsiType, JeTypeWithManager, Disposable {
+    private val disposableManager = manager.toDisposable()
+    private val disposableType = psiType.toDisposable()
+
+    init {
+        ServiceManager.getService(manager.project, JeElementRegistry::class.java).register(this)
+    }
+
+    override fun dispose() = dispose(disposableManager, disposableType)
+
+    override val psiManager: PsiManager
+        get() = disposableManager()
+
+    override val psiType: T
+        get() = disposableType()
 }

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JeTypeVariableType.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.java.model.types
 
 import org.jetbrains.kotlin.java.model.elements.JeTypeParameterElement
 import com.intellij.psi.*
+import org.jetbrains.kotlin.annotation.processing.impl.toDisposable
 import org.jetbrains.kotlin.java.model.toJeElement
 import javax.lang.model.type.TypeKind
 import javax.lang.model.type.TypeMirror
@@ -25,9 +26,20 @@ import javax.lang.model.type.TypeVariable
 import javax.lang.model.type.TypeVisitor
 
 class JeTypeVariableType(
-        override val psiType: PsiClassType, 
+        psiType: PsiClassType,
-        val parameter: PsiTypeParameter
+        parameter: PsiTypeParameter
-) : JePsiType(), JeTypeWithManager, TypeVariable {
+) : JePsiTypeBase<PsiClassType>(psiType, parameter.manager), TypeVariable {
+    private val disposableParameter = parameter.toDisposable()
+
+    val parameter: PsiTypeParameter
+        get() = disposableParameter()
+
+    // JeElementRegistry registration is done in JeDisposablePsiElementOwner
+    override fun dispose() {
+        super.dispose()
+        disposableParameter.dispose()
+    }
+
     override fun getKind() = TypeKind.TYPEVAR
     
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitTypeVariable(this, p)

plugins/java-model-wrappers/src/org/jetbrains/kotlin/java/model/types/JeWildcardType.kt

@@ -24,18 +24,15 @@ import javax.lang.model.type.TypeVisitor
 import javax.lang.model.type.WildcardType
 
 class JeWildcardType(
-        override val psiType: PsiWildcardType, 
+        psiType: PsiWildcardType,
         private val isRaw: Boolean
-) : JePsiType(), JeTypeWithManager, WildcardType {
+) : JePsiTypeBase<PsiWildcardType>(psiType, psiType.manager), WildcardType {
     override fun getKind() = TypeKind.WILDCARD
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitWildcard(this, p)
     
     override fun getSuperBound() = psiType.superBound.toJeType(psiManager, isRaw = isRaw)
     override fun getExtendsBound() = psiType.extendsBound.toJeType(psiManager, isRaw = isRaw)
 
-    override val psiManager: PsiManager
-        get() = psiType.manager
-
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (other?.javaClass != javaClass) return false
@@ -61,10 +58,10 @@ class JeWildcardType(
 }
 
 class JeCapturedWildcardType(
-        override val psiType: PsiCapturedWildcardType, 
+        psiType: PsiCapturedWildcardType,
-        override val psiManager: PsiManager,
+        psiManager: PsiManager,
         private val isRaw: Boolean
-) : JePsiType(), JeTypeWithManager, WildcardType {
+) : JePsiTypeBase<PsiCapturedWildcardType>(psiType, psiManager), WildcardType {
     override fun getKind() = TypeKind.WILDCARD
     override fun <R : Any?, P : Any?> accept(v: TypeVisitor<R, P>, p: P) = v.visitWildcard(this, p)
 

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

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.annotation.processing.test.processor
 
+import com.intellij.testFramework.registerServiceInstance
 import org.jetbrains.kotlin.annotation.AbstractAnnotationProcessingExtension
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
@@ -26,6 +27,7 @@ import org.jetbrains.kotlin.java.model.elements.JeAnnotationMirror
 import org.jetbrains.kotlin.java.model.elements.JeMethodExecutableElement
 import org.jetbrains.kotlin.java.model.elements.JeTypeElement
 import org.jetbrains.kotlin.java.model.elements.JeVariableElement
+import org.jetbrains.kotlin.java.model.internal.JeElementRegistry
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import org.jetbrains.kotlin.test.ConfigurationKind
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -67,6 +69,8 @@ abstract class AbstractProcessorTest : AbstractBytecodeTextTest() {
         val apExtension = AnnotationProcessingExtensionForTests(processors)
         AnalysisCompletedHandlerExtension.registerExtension(project, apExtension)
 
+        project.registerServiceInstance(JeElementRegistry::class.java, JeElementRegistry())
+
         return environment
     }
 

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

@@ -122,7 +122,7 @@ class ProcessorTests : AbstractProcessorTest() {
         fun AnnotationMirror.getParam(name: String) = elementValues.entries.first { it.key.simpleName.toString() == name }.value
         assertTrue(anno2.getParam("a") is JeAnnotationAnnotationValue)
         assertTrue((anno2.getParam("b") as JeArrayAnnotationValue).value.first() is JeAnnotationAnnotationValue)
-        assertTrue(anno2.getParam("c") is JePrimitiveAnnotationValue)
+        assertTrue(anno2.getParam("c") is JePrimitiveAnnotationValue<*>)
         assertTrue(anno2.getParam("d") is JeTypeAnnotationValue)
         assertTrue((anno2.getParam("e") as JeArrayAnnotationValue).value.first() is JeTypeAnnotationValue)
     }

plugins/plugins-tests/tests/org/jetbrains/kotlin/annotation/processing/test/wrappers/AbstractAnnotationProcessingTest.kt

@@ -19,10 +19,12 @@ package org.jetbrains.kotlin.annotation.processing.test.wrappers
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.testFramework.registerServiceInstance
 import org.jetbrains.kotlin.codegen.AbstractBytecodeTextTest
 import org.jetbrains.kotlin.codegen.CodegenTestCase
 import org.jetbrains.kotlin.codegen.CodegenTestUtil
 import org.jetbrains.kotlin.java.model.elements.DefaultJeElementRenderer
+import org.jetbrains.kotlin.java.model.internal.JeElementRegistry
 import org.jetbrains.kotlin.java.model.toJeElement
 import org.jetbrains.kotlin.test.ConfigurationKind
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -38,6 +40,8 @@ abstract class AbstractAnnotationProcessingTest : AbstractBytecodeTextTest() {
 
     override fun doMultiFileTest(wholeFile: File, files: List<CodegenTestCase.TestFile>, javaFilesDir: File?) {
         createEnvironmentWithMockJdkAndIdeaAnnotations(ConfigurationKind.ALL, files, javaFilesDir)
+        myEnvironment.project.registerServiceInstance(JeElementRegistry::class.java, JeElementRegistry())
+
         loadMultiFiles(files)
 
         val text = FileUtil.loadFile(wholeFile, true)