Move: Make MoveKotlinFileHandler stateless to prevent memory leaks. MoveKotlinTopLevelDeclarationsProcessor is made part of transient state, PackageNameInfo is recomputed based on file user data and container

idea/src/org/jetbrains/kotlin/idea/refactoring/move/moveFilesOrDirectories/MoveKotlinFileHandler.kt

@@ -17,14 +17,15 @@
 package org.jetbrains.kotlin.idea.refactoring.move.moveFilesOrDirectories
 
 import com.intellij.openapi.roots.JavaProjectRootsUtil
-import com.intellij.psi.PsiCompiledElement
-import com.intellij.psi.PsiDirectory
-import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiFile
+import com.intellij.psi.*
+import com.intellij.psi.impl.light.LightElement
 import com.intellij.refactoring.move.moveFilesOrDirectories.MoveFileHandler
 import com.intellij.refactoring.move.moveFilesOrDirectories.MoveFilesOrDirectoriesUtil
 import com.intellij.usageView.UsageInfo
+import org.jetbrains.kotlin.idea.JetLanguage
 import org.jetbrains.kotlin.idea.codeInsight.shorten.runWithElementsToShortenIsEmptyIgnored
+import org.jetbrains.kotlin.idea.core.getPackage
+import org.jetbrains.kotlin.idea.core.packageMatchesDirectory
 import org.jetbrains.kotlin.idea.refactoring.move.PackageNameInfo
 import org.jetbrains.kotlin.idea.refactoring.move.getInternalReferencesToUpdateOnPackageNameChange
 import org.jetbrains.kotlin.idea.refactoring.move.moveTopLevelDeclarations.DeferredJetFileKotlinMoveTarget
@@ -36,28 +37,37 @@ import org.jetbrains.kotlin.idea.refactoring.move.updatePackageDirective
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.JetFile
 import org.jetbrains.kotlin.psi.JetNamedDeclaration
-import org.jetbrains.kotlin.idea.core.getPackage
-import org.jetbrains.kotlin.idea.core.packageMatchesDirectory
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 public class MoveKotlinFileHandler : MoveFileHandler() {
-    private data class MoveContext(
-            val packageNameInfo: PackageNameInfo,
+    // This is special 'PsiElement' whose purpose is to wrap MoveKotlinTopLevelDeclarationsProcessor
+    // so that it can be kept in the transition map
+    class MoveContext(
+            psiManager: PsiManager,
             val declarationMoveProcessor: MoveKotlinTopLevelDeclarationsProcessor
-    )
+    ): LightElement(psiManager, JetLanguage.INSTANCE) {
+        override fun toString() = ""
+    }
 
-    private var moveContext: MoveContext? = null
-
-    private fun initMoveContext(psiFile: PsiFile, newParent: PsiDirectory): MoveContext? {
-        this.moveContext = null
-
-        if (psiFile !is JetFile || !(psiFile.updatePackageDirective ?: psiFile.packageMatchesDirectory())) return null
+    private fun JetFile.getPackageNameInfo(newParent: PsiDirectory, clearUserData: Boolean): PackageNameInfo? {
+        val shouldUpdatePackageDirective = updatePackageDirective ?: packageMatchesDirectory()
+        updatePackageDirective = if (clearUserData) null else shouldUpdatePackageDirective
 
+        if (!shouldUpdatePackageDirective) return null
         val newPackage = newParent.getPackage() ?: return null
 
-        val packageNameInfo = PackageNameInfo(psiFile.getPackageFqName(), FqName(newPackage.getQualifiedName()))
-        if (packageNameInfo.oldPackageName == packageNameInfo.newPackageName) return null
-        val project = psiFile.getProject()
+        val oldPackageName = getPackageFqName()
+        val newPackageName = FqName(newPackage.getQualifiedName())
+        if (oldPackageName == newPackageName) return null
 
+        return PackageNameInfo(oldPackageName, newPackageName)
+    }
+
+    private fun initMoveProcessor(psiFile: PsiFile, newParent: PsiDirectory): MoveKotlinTopLevelDeclarationsProcessor? {
+        if (psiFile !is JetFile) return null
+        val packageNameInfo = psiFile.getPackageNameInfo(newParent, false) ?: return null
+
+        val project = psiFile.getProject()
         val declarationMoveProcessor = MoveKotlinTopLevelDeclarationsProcessor(
                 project,
                 MoveKotlinTopLevelDeclarationsOptions(
@@ -70,10 +80,7 @@ public class MoveKotlinFileHandler : MoveFileHandler() {
                 ),
                 Mover.Idle
         )
-
-        val moveContext = MoveContext(packageNameInfo, declarationMoveProcessor)
-        this.moveContext = moveContext
-        return moveContext
+        return declarationMoveProcessor
     }
 
     override fun canProcessElement(element: PsiFile?): Boolean {
@@ -81,41 +88,31 @@ public class MoveKotlinFileHandler : MoveFileHandler() {
         return !JavaProjectRootsUtil.isOutsideJavaSourceRoot(element)
     }
 
-    override fun findUsages(
-            psiFile: PsiFile,
-            newParent: PsiDirectory,
-            searchInComments: Boolean,
-            searchInNonJavaFiles: Boolean
-    ): List<UsageInfo> {
-        val moveContext = initMoveContext(psiFile, newParent) ?: return emptyList()
-        return moveContext.declarationMoveProcessor.findUsages().toList()
-    }
+    override fun findUsages(psiFile: PsiFile, newParent: PsiDirectory, searchInComments: Boolean, searchInNonJavaFiles: Boolean) =
+            initMoveProcessor(psiFile, newParent)?.findUsages()?.toList() ?: emptyList()
 
-    override fun prepareMovedFile(file: PsiFile, moveDestination: PsiDirectory, oldToNewMap: Map<PsiElement, PsiElement>) {
-        initMoveContext(file, moveDestination)
+    override fun prepareMovedFile(file: PsiFile, moveDestination: PsiDirectory, oldToNewMap: MutableMap<PsiElement, PsiElement>) {
+        val moveProcessor = initMoveProcessor(file, moveDestination) ?: return
+        val moveContext = MoveContext(file.getManager(), moveProcessor)
+        oldToNewMap[moveContext] = moveContext
     }
 
     override fun updateMovedFile(file: PsiFile) {
         if (file !is JetFile) return
-        val moveContext = this.moveContext ?: return
 
-        file.updatePackageDirective = null
-        val packageNameInfo = moveContext.packageNameInfo
+        val newDirectory = file.getParent() ?: return
+        val packageNameInfo = file.getPackageNameInfo(newDirectory, true) ?: return
 
         val internalUsages = file.getInternalReferencesToUpdateOnPackageNameChange(packageNameInfo)
         file.getPackageDirective()?.setFqName(packageNameInfo.newPackageName)
         postProcessMoveUsages(internalUsages)
     }
 
-    override fun retargetUsages(usageInfos: List<UsageInfo>?, oldToNewMap: Map<PsiElement, PsiElement>?) {
-        val moveContext = this.moveContext ?: return
+    override fun retargetUsages(usageInfos: List<UsageInfo>?, oldToNewMap: Map<PsiElement, PsiElement>) {
+        val moveContext = oldToNewMap.keySet().firstIsInstanceOrNull<MoveContext>() ?: return
         val processor = moveContext.declarationMoveProcessor
         processor.project.runWithElementsToShortenIsEmptyIgnored {
-            try {
-                usageInfos?.let { processor.execute(it) }
-            } finally {
-                this.moveContext = null
-            }
+            usageInfos?.let { processor.execute(it) }
         }
     }
 }