On-the-fly import optimization

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/ImportInsertHelper.kt

@@ -26,9 +26,6 @@ import kotlin.platform.platformStatic
 import java.util.*
 
 public abstract class ImportInsertHelper {
-
-    public abstract fun optimizeImportsOnTheFly(file: JetFile): Boolean
-
     /*TODO: implementation is not quite correct*/
     public abstract fun isImportedWithDefault(importPath: ImportPath, contextFile: JetFile): Boolean
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/ShortenReferences.kt

@@ -141,7 +141,7 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
         //TODO: that's not correct since we have options!
         val elementsToUse = dropNestedElements(elements)
 
-        val importInserter = ImportInserter(file)
+        val helper = ImportInsertHelper.getInstance(file.getProject())
 
         val failedToImportDescriptors = LinkedHashSet<DeclarationDescriptor>()
 
@@ -161,11 +161,7 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
             for (descriptor in descriptorsToImport) {
                 assert(descriptor !in failedToImportDescriptors)
 
-                if (importInserter.optimizeImports()) {
-                    anyChange = true
-                }
-
-                val result = importInserter.addImport(descriptor)
+                val result = helper.importDescriptor(file, descriptor)
                 if (result != ImportInsertHelper.ImportDescriptorResult.ALREADY_IMPORTED) {
                     anyChange = true
                 }
@@ -457,20 +453,4 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
             return element.replace(element.getReceiverExpression()) as JetElement
         }
     }
-
-    // this class is needed to optimize imports only when we actually insert any import (optimization)
-    private class ImportInserter(val file: JetFile) {
-        private var optimizeImports = true
-        private val helper = ImportInsertHelper.getInstance(file.getProject())
-
-        fun addImport(target: DeclarationDescriptor): ImportInsertHelper.ImportDescriptorResult {
-            return helper.importDescriptor(file, target)
-        }
-
-        fun optimizeImports(): Boolean {
-            if (!optimizeImports) return false
-            optimizeImports = false
-            return helper.optimizeImportsOnTheFly(file)
-        }
-    }
 }

idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt

@@ -18,7 +18,6 @@ package org.jetbrains.kotlin.idea.imports
 
 import com.google.common.collect.HashMultimap
 import com.intellij.lang.ImportOptimizer
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
@@ -29,11 +28,11 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.core.formatter.JetCodeStyleSettings
 import org.jetbrains.kotlin.idea.references.JetReference
-import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.idea.util.ImportInsertHelper
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
@@ -56,32 +55,99 @@ public class KotlinImportOptimizer() : ImportOptimizer {
         OptimizeProcess(file as JetFile).execute()
     }
 
