Kapt: Do not retain references to Kotlin compiler classes in customized Javac components

This should fix potential memory leakages.

plugins/kapt3/kapt3-compiler/src/org/jetbrains/kotlin/kapt3/KaptContext.kt

@@ -20,6 +20,7 @@ import com.intellij.openapi.project.Project
 import com.sun.tools.javac.file.JavacFileManager
 import com.sun.tools.javac.jvm.ClassReader
 import com.sun.tools.javac.main.JavaCompiler
+import com.sun.tools.javac.tree.TreeMaker
 import com.sun.tools.javac.util.Context
 import com.sun.tools.javac.util.Options
 import org.jetbrains.kotlin.codegen.state.GenerationState
@@ -34,21 +35,22 @@ import org.jetbrains.org.objectweb.asm.tree.ClassNode
 import javax.tools.JavaFileManager
 
 class KaptContext<out GState : GenerationState?>(
-        val logger: KaptLogger,
-        val project: Project,
-        val bindingContext: BindingContext,
-        val compiledClasses: List<ClassNode>,
-        val origins: Map<Any, JvmDeclarationOrigin>,
-        val generationState: GState,
-        mapDiagnosticLocations: Boolean,
-        processorOptions: Map<String, String>,
-        javacOptions: Map<String, String> = emptyMap()
+    val logger: KaptLogger,
+    val project: Project,
+    val bindingContext: BindingContext,
+    val compiledClasses: List<ClassNode>,
+    val origins: Map<Any, JvmDeclarationOrigin>,
+    val generationState: GState,
+    mapDiagnosticLocations: Boolean,
+    processorOptions: Map<String, String>,
+    javacOptions: Map<String, String> = emptyMap()
 ) : AutoCloseable {
     val context = Context()
     val compiler: KaptJavaCompiler
     val fileManager: JavacFileManager
     val options: Options
     val javaLog: KaptJavaLog
+    val treeMaker: KaptTreeMaker
 
     init {
         KaptJavaLog.preRegister(this, logger.messageCollector, mapDiagnosticLocations)
@@ -81,9 +83,11 @@ class KaptContext<out GState : GenerationState?>(
         ClassReader.instance(context).saveParameterNames = true
 
         javaLog = compiler.log as KaptJavaLog
+        treeMaker = TreeMaker.instance(context) as KaptTreeMaker
     }
 
     override fun close() {
+        treeMaker.dispose()
         compiler.close()
         fileManager.close()
         generationState?.destroy()

plugins/kapt3/kapt3-compiler/src/org/jetbrains/kotlin/kapt3/javac/KaptJavaLog.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.kapt3.javac
 
-import com.intellij.openapi.project.Project
 import com.sun.tools.javac.tree.JCTree
 import com.sun.tools.javac.util.*
 import com.sun.tools.javac.util.JCDiagnostic.DiagnosticType
@@ -36,7 +35,7 @@ import javax.tools.SimpleJavaFileObject
 import com.sun.tools.javac.util.List as JavacList
 
 class KaptJavaLog(
-        private val project: Project,
+        private val projectBasePath: String?,
         context: Context,
         errWriter: PrintWriter,
         warnWriter: PrintWriter,
@@ -184,7 +183,7 @@ class KaptJavaLog(
 
     private fun getKotlinSourceFile(pos: KotlinPosition): File? {
         return if (pos.isRelativePath) {
-            val basePath = project.basePath
+            val basePath = this.projectBasePath
             if (basePath != null) File(basePath, pos.path) else null
         }
         else {
@@ -242,7 +241,7 @@ class KaptJavaLog(
                 val noticeWriter = makeWriter(WARNING)
 
                 KaptJavaLog(
-                    kaptContext.project, newContext, errWriter, warnWriter, noticeWriter,
+                    kaptContext.project.basePath, newContext, errWriter, warnWriter, noticeWriter,
                     interceptorData, mapDiagnosticLocations)
             })
         }

plugins/kapt3/kapt3-compiler/src/org/jetbrains/kotlin/kapt3/javac/KaptTreeMaker.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.kapt3.javac
 
+import com.intellij.openapi.Disposable
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiClass
 import com.intellij.psi.search.GlobalSearchScope
@@ -32,7 +33,9 @@ import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.Type.*
 import org.jetbrains.org.objectweb.asm.tree.ClassNode
 
-class KaptTreeMaker(context: Context, private val kaptContext: KaptContext<*>) : TreeMaker(context) {
+class KaptTreeMaker(context: Context, kaptContext: KaptContext<*>) : TreeMaker(context), Disposable {
+    private var kaptContext = DisposableReference(kaptContext)
+
     val nameTable: Name.Table = Names.instance(context).table
 
     fun Type(type: Type): JCTree.JCExpression {
@@ -70,6 +73,8 @@ class KaptTreeMaker(context: Context, private val kaptContext: KaptContext<*>) :
         // This is a top-level class
         if ('$' !in nameWithDots) return nameWithDots
 
+        val kaptContext = this.kaptContext.get()
+
         // Maybe it's in our sources?
         val classFromSources = kaptContext.compiledClasses.firstOrNull { it.name == internalName }
         if (classFromSources != null) {
@@ -161,9 +166,22 @@ class KaptTreeMaker(context: Context, private val kaptContext: KaptContext<*>) :
 
     fun name(name: String): Name = nameTable.fromString(name)
 
+    override fun dispose() {
+        kaptContext.dispose()
+    }
+
     companion object {
         internal fun preRegister(context: Context, kaptContext: KaptContext<*>) {
             context.put(treeMakerKey, Context.Factory<TreeMaker> { KaptTreeMaker(it, kaptContext) })
         }
     }
+}
+
+private class DisposableReference<T : Any>(obj: T) : Disposable {
+    private var obj: T? = obj
+    fun get() = obj!!
+
+    override fun dispose() {
+        obj = null
+    }
 }

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3Test.kt

@@ -208,7 +208,13 @@ open class AbstractClassFileToSourceStubConverterTest : AbstractKotlinKapt3Test(
         @JvmStatic
         fun main(args: Array<String>) {
             if (args.isEmpty()) error("1 argument expected, 0 passed")
-            AbstractClassFileToSourceStubConverterTest().doTest(args[0])
+            val test = AbstractClassFileToSourceStubConverterTest()
+            try {
+                test.setUp()
+                test.doTest(args[0])
+            } finally {
+                test.tearDown()
+            }
         }
     }