Reference shortening: performance optimization

idea/src/org/jetbrains/jet/plugin/codeInsight/ShortenReferences.kt

@@ -14,6 +14,8 @@ import org.jetbrains.jet.renderer.DescriptorRenderer;
 import java.util.Collections;
 import java.util.HashSet;
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.containers.HashMap
+import java.util.ArrayList
 
 public object ShortenReferences {
     public fun process(element: JetElement) {
@@ -23,59 +25,131 @@ public object ShortenReferences {
     public fun process(elements: Iterable<JetElement>) {
         val first = elements.firstOrNull()
         if (first == null) return
+        val file = first.getContainingFile() as JetFile
 
-        val visitor = Visitor(first.getContainingFile() as JetFile)
+        // first resolve all qualified references - optimization
+        val resolveAllVisitor = ResolveAllReferencesVisitor(file)
+        processElements(elements, resolveAllVisitor)
+
+        val shortenTypesVisitor = ShortenTypesVisitor(file, resolveAllVisitor.result)
+        processElements(elements, shortenTypesVisitor)
+        shortenTypesVisitor.finish()
+
+        processElements(elements, ShortenQualifiedExpressionsVisitor(file, resolveAllVisitor.result))
+    }
+
+    private fun processElements(elements: Iterable<JetElement>, visitor: JetVisitorVoid) {
         for (element in elements) {
             element.accept(visitor)
         }
     }
 
-    private class Visitor(val file: JetFile) : JetVisitorVoid() {
+    private class ResolveAllReferencesVisitor(file: JetFile) : JetVisitorVoid() {
-        private val resolveSession : CancelableResolveSession
+        private val resolveSession = AnalyzerFacadeWithCache.getLazyResolveSessionForFile(file)
-            get() = AnalyzerFacadeWithCache.getLazyResolveSessionForFile(file)
+        private val resolveMap = HashMap<JetReferenceExpression, BindingContext>()
+
+        public val result: Map<JetReferenceExpression, BindingContext> = resolveMap
 
         override fun visitJetElement(element : JetElement) {
-            acceptChildren(element)
+            element.acceptChildren(this)
         }
 
         override fun visitUserType(userType: JetUserType) {
-            val resultElement = processType(userType)
+            userType.acceptChildren(this)
-            acceptChildren(resultElement)
+
+            if (userType.getQualifier() != null) {
+                val referenceExpression = userType.getReferenceExpression()
+                if (referenceExpression != null) {
+                    resolveMap[referenceExpression] = resolveSession.resolveToElement(referenceExpression)
+                }
+            }
         }
 
-        private fun processType(userType: JetUserType): PsiElement {
+        override fun visitDotQualifiedExpression(expression: JetDotQualifiedExpression) {
-            if (userType.getQualifier() == null) return userType
+            expression.acceptChildren(this)
 
-            val bindingContext = resolveSession.resolveToElement(userType)
+            val referenceExpression = referenceExpression(expression.getSelectorExpression())
-            val target = bindingContext.get(BindingContext.REFERENCE_TARGET, userType.getReferenceExpression())
+            if (referenceExpression != null) {
-            if (target == null) return userType
+                resolveMap[referenceExpression] = resolveSession.resolveToElement(referenceExpression)
+            }
+        }
+    }
+
+    private class ShortenTypesVisitor(val file: JetFile, val resolveMap: Map<JetReferenceExpression, BindingContext>) : JetVisitorVoid() {
+        private val resolveSession : CancelableResolveSession
+            get() = AnalyzerFacadeWithCache.getLazyResolveSessionForFile(file)
+
+        private val typesToShorten = ArrayList<JetUserType>()
+
+        public fun finish() {
+            for (userType in typesToShorten) {
+                shortenType(userType)
+            }
+        }
+
+        private fun bindingContext(expression: JetReferenceExpression): BindingContext = resolveMap[expression]!!
+
+        override fun visitJetElement(element : JetElement) {
+            element.acceptChildren(this)
+        }
+
+        override fun visitUserType(userType: JetUserType) {
+            userType.getTypeArgumentList()?.accept(this)
+
+            if (canShortenType(userType)) {
+                typesToShorten.add(userType)
+            }
+            else{
+                userType.getQualifier()?.accept(this)
+            }
+        }
+
+        private fun canShortenType(userType: JetUserType): Boolean {
+            if (userType.getQualifier() == null) return false
+            val referenceExpression = userType.getReferenceExpression()
+            if (referenceExpression == null) return false
+
+            val target = bindingContext(referenceExpression).get(BindingContext.REFERENCE_TARGET, referenceExpression)
+            if (target == null) return false
             // references to nested classes should be shortened when visiting qualifier
-            if (target.getContainingDeclaration() is ClassDescriptor) return userType
+            if (target.getContainingDeclaration() is ClassDescriptor) return false
 
             val typeReference = PsiTreeUtil.getParentOfType(userType, javaClass<JetTypeReference>())!!
-            val scope = bindingContext.get(BindingContext.TYPE_RESOLUTION_SCOPE, typeReference)!!
+            val scope = resolveSession.resolveToElement(typeReference).get(BindingContext.TYPE_RESOLUTION_SCOPE, typeReference)!!
             val name = target.getName()
             val targetByName = scope.getClassifier(name) ?: scope.getPackage(name)
-            if (target == targetByName) {
+            if (targetByName == null) {
-                return shortenType(userType)
+                addImportIfNeeded(target, file)
+                return true
             }
-            else if (targetByName == null) {
+            else if (target.asString() == targetByName.asString()) {
-                addImportIfNeeded(target)
+                return true
-                return shortenType(userType)
             }
             else {
                 // leave FQ name
-                return userType
+                return false
             }
         }
 
-        private fun shortenType(userType: JetUserType): JetUserType {
+        private fun shortenType(userType: JetUserType) {
             val referenceExpression = userType.getReferenceExpression()
-            if (referenceExpression == null) return userType
+            if (referenceExpression == null) return
             val typeArgumentList = userType.getTypeArgumentList()
             val text = referenceExpression.getText() + (if (typeArgumentList != null) typeArgumentList.getText() else "")
-            val newUserType = JetPsiFactory.createType(userType.getProject(), text).getTypeElement() as JetUserType
+            val newUserType = JetPsiFactory.createType(userType.getProject(), text).getTypeElement()!!
-            return userType.replace(newUserType) as JetUserType
+            userType.replace(newUserType)
+        }
+    }
+
+    private class ShortenQualifiedExpressionsVisitor(val file: JetFile, val resolveMap: Map<JetReferenceExpression, BindingContext>) : JetVisitorVoid() {
+        private val resolveSession : CancelableResolveSession
+            get() = AnalyzerFacadeWithCache.getLazyResolveSessionForFile(file)
+
+        private fun bindingContext(expression: JetReferenceExpression): BindingContext
+                = resolveMap[expression] ?: resolveSession.resolveToElement(expression) // binding context can be absent in the map if some references have been shortened already
+
+        override fun visitJetElement(element : JetElement) {
+            acceptChildren(element)
         }
 
         override fun visitDotQualifiedExpression(expression: JetDotQualifiedExpression) {
@@ -90,22 +164,22 @@ public object ShortenReferences {
                 if (calleeExpression is JetReferenceExpression) {
                     val targetClass = instantiatedClass(calleeExpression)
                     if (targetClass != null) {
-                        return shortenIfPossible(qualifiedExpression, targetClass)
+                        return shortenIfPossible(qualifiedExpression, targetClass, bindingContext(calleeExpression))
                     }
                 }
             }
             else if (selectorExpression is JetReferenceExpression) {
-                val bindingContext = resolveSession.resolveToElement(selectorExpression)
+                val bindingContext = bindingContext(selectorExpression)
                 val target = bindingContext.get(BindingContext.REFERENCE_TARGET, selectorExpression)
                 if (target is ClassDescriptor || target is PackageViewDescriptor) { //TODO: should we ever add imports to real packages?
-                    return shortenIfPossible(qualifiedExpression, target)
+                    return shortenIfPossible(qualifiedExpression, target, bindingContext)
                 }
             }
             return qualifiedExpression
         }
 
         private fun instantiatedClass(calleeExpression: JetReferenceExpression): ClassDescriptor? {
-            val bindingContext = resolveSession.resolveToElement(calleeExpression)
+            val bindingContext = bindingContext(calleeExpression)
             val target = bindingContext.get(BindingContext.REFERENCE_TARGET, calleeExpression)
             if (target != null) {
                 if (target is ConstructorDescriptor) {
@@ -132,47 +206,40 @@ public object ShortenReferences {
             return null
         }
 
-        private fun shortenIfPossible(qualifiedExpression: JetDotQualifiedExpression, targetClassOrPackage: DeclarationDescriptor): PsiElement {
+        private fun shortenIfPossible(qualifiedExpression: JetDotQualifiedExpression, targetClassOrPackage: DeclarationDescriptor, bindingContext: BindingContext): PsiElement {
             // references to nested classes should be shortened when visiting qualifier
             if (targetClassOrPackage.getContainingDeclaration() is ClassDescriptor) return qualifiedExpression
 
-            var bindingContext = resolveSession.resolveToElement(qualifiedExpression)
+            val referenceExpression = referenceExpression(qualifiedExpression.getSelectorExpression())!!
-
-            val referenceExpression = referenceExpression(qualifiedExpression.getSelectorExpression()!!)
             val resolveBefore = resolveState(referenceExpression, bindingContext)
 
             val copy = qualifiedExpression.copy()
 
             val selectorExpression = qualifiedExpression.getSelectorExpression()!!
             val newExpression = qualifiedExpression.replace(selectorExpression) as JetExpression
-            val newReferenceExpression = referenceExpression(newExpression)
+            val newReferenceExpression = referenceExpression(newExpression)!!
 
-            bindingContext = resolveSession.resolveToElement(newReferenceExpression)
+            val newBindingContext = resolveSession.resolveToElement(newReferenceExpression)
-            val resolveAfter = resolveState(newReferenceExpression, bindingContext)
+            val resolveAfter = resolveState(newReferenceExpression, newBindingContext)
             if (resolveAfter != null) {
                 if (resolveBefore == resolveAfter) return newExpression
                 return newExpression.replace(copy) // revert shortening
             }
 
-            addImportIfNeeded(targetClassOrPackage)
+            addImportIfNeeded(targetClassOrPackage, file)
             return newExpression
         }
 
         private fun resolveState(referenceExpression: JetReferenceExpression, bindingContext: BindingContext): Any? {
             val target = bindingContext.get(BindingContext.REFERENCE_TARGET, referenceExpression)
-            if (target != null) return DescriptorRenderer.TEXT.render(target.getOriginal())
+            if (target != null) return target.asString()
 
             val targets = bindingContext.get(BindingContext.AMBIGUOUS_REFERENCE_TARGET, referenceExpression)
-            if (targets != null) return HashSet(targets.map{DescriptorRenderer.TEXT.render(it!!)})
+            if (targets != null) return HashSet(targets.map{it!!.asString()})
 
             return null
         }
 
-        //TODO: do we need this "IfNeeded" check?
-        private fun addImportIfNeeded(descriptor : DeclarationDescriptor) {
-            ImportInsertHelper.addImportDirectiveIfNeeded(DescriptorUtils.getFqNameSafe(descriptor), file)
-        }
-
         // we do not use standard PsiElement.acceptChildren because it won't work correctly if the element is replaced by the visitor
         private fun acceptChildren(element: PsiElement) {
             var child = element.getFirstChild()
@@ -184,10 +251,18 @@ public object ShortenReferences {
         }
     }
 
-    private fun referenceExpression(selectorExpression: JetExpression) = if (selectorExpression is JetCallExpression)
+    private fun referenceExpression(selectorExpression: JetExpression?): JetReferenceExpression?
-            selectorExpression.getCalleeExpression() as JetReferenceExpression
+            = if (selectorExpression is JetCallExpression)
-        else
+                  selectorExpression.getCalleeExpression() as? JetReferenceExpression
-            selectorExpression as JetReferenceExpression
+              else
+                  selectorExpression as? JetReferenceExpression
+
+    private fun DeclarationDescriptor.asString() = DescriptorRenderer.TEXT.render(this)
+
+    //TODO: do we need this "IfNeeded" check?
+    private fun addImportIfNeeded(descriptor: DeclarationDescriptor, file: JetFile) {
+        ImportInsertHelper.addImportDirectiveIfNeeded(DescriptorUtils.getFqNameSafe(descriptor), file)
+    }
 }
 
 //TODO: how about such function in stdlib?

idea/testData/completion/handlers/smart/ConstructorInsertsImport2.kt.after

@@ -1,6 +1,6 @@
 import java.util.HashMap
-import java.util.AbstractMap
 import java.io.File
+import java.util.AbstractMap
 
 class X<T> {
     fun foo(p: java.util.HashMap<T, java.util.AbstractMap<T, java.io.File>>){}

idea/testData/shortenRefs/constructor/GenericType.kt.after

@@ -1,6 +1,6 @@
-import java.util.HashMap
-import java.util.ArrayList
 import java.io.File
+import java.util.ArrayList
+import java.util.HashMap
 
 class A {
     val x = HashMap<ArrayList<File>, String>()

idea/testData/shortenRefs/type/GenericType.kt.after

@@ -1,6 +1,6 @@
-import java.util.HashMap
-import java.util.ArrayList
 import java.io.File
+import java.util.ArrayList
+import java.util.HashMap
 
 class A {
     val x: HashMap<ArrayList<File>, String>

idea/testData/shortenRefs/type/SameClassTwice.kt.after

@@ -1,5 +1,5 @@
-import java.util.ArrayList
 import java.io.File
+import java.util.ArrayList
 
 class A {
     val x: ArrayList<ArrayList<File>>

idea/testData/shortenRefs/type/TwoClassesSameNames.kt.after

@@ -1,6 +1,6 @@
-import java.util.HashMap
 import java.util.Date
 import java.sql
+import java.util.HashMap
 
 class A {
     val x: HashMap<Date, sql.Date>