INSP: add attach file to module quick fix

src/main/kotlin/org/rust/ide/inspections/RsDetachedFileInspection.kt

@@ -0,0 +1,85 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.inspections
+
+import com.intellij.codeInspection.InspectionManager
+import com.intellij.codeInspection.ProblemDescriptor
+import com.intellij.codeInspection.ProblemHighlightType
+import com.intellij.codeInspection.SuppressQuickFix
+import com.intellij.ide.util.PropertiesComponent
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import org.rust.cargo.project.model.cargoProjects
+import org.rust.ide.inspections.fixes.AttachFileToModuleFix
+import org.rust.lang.core.psi.RsFile
+
+class RsDetachedFileInspection : RsLocalInspectionTool() {
+    override fun checkFile(file: PsiFile, manager: InspectionManager, isOnTheFly: Boolean): Array<ProblemDescriptor>? {
+        val rsFile = file as? RsFile ?: return null
+        val project = file.project
+
+        if (!isInspectionEnabled(project, file.virtualFile)) return null
+
+        val cargoProjects = project.cargoProjects
+        if (!cargoProjects.initialized) return null
+
+        // Handled by [NoCargoProjectNotificationProvider]
+        if (cargoProjects.findProjectForFile(file.virtualFile) == null) return null
+
+        if (rsFile.crateRoot == null) {
+            val availableModules = AttachFileToModuleFix.findAvailableModulesForFile(project, rsFile)
+            val attachFix = if (availableModules.isNotEmpty()) {
+                val moduleLabel = if (availableModules.size == 1) {
+                    availableModules[0].name
+                } else {
+                    null
+                }
+                AttachFileToModuleFix(rsFile, moduleLabel)
+            } else {
+                null
+            }
+
+            val fixes = listOfNotNull(
+                attachFix,
+                SuppressFix()
+            )
+
+            return arrayOf(
+                manager.createProblemDescriptor(file,
+                    "File is not included in module tree, analysis is not available",
+                    isOnTheFly,
+                    fixes.toTypedArray(),
+                    ProblemHighlightType.WARNING
+                )
+            )
+        }
+
+        return null
+    }
+
+    private fun isInspectionEnabled(project: Project, file: VirtualFile): Boolean =
+        !PropertiesComponent.getInstance(project).getBoolean(file.disablingKey, false)
+
+    private class SuppressFix : SuppressQuickFix {
+        override fun getFamilyName(): String = "Do not show again"
+        override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+            val file = descriptor.startElement as? RsFile ?: return
+            PropertiesComponent.getInstance(project).setValue(file.virtualFile.disablingKey, true)
+        }
+
+        override fun isAvailable(project: Project, context: PsiElement): Boolean = true
+        override fun isSuppressAll(): Boolean = false
+    }
+
+    companion object {
+        private const val NOTIFICATION_STATUS_KEY = "org.rust.disableDetachedFileInspection"
+
+        private val VirtualFile.disablingKey: String
+            get() = NOTIFICATION_STATUS_KEY + path
+    }
+}

src/main/kotlin/org/rust/ide/inspections/fixes/AttachFileToModuleFix.kt

