Initial implementation of new optimize imports algorithm

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

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.idea.util
 
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.JetFile
 import org.jetbrains.kotlin.psi.JetImportDirective
 import org.jetbrains.kotlin.resolve.ImportPath
@@ -24,6 +23,7 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.components.ServiceManager
 import kotlin.platform.platformStatic
+import java.util.*
 
 public abstract class ImportInsertHelper {
 
@@ -31,10 +31,10 @@ public abstract class ImportInsertHelper {
 
     public abstract fun isImportedWithDefault(importPath: ImportPath, contextFile: JetFile): Boolean
 
-    public abstract fun needImport(importPath: ImportPath, file: JetFile, importDirectives: List<JetImportDirective> = file.getImportDirectives()): Boolean
-
     public abstract fun mayImportByCodeStyle(descriptor: DeclarationDescriptor): Boolean
 
+    public abstract val importSortComparator: Comparator<ImportPath>
+
     public enum class ImportDescriptorResult {
         FAIL
         IMPORT_ADDED

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

@@ -60,8 +60,3 @@ public fun JetType.canBeReferencedViaImport(): Boolean {
     val descriptor = getConstructor().getDeclarationDescriptor()
     return descriptor != null && descriptor.canBeReferencedViaImport()
 }
-
-public fun isInReceiverScope(referenceElement: PsiElement, referencedDescriptor: DeclarationDescriptor): Boolean {
-    val isExpressionWithReceiver = referenceElement is JetSimpleNameExpression && referenceElement.getReceiverExpression() != null
-    return isExpressionWithReceiver && !referencedDescriptor.isExtension
-}

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

@@ -22,82 +22,137 @@ import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.resolve.ImportPath
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.idea.references.JetReference
-import java.util.HashSet
-import java.util.ArrayList
 import org.jetbrains.kotlin.idea.util.ImportInsertHelper
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import java.util.*
-import org.jetbrains.kotlin.idea.refactoring.fqName.isImported
+import org.jetbrains.kotlin.resolve.descriptorUtil.*
+import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.idea.formatter.JetCodeStyleSettings
+import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.psi.psiUtil.*
 
 public class KotlinImportOptimizer() : ImportOptimizer {
 
     override fun supports(file: PsiFile?) = file is JetFile
 
     override fun processFile(file: PsiFile?) = Runnable() {
-        val jetFile = file as JetFile
+        OptimizeProcess(file as JetFile).execute()
-        val usedQualifiedNames = extractUsedQualifiedNames(jetFile)
+    }
 
-        val directives = jetFile.getImportDirectives()
+    private class OptimizeProcess(val file: JetFile) {
+        private val codeStyleSettings = JetCodeStyleSettings.getInstance(file.getProject())
+        private val aliasImports: Map<Name, FqName>
 
-        val directivesBeforeCurrent = ArrayList<JetImportDirective>()
+        ;{
-        val directivesAfterCurrent = ArrayList(jetFile.getImportDirectives())
+            val imports = file.getImportDirectives()
-
+            val aliasImports = HashMap<Name, FqName>()
-        ApplicationManager.getApplication()!!.runWriteAction(Runnable {
+            for (import in imports) {
-            val importInsertHelper = ImportInsertHelper.getInstance(file.getProject())
+                val path = import.getImportPath() ?: continue
-            // Remove only unnecessary imports
+                val aliasName = path.getAlias()
-            for (anImport in directives) {
+                if (aliasName != null) {
-                directivesAfterCurrent.remove(anImport)
+                    aliasImports.put(aliasName, path.fqnPart())
-
-                val importPath = anImport.getImportPath()
-                if (importPath == null) {
-                    continue
-                }
-
-                if (isUseful(importPath, usedQualifiedNames)
-                    && importInsertHelper.needImport(importPath, jetFile, directivesBeforeCurrent)
-                    && importInsertHelper.needImport(importPath, jetFile, directivesAfterCurrent)
-                ) {
-                    directivesBeforeCurrent.add(anImport)
                 }
                 else {
-                    anImport.delete()
+                    //TODO: detect other "special" imports
                 }
             }
-        })
+            this.aliasImports = aliasImports
-    }
+        }
 
-    private fun isUseful(importPath: ImportPath, usedNames: Collection<FqName>): Boolean {
+        public fun execute() {
-        // TODO: Add better analysis for aliases
+            //TODO: keep existing imports? at least aliases (comments)
-        return importPath.hasAlias() || usedNames.any { it.isImported(importPath) }
-    }
 
-    private fun extractUsedQualifiedNames(jetFile: JetFile): Set<FqName> {
+            val importInsertHelper = ImportInsertHelper.getInstance(file.getProject())
-        val usedQualifiedNames = HashSet<FqName>()
+            val currentPackageName = file.getPackageFqName()
-        jetFile.accept(object : JetVisitorVoid() {
+
-            override fun visitElement(element: PsiElement) {
+            val descriptorsToImport = detectDescriptorsToImport()
-                ProgressIndicatorProvider.checkCanceled()
+
-                element?.acceptChildren(this)
+            val importsToGenerate = ArrayList<ImportPath>()
+            if (codeStyleSettings.PREFER_ALL_UNDER_IMPORTS) {
+                //TODO
+                //TODO: don't forget that packages may not be imported by *
+            }
+            else {
+                for (descriptor in descriptorsToImport) {
+                    val fqName = descriptor.importableFqNameSafe
+                    if (descriptor !is PackageViewDescriptor && fqName.parent() == currentPackageName) continue
+                    val importPath = ImportPath(fqName, false)
+                    if (importInsertHelper.isImportedWithDefault(importPath, file)/*TODO: implementation is not quite correct*/) continue
+                    importsToGenerate.add(importPath)
+                }
             }
 
-            override fun visitJetElement(element: JetElement) {
+            //TODO: drop unused aliases
-                if (element.getStrictParentOfType<JetImportDirective>() != null ||
+            aliasImports.mapTo(importsToGenerate) { ImportPath(it.getValue(), false, it.getKey())}
-                    element.getStrictParentOfType<JetPackageDirective>() != null) {
+
-                    return
+            val sortedImportsToGenerate = importsToGenerate.sortBy(importInsertHelper.importSortComparator)
+
+            //TODO: do not touch file if everything is already correct?
+
+            ApplicationManager.getApplication()!!.runWriteAction(Runnable {
+                for (import in file.getImportDirectives()) {
+                    import.delete()
                 }
-                val reference = element.getReference()
+
-                if (reference is JetReference) {
+                val importList = file.getImportList()!!
-                    val referencedDescriptors = reference.resolveToDescriptors()
+                val psiFactory = JetPsiFactory(file.getProject())
-                    val importableDescriptors = referencedDescriptors.filter {
+                for (importPath in sortedImportsToGenerate) {
-                        it.canBeReferencedViaImport() && !isInReceiverScope(element, it)
+                    importList.add(psiFactory.createImportDirective(importPath))
+                }
+            })
+        }
+
+        private fun detectDescriptorsToImport(): Set<DeclarationDescriptor> {
+            val usedDescriptors = HashSet<DeclarationDescriptor>()
+            file.accept(object : JetVisitorVoid() {
+                override fun visitElement(element: PsiElement) {
+                    ProgressIndicatorProvider.checkCanceled()
+                    element.acceptChildren(this)
+                }
+
+                override fun visitImportList(importList: JetImportList) {
+                }
+
+                override fun visitPackageDirective(directive: JetPackageDirective) {
+                }
+
+                override fun visitJetElement(element: JetElement) {
+                    val reference = element.getReference()
+                    if (reference is JetReference) {
+                        val targets = reference.resolveToDescriptors()
+                        //TODO: check if it's resolved through alias
+                        for (target in targets) {
+                            if (!target.canBeReferencedViaImport()) continue
+                            if (target is PackageViewDescriptor && target.getFqName().parent() == FqName.ROOT) continue // no need to import top-level packages (TODO: is that always true?)
+
+                            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 (isAccessibleAsMember(target, element)) continue
+
+                            usedDescriptors.add(target.getImportableDescriptor())
+                        }
                     }
-                    usedQualifiedNames.addAll(importableDescriptors.map { it.importableFqName }.filterNotNull())
-                }
-                super.visitJetElement(element)
-            }
-        })
 
-        return usedQualifiedNames
+                    super.visitJetElement(element)
+                }
+            })
+            return usedDescriptors
+        }
+
+        private fun isAccessibleAsMember(target: DeclarationDescriptor, place: JetElement): Boolean {
+            val container = target.getContainingDeclaration()
+            if (container !is ClassDescriptor) return false
+            val scope = if (container.getKind() == ClassKind.CLASS_OBJECT)
+                container.getContainingDeclaration() as? ClassDescriptor ?: return false
+            else
+                container
+            val declaration = DescriptorToSourceUtils.classDescriptorToDeclaration(scope)
+            return declaration != null && declaration.getContainingFile() == file && declaration.isAncestor(place)
+        }
     }
 }

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

@@ -102,6 +102,9 @@ public class ImportInsertHelperImpl(private val project: Project) : ImportInsert
         }
     }
 
+    override val importSortComparator: Comparator<ImportPath>
+        get() = ImportPathComparator
+
     private object ImportPathComparator : Comparator<ImportPath> {
         override fun compare(import1: ImportPath, import2: ImportPath): Int {
             // alias imports placed last
@@ -125,29 +128,6 @@ public class ImportInsertHelperImpl(private val project: Project) : ImportInsert
         }
     }
 
-    /**
-     * Check that import is useless.
-     */
-    private fun isImportedByDefault(importPath: ImportPath, jetFile: JetFile): Boolean {
-        if (importPath.fqnPart().isRoot()) {
-            return true
-        }
-
-        if (!importPath.isAllUnder() && !importPath.hasAlias()) {
-            // Single element import without .* and alias is useless
-            if (importPath.fqnPart().isOneSegmentFQN()) {
-                return true
-            }
-
-            // There's no need to import a declaration from the package of current file
-            if (jetFile.getPackageFqName() == importPath.fqnPart().parent()) {
-                return true
-            }
-        }
-
-        return isImportedWithDefault(importPath, jetFile)
-    }
-
     override fun isImportedWithDefault(importPath: ImportPath, contextFile: JetFile): Boolean {
         val defaultImports = if (ProjectStructureUtil.isJsKotlinModule(contextFile))
             TopDownAnalyzerFacadeForJS.DEFAULT_IMPORTS
@@ -156,24 +136,6 @@ public class ImportInsertHelperImpl(private val project: Project) : ImportInsert
         return importPath.isImported(defaultImports)
     }
 
-    override fun needImport(importPath: ImportPath, file: JetFile, importDirectives: List<JetImportDirective>): Boolean {
-        if (isImportedByDefault(importPath, file)) {
-            return false
-        }
-
-        if (!importDirectives.isEmpty()) {
-            // Check if import is already present
-            for (directive in importDirectives) {
-                val existentImportPath = directive.getImportPath()
-                if (existentImportPath != null && importPath.isImported(existentImportPath)) {
-                    return false
-                }
-            }
-        }
-
-        return true
-    }
-
     override fun mayImportByCodeStyle(descriptor: DeclarationDescriptor): Boolean {
         val importable = descriptor.getImportableDescriptor()
         return when (importable) {

idea/testData/editor/optimizeImports/ClassMemberImported/ClassMemberImported.after.kt

@@ -1,6 +1,7 @@
 package test
 
-import dependency.*
+import dependency.A
+import dependency.T
 
 fun f(a: A, t: T) {
     a.f()

idea/testData/editor/optimizeImports/CurrentPackage/CurrentPackage.after.kt

@@ -0,0 +1,8 @@
+package ppp
+
+import ppp.ppp1
+
+open class C
+
+class D : C()
+class E : ppp1.C1()

idea/testData/editor/optimizeImports/CurrentPackage/CurrentPackage.dependency.kt

@@ -0,0 +1,3 @@
+package ppp.ppp1
+
+class C1

idea/testData/editor/optimizeImports/CurrentPackage/CurrentPackage.kt

@@ -0,0 +1,9 @@
+package ppp
+
+import ppp.*
+import ppp.ppp1
+
+open class C
+
+class D : C()
+class E : ppp1.C1()

idea/testData/editor/optimizeImports/Kt2709/Kt2709.after.kt

@@ -1,4 +1,4 @@
-import MyClass.*
+import MyClass.TEST
 
 fun main() {
     TEST

idea/testData/editor/optimizeImports/MembersInScope/MembersInScope.after.kt

@@ -0,0 +1,18 @@
+import Outer.C.Inner2
+
+class Outer {
+    class Nested
+
+    class C {
+        fun foo(p1: Nested, p2: ClassObjectNested, p3: Nested2) { }
+
+        inner class Inner
+        inner class Inner2
+    }
+
+    class object {
+        class ClassObjectNested
+    }
+
+    fun f(i: Inner2){}
+}

idea/testData/editor/optimizeImports/MembersInScope/MembersInScope.kt

@@ -0,0 +1,21 @@
+import Outer.Nested
+import Outer.ClassObjectNested
+import Outer.C.Inner
+import Outer.C.Inner2
+
+class Outer {
+    class Nested
+
+    class C {
+        fun foo(p1: Nested, p2: ClassObjectNested, p3: Nested2) { }
+
+        inner class Inner
+        inner class Inner2
+    }
+
+    class object {
+        class ClassObjectNested
+    }
+
+    fun f(i: Inner2){}
+}

idea/testData/editor/optimizeImports/RemoveImportsIfGeneral/RemoveImportsIfGeneral.after.kt

@@ -1,4 +1,5 @@
-import java.util.*
+import java.util.ArrayList
+import java.util.HashSet
 
 class Action {
     fun test(hash: HashSet<Int>) {

idea/testData/editor/optimizeImports/RemoveImportsIfGeneralBefore/RemoveImportsIfGeneralBefore.after.kt

@@ -1,4 +1,5 @@
-import java.util.*
+import java.util.ArrayList
+import java.util.HashSet
 
 class Action {
     fun test(hash: HashSet<Int>) {

idea/testData/editor/optimizeImports/SamConstructor/SamConstructor.after.kt

@@ -1,3 +1,3 @@
-import java.util.*
+import java.util.Comparator
 
 val c = Comparator {(x: Int, y: Int) -> 1}

idea/testData/editor/optimizeImports/WithAliases/WithAliases.after.kt

@@ -1,5 +1,5 @@
-import java.io as JavaIO
 import java.util.ArrayList
+import java.io as JavaIO
 import java.util.ArrayList as SomeThing
 
 class Action {

idea/tests/org/jetbrains/kotlin/idea/imports/OptimizeImportsTestGenerated.java

@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.idea.imports;
 
 import com.intellij.testFramework.TestDataPath;
-import org.jetbrains.kotlin.test.InnerTestClasses;
 import org.jetbrains.kotlin.test.JUnit3RunnerWithInners;
 import org.jetbrains.kotlin.test.JetTestUtils;
 import org.jetbrains.kotlin.test.TestMetadata;
@@ -60,6 +59,12 @@ public class OptimizeImportsTestGenerated extends AbstractOptimizeImportsTest {
         doTest(fileName);
     }
 
+    @TestMetadata("CurrentPackage")
+    public void testCurrentPackage() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/CurrentPackage/");
+        doTest(fileName);
+    }
+
     @TestMetadata("DefaultJsImports")
     public void testDefaultJsImports() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/DefaultJsImports/");
@@ -126,6 +131,12 @@ public class OptimizeImportsTestGenerated extends AbstractOptimizeImportsTest {
         doTest(fileName);
     }
 
+    @TestMetadata("MembersInScope")
+    public void testMembersInScope() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/MembersInScope/");
+        doTest(fileName);
+    }
+
     @TestMetadata("PartiallyQualified")
     public void testPartiallyQualified() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/PartiallyQualified/");