Reworked k2k copy-paste action.

Get rid of resolve from copy phase, paste phase performs resolve in bg #KT-33939 Fixed
2019-12-13 23:41:51 +01:00
parent 9f15c0f4d7
commit 5730a21e7b
20 changed files with 603 additions and 149 deletions
@@ -9,6 +9,7 @@ import com.intellij.ide.highlighter.JavaFileType
 import com.intellij.openapi.actionSystem.IdeActions
 import com.intellij.openapi.application.runWriteAction
 import com.intellij.psi.PsiDocumentManager
+import com.intellij.util.ui.UIUtil
 import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.idea.conversion.copy.AbstractJavaToKotlinCopyPasteConversionTest
 import org.jetbrains.kotlin.idea.conversion.copy.ConvertJavaCopyPasteProcessor
@@ -118,6 +119,7 @@ abstract class AbstractPerformanceJavaToKotlinCopyPasteConversionTest(private va

    private fun perfTestCore() {
        myFixture.performEditorAction(IdeActions.ACTION_PASTE)
+        UIUtil.dispatchAllInvocationEvents()
    }

    private fun stats() = stats[j2kIndex()]
@@ -0,0 +1,52 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.codeInsight
+
+import com.intellij.codeInsight.editorActions.TextBlockTransferableData
+import com.intellij.openapi.util.TextRange
+import java.awt.datatransfer.DataFlavor
+import java.io.Serializable
+
+class BasicKotlinReferenceTransferableData(
+    val sourceFileUrl: String,
+    val packageName: String,
+    val imports: List<String>,
+    val sourceText: String,
+    val textRanges: List<TextRange>
+) : TextBlockTransferableData, Cloneable, Serializable {
+    override fun getFlavor() = dataFlavor
+    override fun getOffsetCount() = 0
+
+    override fun getOffsets(offsets: IntArray?, index: Int) = index
+    override fun setOffsets(offsets: IntArray?, index: Int) = index
+
+    companion object {
+        val dataFlavor: DataFlavor? by lazy {
+            try {
+                val dataClass = KotlinReferenceData::class.java
+                DataFlavor(
+                    DataFlavor.javaJVMLocalObjectMimeType + ";class=" + dataClass.name,
+                    "BasicKotlinReferenceTransferableData",
+                    dataClass.classLoader
+                )
+            } catch (e: NoClassDefFoundError) {
+                null
+            } catch (e: IllegalArgumentException) {
+                null
+            }
+        }
+    }
+
+    public override fun clone(): BasicKotlinReferenceTransferableData {
+        try {
+            return super.clone() as BasicKotlinReferenceTransferableData
+        } catch (e: CloneNotSupportedException) {
+            throw RuntimeException()
+        }
+    }
+
+}
+
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2015 JetBrains s.r.o.
+ * Copyright 2010-2020 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.
@@ -19,21 +19,25 @@ package org.jetbrains.kotlin.idea.codeInsight
 import com.intellij.codeInsight.CodeInsightSettings
 import com.intellij.codeInsight.editorActions.CopyPastePostProcessor
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.ModalityState
+import com.intellij.openapi.application.ReadAction.nonBlocking
 import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.RangeMarker
-import com.intellij.openapi.progress.ProcessCanceledException
+import com.intellij.openapi.progress.ProgressIndicator
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.progress.Task
 import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Ref
 import com.intellij.openapi.util.TextRange
+import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.psi.*
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.concurrency.AppExecutorUtil
 import org.jetbrains.annotations.TestOnly
-import org.jetbrains.kotlin.descriptors.CallableDescriptor
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
+import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.resolve.allowResolveInDispatchThread
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.resolveImportReference
@@ -45,19 +49,19 @@ import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.kdoc.KDocReference
 import org.jetbrains.kotlin.idea.references.*
 import org.jetbrains.kotlin.idea.util.ImportInsertHelper
-import org.jetbrains.kotlin.idea.util.application.runWriteAction
+import org.jetbrains.kotlin.idea.util.ProgressIndicatorUtils
+import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
 import org.jetbrains.kotlin.idea.util.getFileResolutionScope
+import org.jetbrains.kotlin.idea.util.getSourceRoot
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.kdoc.psi.api.KDocElement
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
-import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
-import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
-import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameOrNull
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.utils.findFunction
@@ -69,21 +73,15 @@ import java.awt.datatransfer.Transferable
 import java.awt.datatransfer.UnsupportedFlavorException
 import java.io.IOException
 import java.util.*
+import java.util.function.Consumer

-//NOTE: this class is based on CopyPasteReferenceProcessor and JavaCopyPasteReferenceProcessor
-class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReferenceTransferableData>() {
-    private val LOG = Logger.getInstance(KotlinCopyPasteReferenceProcessor::class.java)
+class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<BasicKotlinReferenceTransferableData>() {

-    private val IGNORE_REFERENCES_INSIDE: Array<Class<out KtElement>> = arrayOf(
-        KtImportList::class.java,
-        KtPackageDirective::class.java
-    )
-
-    override fun extractTransferableData(content: Transferable): List<KotlinReferenceTransferableData> {
+    override fun extractTransferableData(content: Transferable): List<BasicKotlinReferenceTransferableData> {
        if (CodeInsightSettings.getInstance().ADD_IMPORTS_ON_PASTE != CodeInsightSettings.NO) {
            try {
                val flavor = KotlinReferenceData.dataFlavor ?: return listOf()
-                val data = content.getTransferData(flavor) as? KotlinReferenceTransferableData ?: return listOf()
+                val data = content.getTransferData(flavor) as? BasicKotlinReferenceTransferableData ?: return listOf()
                // copy to prevent changing of original by convertLineSeparators
                return listOf(data.clone())
            } catch (ignored: UnsupportedFlavorException) {
@@ -99,30 +97,56 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        editor: Editor,
        startOffsets: IntArray,
        endOffsets: IntArray
-    ): List<KotlinReferenceTransferableData> {
+    ): List<BasicKotlinReferenceTransferableData> {
        if (file !is KtFile || DumbService.getInstance(file.getProject()).isDumb) return listOf()

-        val collectedData = try {
-            collectReferenceData(file, startOffsets, endOffsets)
-        } catch (e: ProcessCanceledException) {
-            // supposedly analysis can only be canceled from another thread
-            // do not log ProcessCanceledException as it is rethrown by IdeaLogger and code won't be copied
-            LOG.debug("ProcessCanceledException while analyzing references in ${file.getName()}. References can't be processed.")
-            return listOf()
-        } catch (e: Throwable) {
-            LOG.error("Exception in processing references for copy paste in file ${file.getName()}}", e)
-            return listOf()
+        check(startOffsets.size == endOffsets.size) {
+            "startOffsets ${startOffsets.contentToString()} has to have same size as endOffsets ${endOffsets.contentToString()}"
        }

-        if (collectedData.isEmpty()) return listOf()
+        val packageName = file.packageDirective?.fqName?.asString() ?: ""
+        val imports = file.importDirectives.map { it.text }
+        val endOfImportsOffset = file.importDirectives.map { it.endOffset }.max() ?: file.packageDirective?.endOffset ?: 0
+        val offsetDelta = if (startOffsets.any { it <= endOfImportsOffset }) 0 else endOfImportsOffset
+        val text = file.text.substring(offsetDelta)
+        val ranges = startOffsets.indices.map {
+            TextRange(startOffsets[it] - offsetDelta, endOffsets[it] - offsetDelta)
+        }

-        return listOf(KotlinReferenceTransferableData(collectedData.toTypedArray()))
+        return listOf(
+            BasicKotlinReferenceTransferableData(
+                sourceFileUrl = file.virtualFile.url,
+                packageName = packageName,
+                imports = imports,
+                sourceText = text,
+                textRanges = ranges
+            )
+        )
    }

+    private fun collectReferenceData(
+        textRanges: List<TextRange>,
+        file: KtFile,
+        targetFile: KtFile,
+        fakePackageName: String,
+        sourcePackageName: String
+    ): List<KotlinReferenceData> =
+        collectReferenceData(
+            file,
+            textRanges.map { it.startOffset }.toIntArray(),
+            textRanges.map { it.endOffset }.toIntArray(),
+            fakePackageName,
+            sourcePackageName,
+            targetFile.packageDirective?.fqName?.asString() ?: ""
+        )
+
    fun collectReferenceData(
        file: KtFile,
        startOffsets: IntArray,
-        endOffsets: IntArray
+        endOffsets: IntArray,
+        fakePackageName: String? = null,
+        sourcePackageName: String? = null,
+        targetPackageName: String? = null
    ): List<KotlinReferenceData> {
        val ranges = toTextRanges(startOffsets, endOffsets)
        val elementsByRange = ranges.associateBy({ it }, { textRange ->
@@ -130,60 +154,90 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        })

        val allElementsToResolve = elementsByRange.values.flatten().flatMap { it.collectDescendantsOfType<KtElement>() }
+        // TODO: allowResolveInDispatchThread could be dropped as soon as
+        //  ConvertJavaCopyPasteProcessor will perform it on non UI thread
        val bindingContext =
            allowResolveInDispatchThread {
                file.getResolutionFacade().analyze(allElementsToResolve, BodyResolveMode.PARTIAL)
            }

-        val result = ArrayList<KotlinReferenceData>()
+        val result = mutableListOf<KotlinReferenceData>()
        for ((range, elements) in elementsByRange) {
-            for (element in elements) {
-                result.addReferenceDataInsideElement(element, file, range.start, startOffsets, endOffsets, bindingContext)
+            elements.forEachDescendant { ktElement ->
+                result.addReferenceDataInsideElement(
+                    ktElement, file, range.start, ranges, bindingContext,
+                    fakePackageName = fakePackageName, sourcePackageName = sourcePackageName, targetPackageName = targetPackageName
+                )
            }
        }
        return result
    }

-    private fun MutableCollection<KotlinReferenceData>.addReferenceDataInsideElement(
-        element: PsiElement,
-        file: KtFile,
-        startOffset: Int,
-        startOffsets: IntArray,
-        endOffsets: IntArray,
-        bindingContext: BindingContext
-    ) {
-        if (PsiTreeUtil.getNonStrictParentOfType(element, *IGNORE_REFERENCES_INSIDE) != null) return
+    private fun List<PsiElement>.forEachDescendant(consumer: (KtElement) -> Unit) {
+        this.forEach { it.forEachDescendant(consumer) }
+    }

-        element.forEachDescendantOfType<KtElement>(canGoInside = { it::class.java as Class<*> !in IGNORE_REFERENCES_INSIDE }) { ktElement ->
-            val reference = ktElement.mainReference ?: return@forEachDescendantOfType
+    private fun PsiElement.forEachDescendant(consumer: (KtElement) -> Unit) {
+        if (PsiTreeUtil.getNonStrictParentOfType(this, *IGNORE_REFERENCES_INSIDE) != null) return

-            val descriptors = resolveReference(reference, bindingContext)
-            //check whether this reference is unambiguous
-            if (reference !is KtMultiReference<*> && descriptors.size > 1) return@forEachDescendantOfType
-
-            for (descriptor in descriptors) {
-                val effectiveReferencedDescriptors = DescriptorToSourceUtils.getEffectiveReferencedDescriptors(descriptor).asSequence()
-                val declaration = effectiveReferencedDescriptors
-                    .map { DescriptorToSourceUtils.getSourceFromDescriptor(it) }
-                    .singleOrNull()
-                if (declaration != null && declaration.isInCopiedArea(file, startOffsets, endOffsets)) continue
-
-                if (!reference.canBeResolvedViaImport(descriptor, bindingContext)) continue
-
-                val fqName = descriptor.importableFqName ?: continue
-                val kind = KotlinReferenceData.Kind.fromDescriptor(descriptor) ?: continue
-                val isQualifiable = KotlinReferenceData.isQualifiable(ktElement, descriptor)
-                val relativeStart = ktElement.range.start - startOffset
-                val relativeEnd = ktElement.range.end - startOffset
-                add(KotlinReferenceData(relativeStart, relativeEnd, fqName.asString(), isQualifiable, kind))
-            }
+        this.forEachDescendantOfType<KtElement>(canGoInside = {
+            it::class.java as Class<*> !in IGNORE_REFERENCES_INSIDE
+        }) { ktElement ->
+            consumer(ktElement)
        }
    }

-    private data class ReferenceToRestoreData(
-        val reference: KtReference,
-        val refData: KotlinReferenceData
-    )
+    private fun MutableCollection<KotlinReferenceData>.addReferenceDataInsideElement(
+        ktElement: KtElement,
+        file: KtFile,
+        startOffset: Int,
+        textRanges: List<TextRange>,
+        bindingContext: BindingContext,
+        fakePackageName: String? = null,
+        sourcePackageName: String? = null,
+        targetPackageName: String? = null
+    ) {
+        val reference = ktElement.mainReference ?: return
+
+        val descriptors = resolveReference(reference, bindingContext)
+        //check whether this reference is unambiguous
+        if (reference !is KtMultiReference<*> && descriptors.size > 1) return
+
+        for (descriptor in descriptors) {
+            val effectiveReferencedDescriptors =
+                DescriptorToSourceUtils.getEffectiveReferencedDescriptors(descriptor).asSequence()
+            val declaration = effectiveReferencedDescriptors
+                .map { DescriptorToSourceUtils.getSourceFromDescriptor(it) }
+                .singleOrNull()
+
+            if (declaration?.isInCopiedArea(file, textRanges) == true ||
+                !reference.canBeResolvedViaImport(descriptor, bindingContext)
+            ) continue
+
+            val importableFqName = descriptor.importableFqName ?: continue
+            val importableName = importableFqName.asString()
+            val pkgName = descriptor.findPackageFqNameSafe()?.asString() ?: ""
+            val importableShortName = importableFqName.shortName().asString()
+
+            val fqName = if (fakePackageName == pkgName) {
+                // It is possible to resolve unnecessary references from a target package (as we resolve it from a fake package)
+                if (sourcePackageName == targetPackageName && importableName == "$fakePackageName.$importableShortName") {
+                    continue
+                }
+                // roll back to original package name when we faced faked pkg name
+                sourcePackageName + importableName.substring(fakePackageName.length)
+            } else importableName
+
+            val kind = KotlinReferenceData.Kind.fromDescriptor(descriptor) ?: continue
+            val isQualifiable = KotlinReferenceData.isQualifiable(ktElement, descriptor)
+            val relativeStart = ktElement.range.start - startOffset
+            val relativeEnd = ktElement.range.end - startOffset
+            add(KotlinReferenceData(relativeStart, relativeEnd, fqName, isQualifiable, kind))
+        }
+    }
+
+    private data class ReferenceToRestoreData(val reference: KtReference, val refData: KotlinReferenceData)
+    private data class PsiElementByTextRange(val textRange: TextRange, val element: PsiElement)

    override fun processTransferableData(
        project: Project,
@@ -191,29 +245,237 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        bounds: RangeMarker,
        caretOffset: Int,
        indented: Ref<Boolean>,
-        values: List<KotlinReferenceTransferableData>
+        values: List<BasicKotlinReferenceTransferableData>
    ) {
-        if (DumbService.getInstance(project).isDumb) return
+        if (DumbService.getInstance(project).isDumb ||
+            CodeInsightSettings.getInstance().ADD_IMPORTS_ON_PASTE == CodeInsightSettings.NO
+        ) return

        val document = editor.document
        val file = PsiDocumentManager.getInstance(project).getPsiFile(document)
        if (file !is KtFile) return

-        val referenceData = values.single().data
+        processReferenceData(project, file, bounds.startOffset, values.single())
+    }

-        processReferenceData(project, file, bounds.startOffset, referenceData)
+    private fun processReferenceData(
+        project: Project,
+        file: KtFile,
+        blockStart: Int,
+        transferableData: BasicKotlinReferenceTransferableData
+    ) {
+        PsiDocumentManager.getInstance(project).commitAllDocuments()
+
+        // figure out candidate elements in UI thread in paste phase
+        // as psi elements could be changed later on - relative offsets are used for mapping purposes
+        val elementsByRange = transferableData.textRanges.flatMap { originalTextRange ->
+            val textRange = TextRange(blockStart, blockStart + originalTextRange.endOffset - originalTextRange.startOffset)
+            file.elementsInRange(textRange)
+                .filter { it is KtElement || it is KDocElement }
+                .map { PsiElementByTextRange(it.textRange, it) }
+        }
+
+        processReferenceData(project, file) { indicator: ProgressIndicator ->
+            findReferenceDataToRestore(
+                file,
+                blockStart,
+                indicator,
+                elementsByRange,
+                transferableData
+            )
+        }
    }

    fun processReferenceData(project: Project, file: KtFile, blockStart: Int, referenceData: Array<KotlinReferenceData>) {
-        PsiDocumentManager.getInstance(project).commitAllDocuments()
+        processReferenceData(project, file) { indicator: ProgressIndicator ->
+            findReferencesToRestore(file, blockStart, referenceData)
+        }
+    }

-        val referencesPossibleToRestore = findReferencesToRestore(file, blockStart, referenceData)
+    private fun processReferenceData(
+        project: Project,
+        file: KtFile,
+        findReferenceProvider: (indicator: ProgressIndicator) -> List<ReferenceToRestoreData>
+    ) {
+        val task: Task.Backgroundable = object : Task.Backgroundable(project, "Resolve pasted references", true) {
+            override fun run(indicator: ProgressIndicator) {
+                assert(!ApplicationManager.getApplication().isWriteAccessAllowed) {
+                    "Resolving references on dispatch thread leads to live lock"
+                }
+                ProgressIndicatorUtils.awaitWithCheckCanceled(
+                    nonBlocking<List<ReferenceToRestoreData>> {
+                        return@nonBlocking findReferenceProvider(indicator)
+                    }
+                        .finishOnUiThread(
+                            ModalityState.defaultModalityState(),
+                            Consumer<List<ReferenceToRestoreData>> { referencesPossibleToRestore ->
+                                val selectedReferencesToRestore =
+                                    showRestoreReferencesDialog(project, referencesPossibleToRestore)
+                                if (selectedReferencesToRestore.isEmpty()) return@Consumer

-        val selectedReferencesToRestore = showRestoreReferencesDialog(project, referencesPossibleToRestore)
-        if (selectedReferencesToRestore.isEmpty()) return
+                                project.executeWriteCommand("resolve pasted references") {
+                                    restoreReferences(selectedReferencesToRestore, file)
+                                }
+                            }
+                        )
+                        .withDocumentsCommitted(project)
+                        .submit(AppExecutorUtil.getAppExecutorService())
+                )
+            }
+        }
+        ProgressManager.getInstance().run(task)
+    }

-        runWriteAction {
-            restoreReferences(selectedReferencesToRestore, file)
+    private fun findReferenceDataToRestore(
+        file: PsiFile,
+        blockStart: Int,
+        indicator: ProgressIndicator,
+        elementsByRange: List<PsiElementByTextRange>,
+        transferableData: BasicKotlinReferenceTransferableData
+    ): List<ReferenceToRestoreData> {
+        if (file !is KtFile) return emptyList()
+
+        val referencesByRange = elementsByRange.flatMap { elementByTextRange ->
+            val element = elementByTextRange.element
+            if (!element.isValid) return@flatMap emptyList<Pair<TextRange, KtReference>>()
+            val refElementsRanges = mutableListOf<Pair<TextRange, KtReference>>()
+            val offsetDelta = element.startOffset - elementByTextRange.textRange.startOffset
+            element.forEachDescendant { ktElement ->
+                indicator.checkCanceled()
+
+                ktElement.mainReference?.let { ref ->
+                    val textRange = ref.element.textRange
+                    findReference(file, textRange)?.let {
+                        // remap reference to the original (as it was on paste phase) text range
+                        val range = if (offsetDelta == 0) textRange
+                        else
+                            TextRange(
+                                textRange.startOffset - offsetDelta,
+                                textRange.endOffset - offsetDelta
+                            )
+                        refElementsRanges.add(range to it)
+                    }
+                }
+            }
+            refElementsRanges
+        }.toMap()
+
+        val sourcePackageName = transferableData.packageName
+        val imports: List<String> = transferableData.imports
+
+        // Step 0. Recreate original source file (i.e. source file as it was on copy action) and resolve references from it
+        val ctxFile = sourceFile(file.project, transferableData) ?: file
+
+        // put original source file to some fake package to avoid ambiguous resolution ( a real file VS a virtual file )
+        val fakePackageName = "__kotlin.__some.__funky.__package"
+
+        val dummyOrigFileProlog =
+            """
+            package $fakePackageName
+            
+            ${buildDummySourceScope(sourcePackageName, imports, fakePackageName, file, transferableData, ctxFile)}
+             
+            """.trimIndent()
+
+        val dummyOriginalFile = KtPsiFactory(file.project)
+            .createAnalyzableFile(
+                "dummy-original.kt",
+                "$dummyOrigFileProlog${transferableData.sourceText}",
+                ctxFile
+            )
+
+        val offsetDelta = dummyOrigFileProlog.length
+
+        val dummyOriginalFileTextRanges =
+            // it is required as it is shifted by dummy prolog
+            transferableData.textRanges.map { TextRange(it.startOffset + offsetDelta, it.endOffset + offsetDelta) }
+
+        // Step 1. Find references in copied blocks of (recreated) source file
+        val sourceFileBasedReferences =
+            collectReferenceData(dummyOriginalFileTextRanges, dummyOriginalFile, file, fakePackageName, sourcePackageName)
+
+        indicator.checkCanceled()
+
+        // Step 2. Find references to restore in a target file
+        return findReferencesToRestore(file, blockStart, sourceFileBasedReferences, referencesByRange)
+    }
+
+    private fun buildDummySourceScope(
+        sourcePkgName: String,
+        imports: List<String>,
+        fakePkgName: String,
+        file: KtFile,
+        transferableData: BasicKotlinReferenceTransferableData,
+        ctxFile: KtFile
+    ): String {
+        // it could be that copied block contains inner classes or enums
+        // to solve this problem a special file is build:
+        // it contains imports from source package replaced with a fake package prefix
+        //
+        // result scope has to contain
+        // - those fake package imports those are successfully resolved (i.e. present in copied block)
+        // - those source package imports those are not present in a fake package
+        // - all rest imports
+
+        val sourceImportPrefix = "import $sourcePkgName"
+        val fakeImportPrefix = "import $fakePkgName"
+
+        val affectedSourcePkgImports = imports.filter { it.startsWith(sourceImportPrefix) }
+        val fakePkgImports = affectedSourcePkgImports.map { fakeImportPrefix + it.substring(sourceImportPrefix.length) }
+
+        fun joinLines(items: Collection<String>) = items.joinToString("\n")
+
+        val dummyImportsFile = KtPsiFactory(file.project)
+            .createAnalyzableFile(
+                "dummy-imports.kt",
+                "package $fakePkgName\n" +
+                        "${joinLines(fakePkgImports)}\n" +
+                        transferableData.sourceText,
+                ctxFile
+            )
+
+        val dummyFileImports = dummyImportsFile.collectDescendantsOfType<KtImportDirective>().mapNotNull { directive ->
+            val importedReference =
+                directive.importedReference?.getQualifiedElementSelector()?.mainReference?.resolve() as? KtNamedDeclaration
+            importedReference?.let { directive.text }
+        }
+
+        val dummyFileImportsSet = dummyFileImports.toSet()
+        val filteredImports = imports.filter {
+            !it.startsWith(sourceImportPrefix) ||
+                    !dummyFileImportsSet.contains(fakeImportPrefix + it.substring(sourceImportPrefix.length))
+        }
+
+        return """
+            ${joinLines(dummyFileImports)}
+            import ${sourcePkgName}.*
+            ${joinLines(filteredImports)}
+        """
+    }
+
+    private fun sourceFile(project: Project, transferableData: BasicKotlinReferenceTransferableData): KtFile? {
+        val sourceFile = VirtualFileManager.getInstance().findFileByUrl(transferableData.sourceFileUrl) ?: return null
+        if (sourceFile.getSourceRoot(project) == null) return null
+
+        return PsiManager.getInstance(project).findFile(sourceFile) as? KtFile
+    }
+
+    private fun filterReferenceData(
+        refData: KotlinReferenceData,
+        fileResolutionScope: LexicalScope
+    ): Boolean {
+        if (refData.isQualifiable) return true
+
+        val originalFqName = FqName(refData.fqName)
+        val name = originalFqName.shortName()
+        return when (refData.kind) {
+            // filter if function is already imported
+            KotlinReferenceData.Kind.FUNCTION -> fileResolutionScope
+                .findFunction(name, NoLookupLocation.FROM_IDE) { it.importableFqName == originalFqName } == null
+            // filter if property is already imported
+            KotlinReferenceData.Kind.PROPERTY -> fileResolutionScope
+                .findVariable(name, NoLookupLocation.FROM_IDE) { it.importableFqName == originalFqName } == null
+            else -> true
        }
    }

@@ -221,33 +483,52 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        file: PsiFile,
        blockStart: Int,
        referenceData: Array<out KotlinReferenceData>
-    ): List<ReferenceToRestoreData> {
-        if (file !is KtFile) return listOf()
+    ): List<ReferenceToRestoreData> = findReferences(file, referenceData.map { it to findReference(it, file as KtFile, blockStart) })
+
+    private fun findReferencesToRestore(
+        file: PsiFile,
+        blockStart: Int,
+        referenceData: List<KotlinReferenceData>,
+        referencesByPosition: Map<TextRange, KtReference>
+    ): List<ReferenceToRestoreData> =
+        // use already found reference candidates - so file could be changed
+        findReferences(file, referenceData
+            .map {
+                val textRange = TextRange(it.startOffset + blockStart, it.endOffset + blockStart)
+                val reference = referencesByPosition[textRange]
+                it to if (reference?.element?.isValid == true) reference else null
+            })
+
+    private fun findReferences(
+        file: PsiFile,
+        references: List<Pair<KotlinReferenceData, KtReference?>>
+    ): List<ReferenceToRestoreData> {
+        val ktFile = file as? KtFile ?: return listOf()

-        val references = referenceData.map { it to findReference(it, file, blockStart) }
        val bindingContext = try {
-            file.getResolutionFacade().analyze(references.mapNotNull { it.second?.element }, BodyResolveMode.PARTIAL)
+            ktFile.getResolutionFacade().analyze(references.mapNotNull { it.second?.element }, BodyResolveMode.PARTIAL)
        } catch (e: Throwable) {
            if (e is ControlFlowException) throw e
            LOG.error("Failed to analyze references after copy paste", e)
            return emptyList()
        }
-        val fileResolutionScope = file.getResolutionFacade().getFileResolutionScope(file)
+        val fileResolutionScope = ktFile.getResolutionFacade().getFileResolutionScope(ktFile)
        return references.mapNotNull { pair ->
            val data = pair.first
            val reference = pair.second
-            if (reference != null)
-                createReferenceToRestoreData(reference, data, file, fileResolutionScope, bindingContext)
-            else
-                null
+            reference?.let { createReferenceToRestoreData(it, data, ktFile, fileResolutionScope, bindingContext) }
        }
    }

-    private fun findReference(data: KotlinReferenceData, file: KtFile, blockStart: Int): KtReference? {
-        val startOffset = data.startOffset + blockStart
-        val endOffset = data.endOffset + blockStart
-        val element = file.findElementAt(startOffset) ?: return null
-        val desiredRange = TextRange(startOffset, endOffset)
+    private fun findReference(data: KotlinReferenceData, file: KtFile, blockStart: Int): KtReference? =
+        findReference(file, TextRange(data.startOffset + blockStart, data.endOffset + blockStart))
+
+    private fun findReference(file: KtFile, desiredRange: TextRange): KtReference? {
+        val element = file.findElementAt(desiredRange.startOffset) ?: return null
+        return findReference(element, desiredRange)
+    }
+
+    private fun findReference(element: PsiElement, desiredRange: TextRange): KtReference? {
        for (current in element.parentsWithSelf) {
            val range = current.range
            if (current is KtElement && range == desiredRange) {
@@ -265,21 +546,9 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        fileResolutionScope: LexicalScope,
        bindingContext: BindingContext
    ): ReferenceToRestoreData? {
+        if (!filterReferenceData(refData, fileResolutionScope)) return null
+
        val originalFqName = FqName(refData.fqName)
-        val name = originalFqName.shortName()
-
-        if (!refData.isQualifiable) {
-            if (refData.kind == KotlinReferenceData.Kind.FUNCTION) {
-                if (fileResolutionScope.findFunction(name, NoLookupLocation.FROM_IDE) { it.importableFqName == originalFqName } != null) {
-                    return null // already imported
-                }
-            } else if (refData.kind == KotlinReferenceData.Kind.PROPERTY) {
-                if (fileResolutionScope.findVariable(name, NoLookupLocation.FROM_IDE) { it.importableFqName == originalFqName } != null) {
-                    return null // already imported
-                }
-            }
-        }
-
        val referencedDescriptors = resolveReference(reference, bindingContext)
        val referencedFqNames = referencedDescriptors
            .filterNot { ErrorUtils.isError(it) }
@@ -295,14 +564,15 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        return ReferenceToRestoreData(reference, refData)
    }

-    private fun resolveReference(reference: KtReference, bindingContext: BindingContext): Collection<DeclarationDescriptor> {
+    private fun resolveReference(reference: KtReference, bindingContext: BindingContext): List<DeclarationDescriptor> {
        val element = reference.element
        if (element is KtNameReferenceExpression && reference is KtSimpleNameReference) {
-            bindingContext[BindingContext.SHORT_REFERENCE_TO_COMPANION_OBJECT, element]
-                ?.let { return listOf(it) }
+            val classifierDescriptor =
+                bindingContext[BindingContext.SHORT_REFERENCE_TO_COMPANION_OBJECT, element]
+            (classifierDescriptor ?: bindingContext[BindingContext.REFERENCE_TARGET, element])?.let { return listOf(it) }
        }

-        return reference.resolveToDescriptors(bindingContext)
+        return reference.resolveToDescriptors(bindingContext).toList()
    }

    private fun restoreReferences(referencesToRestore: Collection<ReferenceToRestoreData>, file: KtFile) {
@@ -318,11 +588,13 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        val descriptorsToImport = ArrayList<DeclarationDescriptor>()

        for ((reference, refData) in referencesToRestore) {
+            if (!reference.element.isValid) continue
            val fqName = FqName(refData.fqName)

            if (refData.isQualifiable) {
                if (reference is KtSimpleNameReference) {
-                    val pointer = smartPointerManager.createSmartPsiElementPointer(reference.element, file)
+                    val pointer =
+                        smartPointerManager.createSmartPsiElementPointer(reference.element, file)
                    bindingRequests.add(BindingRequest(pointer, fqName))
                } else if (reference is KDocReference) {
                    descriptorsToImport.addAll(findImportableDescriptors(fqName, file))
@@ -351,11 +623,19 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
    private fun findCallableToImport(fqName: FqName, file: KtFile): CallableDescriptor? =
        findImportableDescriptors(fqName, file).firstIsInstanceOrNull()

+    private tailrec fun DeclarationDescriptor.findPackageFqNameSafe(): FqName? {
+        return when {
+            this is PackageFragmentDescriptor -> this.fqNameOrNull()
+            containingDeclaration == null -> this.fqNameOrNull()
+            else -> this.containingDeclaration!!.findPackageFqNameSafe()
+        }
+    }
+
    private fun showRestoreReferencesDialog(
        project: Project,
        referencesToRestore: List<ReferenceToRestoreData>
    ): Collection<ReferenceToRestoreData> {
-        val fqNames = referencesToRestore.asSequence().map { it.refData.fqName }.toSortedSet()
+        val fqNames = referencesToRestore.map { it.refData.fqName }.toSortedSet()

        if (ApplicationManager.getApplication().isUnitTestMode) {
            declarationsToImportSuggested = fqNames
@@ -378,13 +658,21 @@ class KotlinCopyPasteReferenceProcessor : CopyPastePostProcessor<KotlinReference
        return startOffsets.indices.map { TextRange(startOffsets[it], endOffsets[it]) }
    }

-    private fun PsiElement.isInCopiedArea(fileCopiedFrom: KtFile, startOffsets: IntArray, endOffsets: IntArray): Boolean {
+    private fun PsiElement.isInCopiedArea(fileCopiedFrom: KtFile, textRanges: List<TextRange>): Boolean {
        if (containingFile != fileCopiedFrom) return false
-        return toTextRanges(startOffsets, endOffsets).any { this.range in it }
+        return textRanges.any { this.range in it }
    }

    companion object {
        @get:TestOnly
        var declarationsToImportSuggested: Collection<String> = emptyList()
+
+        private val LOG = Logger.getInstance(KotlinCopyPasteReferenceProcessor::class.java)
+
+        private val IGNORE_REFERENCES_INSIDE: Array<Class<out KtElement>> = arrayOf(
+            KtImportList::class.java,
+            KtPackageDirective::class.java
+        )
    }
+
 }
@@ -16,7 +16,6 @@

 package org.jetbrains.kotlin.idea.codeInsight

-import com.intellij.codeInsight.editorActions.TextBlockTransferableData
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.psi.KtCallableReferenceExpression
@@ -30,36 +29,7 @@ import org.jetbrains.kotlin.types.expressions.DoubleColonLHS
 import java.awt.datatransfer.DataFlavor
 import java.io.Serializable

-class KotlinReferenceTransferableData(
-    val data: Array<KotlinReferenceData>
-) : TextBlockTransferableData, Cloneable, Serializable {
-
-    override fun getFlavor() = KotlinReferenceData.dataFlavor
-
-    override fun getOffsetCount() = data.size * 2
-
-    override fun getOffsets(offsets: IntArray, index: Int): Int {
-        var i = index
-        for (d in data) {
-            offsets[i++] = d.startOffset
-            offsets[i++] = d.endOffset
-        }
-        return i
-    }
-
-    override fun setOffsets(offsets: IntArray, index: Int): Int {
-        var i = index
-        for (d in data) {
-            d.startOffset = offsets[i++]
-            d.endOffset = offsets[i++]
-        }
-        return i
-    }
-
-    public override fun clone() = KotlinReferenceTransferableData(Array(data.size) { data[it].clone() })
-}
-
-class KotlinReferenceData(
+data class KotlinReferenceData(
    var startOffset: Int,
    var endOffset: Int,
    val fqName: String,
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.util
+
+import java.util.concurrent.Future
+
+object ProgressIndicatorUtils {
+    @JvmStatic
+    fun <T> awaitWithCheckCanceled(future: Future<T>) =
+        com.intellij.openapi.progress.util.BackgroundTaskUtil.awaitWithCheckCanceled(future)
+}
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.util
+
+import java.util.concurrent.Future
+
+object ProgressIndicatorUtils {
+    @JvmStatic
+    fun <T> awaitWithCheckCanceled(future: Future<T>) =
+        com.intellij.openapi.progress.util.ProgressIndicatorUtils.awaitWithCheckCanceled(future)
+}
@@ -0,0 +1,3 @@
+package foo
+
+class Bar(val v: Int){}
@@ -0,0 +1,20 @@
+package to
+
+import d.A
+
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+package d
+
+/*
+ * Some comment in front of imports
+ */
+import foo.Bar
+
+/**
+ * Some kdoc for class Foo
+ */
+class Foo(val a: A) {
+}
@@ -0,0 +1 @@
+d.A
@@ -0,0 +1,19 @@
+// WITH_LIBRARY: copyPaste/imports/KotlinLibrary
+<selection>
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+package d
+
+/*
+ * Some comment in front of imports
+ */
+import foo.Bar
+
+/**
+ * Some kdoc for class Foo
+ */
+class Foo(val a: A) {
+}
+</selection>
@@ -0,0 +1,2 @@
+package to
+<caret>
@@ -0,0 +1,3 @@
+package foo
+
+class Bar(val v: Int){}
@@ -0,0 +1,15 @@
+package to
+
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+import d.A
+import foo.Bar
+
+/**
+ * Some kdoc for class Foo
+ */
+class Foo(val a: A) {
+}
@@ -0,0 +1 @@
+d.A
@@ -0,0 +1,16 @@
+// WITH_LIBRARY: copyPaste/imports/KotlinLibrary
+package d
+<selection>
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+import foo.Bar
+
+/**
+ * Some kdoc for class Foo
+ */
+class Foo(val a: A) {
+}
+</selection>
@@ -0,0 +1,2 @@
+package to
+<caret>
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.codeInsight

 import com.intellij.openapi.actionSystem.IdeActions
 import com.intellij.openapi.util.io.FileUtil
+import com.intellij.util.ui.UIUtil
 import org.jetbrains.kotlin.idea.AbstractCopyPasteTest
 import org.jetbrains.kotlin.idea.test.dumpTextWithErrors
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
@@ -15,6 +16,7 @@ import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils

 abstract class AbstractInsertImportOnPasteTest : AbstractCopyPasteTest() {
+    private val TODO_INVESTIGATE_DIRECTIVE = "// TODO: Investigation is required"
    private val NO_ERRORS_DUMP_DIRECTIVE = "// NO_ERRORS_DUMP"
    private val DELETE_DEPENDENCIES_BEFORE_PASTE_DIRECTIVE = "// DELETE_DEPENDENCIES_BEFORE_PASTE"

@@ -48,6 +50,12 @@ abstract class AbstractInsertImportOnPasteTest : AbstractCopyPasteTest() {

        configureTargetFile(testFileName.replace(".kt", ".to.kt"))
        performNotWriteEditorAction(IdeActions.ACTION_PASTE)
+        UIUtil.dispatchAllInvocationEvents()
+
+        if (InTextDirectivesUtils.isDirectiveDefined(testFileText, TODO_INVESTIGATE_DIRECTIVE)) {
+            println("File $testFile has $TODO_INVESTIGATE_DIRECTIVE")
+            return
+        }

        val namesToImportDump = KotlinCopyPasteReferenceProcessor.declarationsToImportSuggested.joinToString("\n")
        KotlinTestUtils.assertEqualsToFile(testDataFile(testFileName.replace(".kt", ".expected.names")), namesToImportDump)
@@ -9,6 +9,7 @@ import com.intellij.codeInsight.actions.OptimizeImportsProcessor
 import com.intellij.openapi.actionSystem.IdeActions
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.psi.PsiDocumentManager
+import com.intellij.util.ui.UIUtil
 import junit.framework.TestCase
 import org.jetbrains.kotlin.idea.AbstractCopyPasteTest
 import org.jetbrains.kotlin.idea.core.moveCaret
@@ -48,6 +49,7 @@ abstract class AbstractMoveOnCutPasteTest : AbstractCopyPasteTest() {
        val targetFileExists = testDataFile(targetFileName).exists()
        val targetPsiFile = if (targetFileExists) configureTargetFile(targetFileName) else null
        performNotWriteEditorAction(IdeActions.ACTION_PASTE)
+        UIUtil.dispatchAllInvocationEvents()

        val shouldBeAvailable = InTextDirectivesUtils.getPrefixedBoolean(testFileText, IS_AVAILABLE_DIRECTIVE) ?: true
        val cookie = editor.getUserData(MoveDeclarationsEditorCookie.KEY)
@@ -130,6 +130,11 @@ public class InsertImportOnPasteTestGenerated extends AbstractInsertImportOnPast
            runTest("idea/testData/copyPaste/imports/DependencyOnStdLib.kt");
        }

+        @TestMetadata("EntireFile.kt")
+        public void testEntireFile() throws Exception {
+            runTest("idea/testData/copyPaste/imports/EntireFile.kt");
+        }
+
        @TestMetadata("EnumEntries.kt")
        public void testEnumEntries() throws Exception {
            runTest("idea/testData/copyPaste/imports/EnumEntries.kt");
@@ -195,6 +200,11 @@ public class InsertImportOnPasteTestGenerated extends AbstractInsertImportOnPast
            runTest("idea/testData/copyPaste/imports/ImportDirective.kt");
        }

+        @TestMetadata("ImportDirectiveAndClassBody.kt")
+        public void testImportDirectiveAndClassBody() throws Exception {
+            runTest("idea/testData/copyPaste/imports/ImportDirectiveAndClassBody.kt");
+        }
+
        @TestMetadata("ImportableEntityInExtensionLiteral.kt")
        public void testImportableEntityInExtensionLiteral() throws Exception {
            runTest("idea/testData/copyPaste/imports/ImportableEntityInExtensionLiteral.kt");
@@ -473,6 +483,11 @@ public class InsertImportOnPasteTestGenerated extends AbstractInsertImportOnPast
            runTest("idea/testData/copyPaste/imports/DependencyOnStdLib.kt");
        }

+        @TestMetadata("EntireFile.kt")
+        public void testEntireFile() throws Exception {
+            runTest("idea/testData/copyPaste/imports/EntireFile.kt");
+        }
+
        @TestMetadata("EnumEntries.kt")
        public void testEnumEntries() throws Exception {
            runTest("idea/testData/copyPaste/imports/EnumEntries.kt");
@@ -538,6 +553,11 @@ public class InsertImportOnPasteTestGenerated extends AbstractInsertImportOnPast
            runTest("idea/testData/copyPaste/imports/ImportDirective.kt");
        }

+        @TestMetadata("ImportDirectiveAndClassBody.kt")
+        public void testImportDirectiveAndClassBody() throws Exception {
+            runTest("idea/testData/copyPaste/imports/ImportDirectiveAndClassBody.kt");
+        }
+
        @TestMetadata("ImportableEntityInExtensionLiteral.kt")
        public void testImportableEntityInExtensionLiteral() throws Exception {
            runTest("idea/testData/copyPaste/imports/ImportableEntityInExtensionLiteral.kt");
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.conversion.copy

 import com.intellij.openapi.actionSystem.IdeActions
+import com.intellij.util.ui.UIUtil
 import org.jetbrains.kotlin.idea.editor.KotlinEditorOptions
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
@@ -53,6 +54,7 @@ abstract class AbstractJavaToKotlinCopyPasteConversionTest : AbstractJ2kCopyPast
        ConvertJavaCopyPasteProcessor.conversionPerformed = false

        myFixture.performEditorAction(IdeActions.ACTION_PASTE)
+        UIUtil.dispatchAllInvocationEvents()

        assertEquals(
            noConversionExpected, !ConvertJavaCopyPasteProcessor.conversionPerformed,