@@ -0,0 +1,193 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.inspections.fixes
+
+import com.intellij.codeInspection.LocalQuickFixOnPsiElement
+import com.intellij.notification.NotificationType
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.ui.ComboBox
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.openapiext.isUnitTestMode
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.ui.SimpleListCellRenderer
+import com.intellij.ui.components.dialog
+import com.intellij.ui.layout.CCFlags
+import com.intellij.ui.layout.panel
+import com.intellij.util.containers.addIfNotNull
+import org.jetbrains.annotations.TestOnly
+import org.rust.cargo.project.model.cargoProjects
+import org.rust.cargo.project.workspace.CargoWorkspace
+import org.rust.ide.notifications.showBalloon
+import org.rust.lang.RsConstants
+import org.rust.lang.core.psi.RsFile
+import org.rust.lang.core.psi.RsModDeclItem
+import org.rust.lang.core.psi.RsPsiFactory
+import org.rust.lang.core.psi.ext.RsMod
+import org.rust.lang.core.psi.ext.containingCargoPackage
+import org.rust.lang.core.psi.rustFile
+import org.rust.openapiext.pathAsPath
+import org.rust.openapiext.toPsiFile
+
+/**
+ * Attaches a file to a Rust module.
+ *
+ * Before fix:
+ * foo.rs (not attached)
+ * lib.rs
+ *
+ * After fix:
+ * foo.rs
+ * lib.rs
+ *  mod foo;
+ */
+class AttachFileToModuleFix(
+    file: RsFile,
+    private val targetModuleName: String? = null
+) : LocalQuickFixOnPsiElement(file) {
+    override fun getFamilyName(): String = text
+    override fun getText(): String = "Attach file to ${targetModuleName ?: "a module"}"
+
+    override fun invoke(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement) {
+        val rsFile = startElement as? RsFile ?: return
+        val availableModules = findAvailableModulesForFile(project, rsFile)
+        if (availableModules.isEmpty()) return
+
+        if (availableModules.size == 1) {
+            insertFileToModule(rsFile, availableModules[0])
+        } else if (availableModules.size > 1) {
+            selectModule(rsFile, availableModules)?.let { insertFileToModule(rsFile, it) }
+        }
+    }
+
+    override fun startInWriteAction(): Boolean {
+        return false
+    }
+
+    companion object {
+        fun findAvailableModulesForFile(project: Project, file: RsFile): List<RsMod> {
+            val virtualFile = file.virtualFile ?: return emptyList()
+            val pkg = project.cargoProjects.findPackageForFile(virtualFile) ?: return emptyList()
+
+            val directory = virtualFile.parent ?: return emptyList()
+            val modules = mutableListOf<RsMod>()
+
+            if (file.isModuleFile) {
+                // package target roots in parent directory
+                for (target in pkg.targets) {
+                    val crateRoot = target.crateRoot ?: continue
+                    if (crateRoot.parent == directory.parent) {
+                        modules.addIfNotNull(crateRoot.toPsiFile(project)?.rustFile)
+                    }
+                }
+            } else {
+                // mod.rs in the same directory
+                modules.addIfNotNull(findModule(file, project, directory.findFileByRelativePath(RsConstants.MOD_RS_FILE)))
+
+                // module file in parent directory
+                if (pkg.edition == CargoWorkspace.Edition.EDITION_2018) {
+                    val parent = directory.parent
+                    modules.addIfNotNull(findModule(file, project, parent?.findFileByRelativePath("${directory.name}.rs")))
+                }
+
+                // package target roots in the same directory
+                for (target in pkg.targets) {
+                    val crateRoot = target.crateRoot ?: continue
+                    if (crateRoot.parent == directory) {
+                        modules.addIfNotNull(crateRoot.toPsiFile(project)?.rustFile)
+                    }
+                }
+            }
+
+            return modules
+        }
+    }
+}
+
+private fun selectModule(file: RsFile, availableModules: List<RsMod>): RsMod? {
+    if (isUnitTestMode) {
+        val mock = MOCK ?: error("You should set mock module selector via withMockModuleAttachSelector")
+        return mock(file, availableModules)
+    }
+
+    val box = ComboBox<RsMod>()
+    with(box) {
+        for (module in availableModules) {
+            addItem(module)
+        }
+        renderer = SimpleListCellRenderer.create("") {
+            val root = it.containingCargoPackage?.rootDirectory
+            val path = it.containingFile.virtualFile.pathAsPath
+            (root?.relativize(path) ?: path).toString()
+        }
+    }
+
+    val dialog = dialog("Select a module", panel {
+        row { box(CCFlags.growX) }
+    }, focusedComponent = box)
+
+    return if (dialog.showAndGet()) {
+        box.selectedItem as? RsMod
+    } else {
+        null
+    }
+}
+
+private fun findModule(root: RsFile, project: Project, file: VirtualFile?): RsMod? {
+    if (file == null) return null
+    val module = file.toPsiFile(project)?.rustFile ?: return null
+    if (module == root || module.crateRoot == null) return null
+    return module
+}
+
+private fun insertFileToModule(file: RsFile, mod: RsMod) {
+    val project = file.project
+    val factory = RsPsiFactory(project)
+
+    // if the filename is mod.rs, attach it's parent directory
+    val name = if (file.isModuleFile) {
+        file.virtualFile.parent.name
+    } else {
+        file.virtualFile.nameWithoutExtension
+    }
+
+    val modItem = factory.tryCreateModDeclItem(name)
+    if (modItem == null) {
+        project.showBalloon("Could not create `mod ${name}`", NotificationType.ERROR)
+        return
+    }
+
+    WriteCommandAction.runWriteCommandAction(project) {
+        val child = mod.firstChild
+        val inserted = if (child == null) {
+            mod.add(modItem)
+        } else {
+            mod.addBefore(modItem, child)
+        } as RsModDeclItem
+        inserted.navigate(true)
+    }
+}
+
+private val RsFile.isModuleFile
+    get() = name == RsConstants.MOD_RS_FILE
+
+typealias ModuleAttachSelector = (file: RsFile, availableModules: List<RsMod>) -> RsMod?
+
+private var MOCK: ModuleAttachSelector? = null
+
+@TestOnly
+fun withMockModuleAttachSelector(
+    mock: ModuleAttachSelector,
+    f: () -> Unit
+) {
+    MOCK = mock
+    try {
+        f()
+    } finally {
+        MOCK = null
+    }
+}

src/main/kotlin/org/rust/lang/core/psi/RsPsiFactory.kt

@@ -205,8 +205,10 @@ class RsPsiFactory(
         createType("&${if (mutable) "mut " else ""}$innerTypeText").skipParens() as RsRefLikeType
 
     fun createModDeclItem(modName: String): RsModDeclItem =
+        tryCreateModDeclItem(modName) ?: error("Failed to create mod decl with name: `$modName`")
+
+    fun tryCreateModDeclItem(modName: String): RsModDeclItem? =
         createFromText("mod $modName;")
-            ?: error("Failed to create mod decl with name: `$modName`")
 
     fun createUseItem(text: String, visibility: String = ""): RsUseItem =
         createFromText("$visibility use $text;")

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

@@ -540,6 +540,11 @@
                          enabledByDefault="true" level="ERROR"
                          implementationClass="org.rust.ide.inspections.RsMainFunctionNotFoundInspection"/>
 
+        <localInspection language="Rust" groupName="Rust"
+                         displayName="Detached file"
+                         enabledByDefault="true" level="WARNING"
+                         implementationClass="org.rust.ide.inspections.RsDetachedFileInspection"/>
+
         <!-- Surrounders -->
 
         <lang.surroundDescriptor language="Rust"
@@ -856,7 +861,6 @@
         <!-- Notification Providers -->
 
         <editorNotificationProvider implementation="org.rust.ide.notifications.MissingToolchainNotificationProvider"/>
-        <editorNotificationProvider implementation="org.rust.ide.notifications.DetachedFileNotificationProvider"/>
         <editorNotificationProvider implementation="org.rust.ide.notifications.NoCargoProjectNotificationProvider"/>
 
         <!-- Editor Tab Title Providers -->

src/main/resources/inspectionDescriptions/RsDetachedFile.html

@@ -0,0 +1,5 @@
+<html>
+<body>
+Detects files that are not attached to any Rust module.
+</body>
+</html>