IDE: add imports after pasting code

src/main/kotlin/org/rust/ide/typing/paste/RsImportCopyPasteProcessor.kt

@@ -0,0 +1,213 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.typing.paste
+
+import com.intellij.codeInsight.editorActions.CopyPastePostProcessor
+import com.intellij.codeInsight.editorActions.TextBlockTransferableData
+import com.intellij.openapi.application.runWriteAction
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.RangeMarker
+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.psi.PsiDocumentManager
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.impl.source.tree.injected.changesHandler.range
+import com.intellij.psi.util.parents
+import org.rust.ide.inspections.import.AutoImportFix
+import org.rust.ide.utils.import.import
+import org.rust.lang.core.psi.*
+import org.rust.lang.core.psi.ext.*
+import org.rust.lang.core.types.inference
+import org.rust.openapiext.toPsiFile
+import java.awt.datatransfer.DataFlavor
+import java.awt.datatransfer.Transferable
+import java.awt.datatransfer.UnsupportedFlavorException
+import java.io.IOException
+
+class RsTextBlockTransferableData(val offsetToFqnMap: Map<Int, String>) : TextBlockTransferableData, Cloneable {
+    override fun getFlavor(): DataFlavor? = RsImportCopyPasteProcessor.dataFlavor
+
+    override fun getOffsetCount(): Int = 0
+
+    override fun getOffsets(offsets: IntArray?, index: Int): Int = index
+    override fun setOffsets(offsets: IntArray?, index: Int): Int = index
+
+    public override fun clone(): RsTextBlockTransferableData {
+        try {
+            return super.clone() as RsTextBlockTransferableData
+        } catch (e: CloneNotSupportedException) {
+            throw RuntimeException()
+        }
+    }
+}
+
+class RsImportCopyPasteProcessor : CopyPastePostProcessor<RsTextBlockTransferableData>() {
+    override fun collectTransferableData(
+        file: PsiFile,
+        editor: Editor,
+        startOffsets: IntArray,
+        endOffsets: IntArray
+    ): List<RsTextBlockTransferableData> {
+        if (file !is RsFile || DumbService.getInstance(file.getProject()).isDumb) return emptyList()
+
+        val ranges = startOffsets.indices.map { TextRange(startOffsets[it], endOffsets[it]) }
+
+        val map = if (ranges.size == 1) {
+            createFqnMap(file, ranges[0])
+        } else {
+            emptyMap()
+        }
+
+        return listOf(
+            RsTextBlockTransferableData(map)
+        )
+    }
+
+    override fun extractTransferableData(content: Transferable): List<RsTextBlockTransferableData> {
+        try {
+            val data = content.getTransferData(dataFlavor) as? RsTextBlockTransferableData ?: return emptyList()
+            // copy to prevent changing of original by convertLineSeparators
+            return listOf(data.clone())
+        } catch (ignored: UnsupportedFlavorException) {
+        } catch (ignored: IOException) {
+        }
+        return emptyList()
+    }
+
+    override fun processTransferableData(
+        project: Project,
+        editor: Editor,
+        bounds: RangeMarker,
+        caretOffset: Int,
+        indented: Ref<in Boolean>,
+        values: List<RsTextBlockTransferableData>
+    ) {
+        PsiDocumentManager.getInstance(project).commitAllDocuments()
+
+        val data = values.getOrNull(0) ?: return
+        val rsFile = editor.document.toPsiFile(project) as? RsFile ?: return
+        val range = bounds.range
+
+        val elements = gatherElements(rsFile, range)
+        val importCtx = elements.firstOrNull { it is RsElement } as? RsElement ?: return
+
+        val visitor = ImportingVisitor(project, importCtx, range, data)
+
+        runWriteAction {
+            for (element in elements) {
+                element.accept(visitor)
+            }
+        }
+    }
+
+    companion object {
+        val dataFlavor: DataFlavor? by lazy {
+            try {
+                val dataClass = RsReferenceData::class.java
+                DataFlavor(
+                    DataFlavor.javaJVMLocalObjectMimeType + ";class=" + dataClass.name,
+                    "RsReferenceData",
+                    dataClass.classLoader
+                )
+            } catch (e: NoClassDefFoundError) {
+                null
+            } catch (e: IllegalArgumentException) {
+                null
+            }
+        }
+    }
+}
+
+private class RsReferenceData
+
+private class ImportingVisitor(
+    private val project: Project,
+    private val importCtx: RsElement,
+    private val range: TextRange,
+    private val data: RsTextBlockTransferableData
+) : RsRecursiveVisitor() {
+    override fun visitPath(path: RsPath) {
+        val ctx = AutoImportFix.findApplicableContext(project, path)
+        handleContext(path, ctx)
+        super.visitPath(path)
+    }
+
+    override fun visitMethodCall(methodCall: RsMethodCall) {
+        val ctx = AutoImportFix.findApplicableContext(project, methodCall)
+        handleContext(methodCall, ctx)
+        super.visitMethodCall(methodCall)
+    }
+
+    private fun handleContext(element: PsiElement, ctx: AutoImportFix.Context?) {
+        if (ctx != null) {
+            if (ctx.candidates.size == 1) {
+                ctx.candidates[0].import(importCtx)
+            } else {
+                val candidate = ctx.candidates.find {
+                    val offset = toRelativeOffset(element, range)
+                    val fqn = data.offsetToFqnMap[offset]
+                    fqn == it.qualifiedNamedItem.item.qualifiedName
+                }
+                candidate?.import(importCtx)
+            }
+        }
+    }
+}
+
+/**
+ * Records mapping between offsets (relative to copy/paste content range) and fully qualified names of resolved items
+ * from paths and method calls.
+ */
+private fun createFqnMap(file: RsFile, range: TextRange): Map<Int, String> {
+    val elements = gatherElements(file, range)
+    val map = mutableMapOf<Int, String>()
+
+    val visitor = object : RsRecursiveVisitor() {
+        override fun visitPath(path: RsPath) {
+            val target = (path.reference?.resolve() as? RsQualifiedNamedElement)?.qualifiedName
+            if (target != null) {
+                map[toRelativeOffset(path, range)] = target
+            }
+
+            super.visitPath(path)
+        }
+
+        override fun visitMethodCall(methodCall: RsMethodCall) {
+            val methods = methodCall.inference?.getResolvedMethod(methodCall)
+            val target = methods?.mapNotNull {
+                it.source.implementedTrait?.element?.qualifiedName
+            }?.firstOrNull()
+
+            if (target != null) {
+                map[toRelativeOffset(methodCall, range)] = target
+            }
+
+            super.visitMethodCall(methodCall)
+        }
+    }
+    for (element in elements) {
+        element.accept(visitor)
+    }
+
+    return map
+}
+
+private fun gatherElements(file: RsFile, range: TextRange): List<PsiElement> {
+    val leafElement = file.findElementAt(range.startOffset) ?: return emptyList()
+    val firstElement = leafElement
+        .parents(true)
+        .takeWhile { it.startOffset >= range.startOffset }
+        .lastOrNull() as? RsElement ?: return emptyList()
+    val siblings = firstElement
+        .rightSiblings
+        .takeWhile { it.endOffset <= range.endOffset }
+    return listOf(firstElement) + siblings
+}
+
+private fun toRelativeOffset(element: PsiElement, range: TextRange): Int = element.startOffset - range.startOffset

src/main/resources/META-INF/rust-core.xml

@@ -138,6 +138,9 @@
         <lang.smartEnterProcessor language="Rust"
                                   implementationClass="org.rust.ide.typing.assist.RsSmartEnterProcessor"/>
 
+        <!-- Copy paste processors -->
+        <copyPastePostProcessor implementation="org.rust.ide.typing.paste.RsImportCopyPasteProcessor"/>
+
         <!-- Imports -->
 
         <lang.importOptimizer language="Rust" implementationClass="org.rust.ide.refactoring.RsImportOptimizer"/>

src/test/kotlin/org/rust/FileTree.kt

@@ -181,7 +181,9 @@ class TestProject(
             else -> error("More than one file with carets found: $filesWithCaret")
         }
 
-    val fileWithCaretOrSelection: String get() = filesWithCaret.singleOrNull() ?: filesWithSelection.single()
+    val fileWithCaretOrSelection: String get() = filesWithCaret.singleOrNull() ?: fileWithSelection
+
+    val fileWithSelection: String get() = filesWithSelection.single()
 
     inline fun <reified T : PsiElement> findElementInFile(path: String): T {
         return doFindElementInFile(path, T::class.java)