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)
2016-09-27 17:31:54 +03:00
parent f83e7509b2
commit 46f3eae32c
26 changed files with 219 additions and 72 deletions
@@ -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

@@ -127,6 +128,8 @@ class AnnotationProcessingComponentRegistrar : ComponentRegistrar {
        val annotationProcessingExtension = ClasspathBasedAnnotationProcessingExtension(
                classpath, generatedOutputDirFile, classesOutputDir, javaRoots, verboseOutput, 
                incrementalDataFile, sourceRetentionAnnotationHandler)
+
+        project.registerService(JeElementRegistry::class.java, JeElementRegistry())
        
        AnalysisCompletedHandlerExtension.registerExtension(project, annotationProcessingExtension)
    }
@@ -31,6 +31,6 @@ class DisposableRef<out T : Any>(initialValue: T): Disposable {

 fun <T : Any> T.toDisposable() = DisposableRef(this)

-fun dispose(vararg refs: DisposableRef<*>) {
-    refs.forEach { it.dispose() }
+fun dispose(vararg refs: DisposableRef<*>?) {
+    refs.forEach { it?.dispose() }
 }
@@ -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.*
@@ -55,13 +57,15 @@ class KotlinProcessingEnvironment(
    internal val javaPsiFacade = javaPsiFacade.toDisposable()
    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()
-        elements.dispose()
+        types().dispose()
+        elements().dispose()
+        registry().dispose()
        dispose(elements, types, messager, filer, processors,
                project, psiManager, javaPsiFacade, projectScope, bindingContext,
                appendJavaSourceRootsHandler, options)
@@ -272,8 +272,8 @@ class KotlinTypes(
    }

    override fun isSubsignature(t1: ExecutableType, t2: ExecutableType): Boolean {
-        val m1 = (t1 as JeClassInitializerExecutableTypeMirror).initializer
-        val m2 = (t2 as JeClassInitializerExecutableTypeMirror).initializer
+        val m1 = (t1 as JeClassInitializerExecutableTypeMirror).psi
+        val m2 = (t2 as JeClassInitializerExecutableTypeMirror).psi
        
        // No parameters
        if (m1 is PsiClassInitializer && m2 is PsiClassInitializer) return true
@@ -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
@@ -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)
@@ -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
 }
@@ -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,
-        ExecutableElement, JeNoAnnotations, JeModifierListOwner
-{
+class JeClassInitializerExecutableElement(
+        psi: PsiClassInitializer
+) : JeAbstractElement<PsiClassInitializer>(psi), ExecutableElement, JeNoAnnotations, JeModifierListOwner {
    val isStaticInitializer = psi.isStatic

    override fun getEnclosingElement() = psi.containingClass?.let { JeTypeElement(it) }
@@ -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 {
@@ -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.JeModifierListOwner
-import org.jetbrains.kotlin.java.model.JeName
-import org.jetbrains.kotlin.java.model.JeNoAnnotations
+import org.jetbrains.kotlin.java.model.*
 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)
@@ -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
@@ -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
@@ -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) }
        }
@@ -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()
+    }
+}
@@ -23,10 +23,10 @@ import javax.lang.model.type.TypeKind
 import javax.lang.model.type.TypeVisitor

 class JeArrayType(
-        override val psiType: PsiArrayType, 
-        override val psiManager: PsiManager,
+        psiType: PsiArrayType,
+        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)
@@ -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 + "." } ?: "") +
-                              (if (initializer.isStatic) "<clinit>" else "<instinit>")
+    override fun toString() = (psi.containingClass?.qualifiedName?.let { it + "." } ?: "") +
+                              (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()
 }
@@ -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,
-        val psiClass: PsiClass,
+        psiType: PsiClassType,
+        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)
@@ -23,10 +23,10 @@ import javax.lang.model.type.TypeKind
 import javax.lang.model.type.TypeVisitor

 class JeIntersectionType(
-        override val psiType: PsiIntersectionType, 
-        override val psiManager: PsiManager,
+        psiType: PsiIntersectionType,
+        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)

@@ -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,
-        val signature: MethodSignature? = null,
-        val returnType: PsiType? = null,
+        psi: PsiMethod,
+        signature: MethodSignature? = 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) }
@@ -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) {
@@ -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 {
-    abstract val psiType: PsiType
+interface JePsiType : JeTypeMirror {
+    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()
 }
@@ -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, 
-        val parameter: PsiTypeParameter
-) : JePsiType(), JeTypeWithManager, TypeVariable {
+        psiType: PsiClassType,
+        parameter: PsiTypeParameter
+) : 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)
@@ -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, 
-        override val psiManager: PsiManager,
+        psiType: PsiCapturedWildcardType,
+        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)

@@ -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
    }

@@ -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)
    }
@@ -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)