-    private class OptimizeProcess(val file: JetFile) {
-        private val codeStyleSettings = JetCodeStyleSettings.getInstance(file.getProject())
-        private val aliasImports: Map<Name, FqName>
-
-        init {
-            val imports = file.getImportDirectives()
-            val aliasImports = HashMap<Name, FqName>()
-            for (import in imports) {
-                val path = import.getImportPath() ?: continue
-                val aliasName = path.getAlias()
-                if (aliasName != null) {
-                    aliasImports.put(aliasName, path.fqnPart())
-                }
-            }
-            this.aliasImports = aliasImports
-        }
-
+    private class OptimizeProcess(private val file: JetFile) {
         public fun execute() {
-            val oldImports = file.getImportDirectives()
+            val oldImports = file.importDirectives
             if (oldImports.isEmpty()) return
 
             //TODO: keep existing imports? at least aliases (comments)
 
-            val importInsertHelper = ImportInsertHelper.getInstance(file.getProject())
+            val visitor = CollectUsedDescriptorsVisitor(file, recursive = true)
+            file.accept(visitor)
+            val descriptorsToImport = visitor.descriptors
 
-            val descriptorsToImport = detectDescriptorsToImport()
+            val imports = prepareOptimizedImports(file, descriptorsToImport) ?: return
+
+            runWriteAction { replaceImports(file, imports) }
+        }
+    }
+
+    public open class CollectUsedDescriptorsVisitor(val file: JetFile, val recursive: Boolean) : JetVisitorVoid() {
+        private val _descriptors = HashSet<DeclarationDescriptor>()
+        private val currentPackageName = file.packageFqName
+
+        public val descriptors: Set<DeclarationDescriptor>
+            get() = _descriptors
+
+        override fun visitElement(element: PsiElement) {
+            if (recursive) {
+                ProgressIndicatorProvider.checkCanceled()
+                element.acceptChildren(this)
+            }
+        }
+
+        override fun visitImportList(importList: JetImportList) {
+        }
+
+        override fun visitPackageDirective(directive: JetPackageDirective) {
+        }
+
+        override fun visitJetElement(element: JetElement) {
+            if (!recursive && element.parents.any { it is JetImportDirective || it is JetPackageDirective }) return
+
+            for (reference in element.references) {
+                if (reference !is JetReference) continue
+
+                val referencedName = (element as? JetNameReferenceExpression)?.getReferencedNameAsName() //TODO: other types of references
+
+                val bindingContext = element.analyze()
+                //class qualifiers that refer to companion objects should be considered (containing) class references
+                val targets = bindingContext[BindingContext.SHORT_REFERENCE_TO_COMPANION_OBJECT, element as? JetReferenceExpression]?.let { listOf(it) }
+                              ?: reference.resolveToDescriptors(bindingContext)
+                for (target in targets) {
+                    if (!target.canBeReferencedViaImport()) continue
+                    val importableDescriptor = target.getImportableDescriptor()
+                    val parentFqName = DescriptorUtils.getFqNameSafe(importableDescriptor).parent()
+                    if (target is PackageViewDescriptor && parentFqName == FqName.ROOT) continue // no need to import top-level packages
+                    if (target !is PackageViewDescriptor && parentFqName == currentPackageName) continue
+
+                    if (!target.isExtension) { // for non-extension targets, count only non-qualified simple name usages
+                        if (element !is JetNameReferenceExpression) continue
+                        if (element.getIdentifier() == null) continue // skip 'this' etc
+                        if (element.getReceiverExpression() != null) continue
+                    }
+
+                    if (referencedName != null && importableDescriptor.name != referencedName) continue // resolved via alias
+
+                    if (isAccessibleAsMember(importableDescriptor, element)) continue
+
+                    _descriptors.add(importableDescriptor)
+                }
+            }
+
+            super.visitJetElement(element)
+        }
+
+        private fun isAccessibleAsMember(target: DeclarationDescriptor, place: JetElement): Boolean {
+            val container = target.containingDeclaration
+            if (container !is ClassDescriptor) return false
+            val scope = if (DescriptorUtils.isCompanionObject(container))
+                container.getContainingDeclaration() as? ClassDescriptor ?: return false
+            else
+                container
+            val classBody = (DescriptorToSourceUtils.getSourceFromDescriptor(scope) as? JetClassOrObject)?.getBody()
+            return classBody != null && classBody.containingFile == file && classBody.isAncestor(place)
+        }
+    }
+
+    companion object {
+        public fun prepareOptimizedImports(
+                file: JetFile,
+                descriptorsToImport: Collection<DeclarationDescriptor>
+        ): List<ImportPath>? {
+            val importInsertHelper = ImportInsertHelper.getInstance(file.project)
+            val codeStyleSettings = JetCodeStyleSettings.getInstance(file.project)
+            val aliasImports = buildAliasImportMap(file)
 
             val importsToGenerate = HashSet<ImportPath>()
 
@@ -127,8 +193,8 @@ public class KotlinImportOptimizer() : ImportOptimizer {
 
             // now check that there are no conflicts and all classes are really imported
             val fileWithImportsText = StringBuilder {
-                append("package ").append(file.getPackageFqName().render()).append("\n")
-                importsToGenerate.filter { it.isAllUnder() }.map { "import " + it.getPathStr() }.joinTo(this, "\n")
+                append("package ").append(file.packageFqName.render()).append("\n")
+                importsToGenerate.filter { it.isAllUnder() }.map { "import " + it.pathStr }.joinTo(this, "\n")
             }.toString()
             val fileWithImports = JetPsiFactory(file).createAnalyzableFile("Dummy.kt", fileWithImportsText, file)
             val scope = fileWithImports.getResolutionFacade().getFileTopLevelScope(fileWithImports)
@@ -156,94 +222,37 @@ public class KotlinImportOptimizer() : ImportOptimizer {
             val sortedImportsToGenerate = importsToGenerate.sortBy(importInsertHelper.importSortComparator)
 
             // check if no changes to imports required
-            if (oldImports.size() == sortedImportsToGenerate.size() && oldImports.map { it.getImportPath() } == sortedImportsToGenerate) return
+            val oldImports = file.importDirectives
+            if (oldImports.size() == sortedImportsToGenerate.size() && oldImports.map { it.importPath } == sortedImportsToGenerate) return null
 
-            runWriteAction {
-                val importList = file.getImportList()!!
-                val psiFactory = JetPsiFactory(file.getProject())
-                for (importPath in sortedImportsToGenerate) {
-                    importList.addBefore(psiFactory.createImportDirective(importPath), oldImports.lastOrNull()) // insert into the middle to keep collapsed state
-                }
+            return sortedImportsToGenerate
+        }
 
-                // remove old imports after adding new ones to keep imports folding state
-                for (import in oldImports) {
-                    import.delete()
+        private fun buildAliasImportMap(file: JetFile): Map<Name, FqName> {
+            val imports = file.importDirectives
+            val aliasImports = HashMap<Name, FqName>()
+            for (import in imports) {
+                val path = import.importPath ?: continue
+                val aliasName = path.alias
+                if (aliasName != null) {
+                    aliasImports.put(aliasName, path.fqnPart())
                 }
             }
+            return aliasImports
         }
 
-        private fun detectDescriptorsToImport(): Set<DeclarationDescriptor> {
-            val visitor = CollectUsedDescriptorsVisitor(file, recursive = true)
-            file.accept(visitor)
-            return visitor.descriptors
-        }
-    }
-
-    public class CollectUsedDescriptorsVisitor(val file: JetFile, val recursive: Boolean) : JetVisitorVoid() {
-        private val _descriptors = HashSet<DeclarationDescriptor>()
-        private val currentPackageName = file.getPackageFqName()
-
-        public val descriptors: Set<DeclarationDescriptor>
-            get() = _descriptors
-
-        override fun visitElement(element: PsiElement) {
-            if (recursive) {
-                ProgressIndicatorProvider.checkCanceled()
-                element.acceptChildren(this)
-            }
-        }
-
-        override fun visitImportList(importList: JetImportList) {
-        }
-
-        override fun visitPackageDirective(directive: JetPackageDirective) {
-        }
-
-        override fun visitJetElement(element: JetElement) {
-            if (!recursive && element.parents.any { it is JetImportDirective || it is JetPackageDirective }) return
-
-            for (reference in element.getReferences()) {
-                if (reference !is JetReference) continue
-
-                val referencedName = (element as? JetNameReferenceExpression)?.getReferencedNameAsName() //TODO: other types of references
-
-                val bindingContext = element.analyze()
-                //class qualifiers that refer to companion objects should be considered (containing) class references
-                val targets = bindingContext[BindingContext.SHORT_REFERENCE_TO_COMPANION_OBJECT, element as? JetReferenceExpression]?.let { listOf(it) }
-                              ?: reference.resolveToDescriptors(bindingContext)
-                for (target in targets) {
-                    if (!target.canBeReferencedViaImport()) continue
-                    val importableDescriptor = target.getImportableDescriptor()
-                    val parentFqName = DescriptorUtils.getFqNameSafe(importableDescriptor).parent()
-                    if (target is PackageViewDescriptor && parentFqName == FqName.ROOT) continue // no need to import top-level packages
-                    if (target !is PackageViewDescriptor && parentFqName == currentPackageName) continue
-
-                    if (!target.isExtension) { // for non-extension targets, count only non-qualified simple name usages
-                        if (element !is JetNameReferenceExpression) continue
-                        if (element.getIdentifier() == null) continue // skip 'this' etc
-                        if (element.getReceiverExpression() != null) continue
-                    }
-
-                    if (referencedName != null && importableDescriptor.getName() != referencedName) continue // resolved via alias
-
-                    if (isAccessibleAsMember(importableDescriptor, element)) continue
-
-                    _descriptors.add(importableDescriptor)
-                }
+        public fun replaceImports(file: JetFile, imports: List<ImportPath>) {
+            val importList = file.importList!!
+            val oldImports = importList.imports
+            val psiFactory = JetPsiFactory(file.project)
+            for (importPath in imports) {
+                importList.addBefore(psiFactory.createImportDirective(importPath), oldImports.lastOrNull()) // insert into the middle to keep collapsed state
             }
 
-            super.visitJetElement(element)
-        }
-
-        private fun isAccessibleAsMember(target: DeclarationDescriptor, place: JetElement): Boolean {
-            val container = target.getContainingDeclaration()
-            if (container !is ClassDescriptor) return false
-            val scope = if (DescriptorUtils.isCompanionObject(container))
-                container.getContainingDeclaration() as? ClassDescriptor ?: return false
-            else
-                container
-            val classBody = (DescriptorToSourceUtils.getSourceFromDescriptor(scope) as? JetClassOrObject)?.getBody()
-            return classBody != null && classBody.getContainingFile() == file && classBody.isAncestor(place)
+            // remove old imports after adding new ones to keep imports folding state
+            for (import in oldImports) {
+                import.delete()
+            }
         }
     }
 }

idea/src/org/jetbrains/kotlin/idea/inspections/UnusedImportInspection.kt

@@ -16,11 +16,26 @@
 
 package org.jetbrains.kotlin.idea.inspections
 
+import com.intellij.codeInsight.CodeInsightSettings
 import com.intellij.codeInsight.actions.OptimizeImportsProcessor
+import com.intellij.codeInsight.daemon.impl.DaemonCodeAnalyzerEx
+import com.intellij.codeInsight.daemon.impl.DaemonListeners
+import com.intellij.codeInsight.daemon.impl.HighlightingSessionImpl
 import com.intellij.codeInspection.*
+import com.intellij.lang.annotation.HighlightSeverity
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.command.undo.UndoManager
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.progress.util.ProgressWrapper
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.Key
+import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElementVisitor
+import com.intellij.util.DocumentUtil
+import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.imports.KotlinImportOptimizer
 import org.jetbrains.kotlin.idea.imports.importableFqNameSafe
@@ -29,26 +44,37 @@ import org.jetbrains.kotlin.psi.JetFile
 import org.jetbrains.kotlin.psi.JetImportDirective
 import org.jetbrains.kotlin.psi.JetSimpleNameExpression
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
+import org.jetbrains.kotlin.resolve.ImportPath
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getReferenceTargets
 import java.util.HashSet
 
 class UnusedImportInspection : AbstractKotlinInspection() {
-    private val visitorKey = Key<KotlinImportOptimizer.CollectUsedDescriptorsVisitor>("UnusedImportInspection.visitorKey")
+    private val visitorKey = Key<MyVisitor>("UnusedImportInspection.visitorKey")
 
     override fun runForWholeFile() = true
 
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean, session: LocalInspectionToolSession): PsiElementVisitor {
-        val visitor = KotlinImportOptimizer.CollectUsedDescriptorsVisitor(session.getFile() as JetFile, recursive = false)
+        val visitor = MyVisitor(session.file as JetFile)
         session.putUserData(visitorKey, visitor)
         return visitor
     }
 
+    private class MyVisitor(file: JetFile) : KotlinImportOptimizer.CollectUsedDescriptorsVisitor(file, recursive = false) {
+        var anyImportVisited = false
+
+        override fun visitImportDirective(importDirective: JetImportDirective) {
+            anyImportVisited = true
+        }
+    }
+
     override fun inspectionFinished(session: LocalInspectionToolSession, problemsHolder: ProblemsHolder) {
-        val usedDescriptors = session.getUserData(visitorKey)!!.descriptors
+        val visitor = session.getUserData(visitorKey)!!
+        if (!visitor.anyImportVisited) return // workaround for the problem (bug?) that IDEA does not call us for all elements when we press Enter
+        val descriptorsToImport = visitor.descriptors
 
         val fqNames = HashSet<FqName>()
         val parentFqNames = HashSet<FqName>()
-        for (descriptor in usedDescriptors) {
+        for (descriptor in descriptorsToImport) {
             val fqName = descriptor.importableFqNameSafe
             fqNames.add(fqName)
             val parentFqName = fqName.parent()
@@ -57,11 +83,11 @@ class UnusedImportInspection : AbstractKotlinInspection() {
             }
         }
 
-        val file = session.getFile() as JetFile
-        val directives = file.getImportDirectives()
+        val file = session.file as JetFile
+        val directives = file.importDirectives
         for (directive in directives) {
-            val importPath = directive.getImportPath() ?: continue
-            if (importPath.getAlias() != null) continue // highlighting of unused alias imports not supported yet
+            val importPath = directive.importPath ?: continue
+            if (importPath.alias != null) continue // highlighting of unused alias imports not supported yet
             val isUsed = if (importPath.isAllUnder()) {
                 importPath.fqnPart() in parentFqNames
             }
@@ -70,7 +96,7 @@ class UnusedImportInspection : AbstractKotlinInspection() {
             }
 
             if (!isUsed) {
-                val nameExpression = directive.getImportedReference()?.getQualifiedElementSelector() as? JetSimpleNameExpression
+                val nameExpression = directive.importedReference?.getQualifiedElementSelector() as? JetSimpleNameExpression
                 if (nameExpression == null || nameExpression.getReferenceTargets(nameExpression.analyze()).isEmpty()) continue // do not highlight unresolved imports as unused
 
                 problemsHolder.registerProblem(
@@ -80,12 +106,70 @@ class UnusedImportInspection : AbstractKotlinInspection() {
                         OptimizeImportsQuickFix(file))
             }
         }
+
+        if (CodeInsightSettings.getInstance().OPTIMIZE_IMPORTS_ON_THE_FLY) {
+            val optimizedImports = KotlinImportOptimizer.prepareOptimizedImports(file, descriptorsToImport) ?: return // return if already optimized
+
+            // unwrap progress indicator
+            val progress = sequence(ProgressManager.getInstance().progressIndicator) {
+                (it as? ProgressWrapper)?.originalProgressIndicator
+            }.last()
+            val highlightingSession = HighlightingSessionImpl.getHighlightingSession(file, progress)
+
+            val project = highlightingSession.project
+            val editor = highlightingSession.editor
+            if (editor != null) {
+                Disposer.register(highlightingSession, object : Disposable {
+                    override fun dispose() {
+                        // later because should invoke when highlighting is finished
+                        ApplicationManager.getApplication().invokeLater {
+                            if (timeToOptimizeImportsOnTheFly(file, editor, project)) {
+                                optimizeImportsOnTheFly(file, optimizedImports, editor, project)
+                            }
+                        }
+                    }
+                })
+            }
+        }
+    }
+
+    private fun timeToOptimizeImportsOnTheFly(file: JetFile, editor: Editor, project: Project): Boolean {
+        if (project.isDisposed() || !file.isValid() || editor.isDisposed()) return false
+
+        if (!file.isWritable()) return false
+
+        // do not optimize imports on the fly during undo/redo
+        val undoManager = UndoManager.getInstance(editor.project)
+        if (undoManager.isUndoInProgress() || undoManager.isRedoInProgress()) return false
+
+        // if we stand inside import statements, do not optimize
+        val importsRange = file.importList?.textRange ?: return false
+        if (importsRange.containsOffset(editor.caretModel.offset)) return false
+
+        val codeAnalyzer = DaemonCodeAnalyzerEx.getInstanceEx(project)
+        if (!codeAnalyzer.isHighlightingAvailable(file)) return false
+        if (!codeAnalyzer.isErrorAnalyzingFinished(file)) return false
+
+        val document = editor.document
+        var hasErrors = false
+        DaemonCodeAnalyzerEx.processHighlights(document, project, HighlightSeverity.ERROR, 0, document.textLength, Processor { hasErrors = true; false })
+        if (hasErrors) return false
+
+        return DaemonListeners.canChangeFileSilently(file)
+    }
+
+    private fun optimizeImportsOnTheFly(file: JetFile, optimizedImports: List<ImportPath>, editor: Editor, project: Project) {
+        PsiDocumentManager.getInstance(file.project).commitAllDocuments()
+        DocumentUtil.writeInRunUndoTransparentAction {
+            KotlinImportOptimizer.replaceImports(file, optimizedImports)
+            PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(editor.document)
+        }
     }
 
     private class OptimizeImportsQuickFix(private val file: JetFile): LocalQuickFix {
         override fun getName() = "Optimize imports"
 
-        override fun getFamilyName() = getName()
+        override fun getFamilyName() = name
 
         override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
             OptimizeImportsProcessor(project, file).run()

idea/src/org/jetbrains/kotlin/idea/util/ImportInsertHelperImpl.kt

@@ -55,16 +55,6 @@ public class ImportInsertHelperImpl(private val project: Project) : ImportInsert
     private val codeStyleSettings: JetCodeStyleSettings
         get() = JetCodeStyleSettings.getInstance(project)
 
-    override fun optimizeImportsOnTheFly(file: JetFile): Boolean {
-        if (CodeInsightSettings.getInstance().OPTIMIZE_IMPORTS_ON_THE_FLY) {
-            OptimizeImportsProcessor(project, file).runWithoutProgress()
-            return true
-        }
-        else {
-            return false
-        }
-    }
-
     override val importSortComparator: Comparator<ImportPath>
         get() = ImportPathComparator