IDE: add imports after pasting code

src/main/kotlin/org/rust/ide/settings/RsAutoImportOptions.kt

@@ -26,6 +26,10 @@ class RsAutoImportOptions : UiDslUnnamedConfigurable.Simple(), AutoImportOptions
                 checkBox(RsBundle.message("settings.rust.auto.import.on.completion"))
                     .bindSelected(settings::importOutOfScopeItems)
             }
+            row {
+                checkBox(RsBundle.message("settings.rust.auto.import.on.paste"))
+                    .bindSelected(settings::importOnPaste)
+            }
             row {
                 checkBox(ApplicationBundle.message("checkbox.add.unambiguous.imports.on.the.fly"))
                     .bindSelected(settings::addUnambiguousImportsOnTheFly)

src/main/kotlin/org/rust/ide/settings/RsCodeInsightSettings.kt

@@ -18,6 +18,7 @@ class RsCodeInsightSettings : PersistentStateComponent<RsCodeInsightSettings> {
     var importOutOfScopeItems: Boolean = true
     var suggestOutOfScopeItems: Boolean = true
     var addUnambiguousImportsOnTheFly: Boolean = false
+    var importOnPaste: Boolean = true
 
     override fun getState(): RsCodeInsightSettings = this
 

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

@@ -0,0 +1,223 @@
+/*
+ * 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.daemon.impl.CollectHighlightsUtil
+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 org.rust.ide.inspections.import.AutoImportFix
+import org.rust.ide.settings.RsCodeInsightSettings
+import org.rust.ide.utils.import.ImportCandidate
+import org.rust.ide.utils.import.import
+import org.rust.lang.core.psi.*
+import org.rust.lang.core.psi.ext.RsElement
+import org.rust.lang.core.psi.ext.RsQualifiedNamedElement
+import org.rust.lang.core.psi.ext.qualifiedName
+import org.rust.lang.core.psi.ext.startOffset
+import org.rust.lang.core.types.inference
+import org.rust.openapiext.toPsiFile
+import java.awt.datatransfer.DataFlavor
+import java.awt.datatransfer.Transferable
+
+data class ImportMap(private val offsetToFqnMap: Map<Int, String>) {
+    fun elementToFqn(element: PsiElement, range: TextRange): String? {
+        val offset = toRelativeOffset(element, range)
+        return offsetToFqnMap[offset]
+    }
+}
+
+class RsTextBlockTransferableData(val importMap: ImportMap) : TextBlockTransferableData {
+    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
+}
+
+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()
+        if (!RsCodeInsightSettings.getInstance().importOnPaste) return emptyList()
+
+        val startOffset = startOffsets.singleOrNull() ?: return emptyList()
+        val endOffset = endOffsets.singleOrNull() ?: return emptyList()
+        val range = TextRange(startOffset, endOffset)
+
+        // If the whole file is copied, it's not useful to add imports
+        if (range == file.textRange) return emptyList()
+
+        val map = createFqnMap(file, range)
+
+        return listOf(RsTextBlockTransferableData(map))
+    }
+
+    override fun extractTransferableData(content: Transferable): List<RsTextBlockTransferableData> {
+        try {
+            val data = content.getTransferData(dataFlavor) as? RsTextBlockTransferableData ?: return emptyList()
+            return listOf(data)
+        } catch (e: Throwable) {
+            return emptyList()
+        }
+    }
+
+    override fun processTransferableData(
+        project: Project,
+        editor: Editor,
+        bounds: RangeMarker,
+        caretOffset: Int,
+        indented: Ref<in Boolean>,
+        values: List<RsTextBlockTransferableData>
+    ) {
+        if (!RsCodeInsightSettings.getInstance().importOnPaste) return
+
+        PsiDocumentManager.getInstance(project).commitAllDocuments()
+
+        val data = values.getOrNull(0) ?: return
+        val file = editor.document.toPsiFile(project) as? RsFile ?: return
+        val range = bounds.range
+
+        val elements = gatherElements(file, range)
+        val importCtx = elements.firstOrNull { it is RsElement } as? RsElement ?: return
+
+        val visitor = ImportingVisitor(range, data.importMap)
+
+        runWriteAction {
+            for (element in elements) {
+                element.accept(visitor)
+            }
+            // We need to import the candidates after visiting all elements, otherwise the relative offsets could be
+            // invalidated after an import has been added
+            for (candidate in visitor.importCandidates) {
+                candidate.import(importCtx)
+            }
+        }
+    }
+
+    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 range: TextRange, private val importMap: ImportMap) : RsRecursiveVisitor() {
+    private val candidates: MutableList<ImportCandidate> = mutableListOf()
+
+    val importCandidates: List<ImportCandidate> = candidates
+
+    override fun visitPath(path: RsPath) {
+        val ctx = AutoImportFix.findApplicableContext(path)
+        handleContext(path, ctx)
+        super.visitPath(path)
+    }
+
+    override fun visitMethodCall(methodCall: RsMethodCall) {
+        val ctx = AutoImportFix.findApplicableContext(methodCall)
+        handleContext(methodCall, ctx)
+        super.visitMethodCall(methodCall)
+    }
+
+    override fun visitPatBinding(binding: RsPatBinding) {
+        if (importMap.elementToFqn(binding, range) != null) {
+            val ctx = AutoImportFix.findApplicableContext(binding)
+            handleContext(binding, ctx)
+        }
+        super.visitPatBinding(binding)
+    }
+
+    private fun handleContext(element: PsiElement, ctx: AutoImportFix.Context?) {
+        if (ctx != null) {
+            val candidate = ctx.candidates.find {
+                val fqn = importMap.elementToFqn(element, range)
+                fqn == it.qualifiedNamedItem.item.qualifiedName
+            }
+            if (candidate != null) {
+                candidates.add(candidate)
+            }
+        }
+    }
+}
+
+/**
+ * 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): ImportMap {
+    val elements = gatherElements(file, range)
+    val fqnMap = hashMapOf<Int, String>()
+
+    val visitor = object : RsRecursiveVisitor() {
+        override fun visitPath(path: RsPath) {
+            val target = (path.reference?.resolve() as? RsQualifiedNamedElement)?.qualifiedName
+            if (target != null) {
+                fqnMap[toRelativeOffset(path, range)] = target
+            }
+
+            super.visitPath(path)
+        }
+
+        override fun visitMethodCall(methodCall: RsMethodCall) {
+            val methods = methodCall.inference?.getResolvedMethod(methodCall)
+            val target = methods?.firstNotNullOfOrNull {
+                it.source.implementedTrait?.element?.qualifiedName
+            }
+
+            if (target != null) {
+                fqnMap[toRelativeOffset(methodCall, range)] = target
+            }
+
+            super.visitMethodCall(methodCall)
+        }
+
+        override fun visitPatBinding(binding: RsPatBinding) {
+            val target = (binding.reference.resolve() as? RsQualifiedNamedElement)?.qualifiedName
+            if (target != null) {
+                fqnMap[toRelativeOffset(binding, range)] = target
+            }
+            super.visitPatBinding(binding)
+        }
+    }
+    for (element in elements) {
+        element.accept(visitor)
+    }
+
+    return ImportMap(fqnMap)
+}
+
+private fun gatherElements(file: RsFile, range: TextRange): List<PsiElement> =
+    CollectHighlightsUtil.getElementsInRange(file, range.startOffset, range.endOffset)
+
+private fun toRelativeOffset(element: PsiElement, range: TextRange): Int = element.startOffset - range.startOffset

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

@@ -138,6 +138,7 @@
 
         <!-- Copy paste processors -->
         <copyPastePostProcessor implementation="org.rust.ide.typing.paste.RsConvertJsonToStructCopyPasteProcessor"/>
+        <copyPastePostProcessor implementation="org.rust.ide.typing.paste.RsImportCopyPasteProcessor"/>
 
         <!-- Imports -->
 

src/main/resources/messages/RsBundle.properties

@@ -105,6 +105,7 @@ run.target.rustc.executable.version.label=Rustc version:
 
 settings.rust.auto.import.on.completion=Import out-of-scope items on completion
 settings.rust.auto.import.show.popup=Show import popup
+settings.rust.auto.import.on.paste=Insert imports on paste
 settings.rust.auto.import.title=Rust
 
 settings.rust.cargo.auto.update.project.label=Update project automatically when Cargo.toml changes

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

@@ -223,7 +223,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)