Move: Fix detection of imported declarations and processing of extension usages

#KT-5787 Fixed
 #KT-4959 Fixed

annotations/com/intellij/refactoring/move/moveClassesOrPackages/annotations.xml

@@ -0,0 +1,6 @@
+<root>
+    <item
+        name='com.intellij.refactoring.move.moveClassesOrPackages.CommonMoveUtil com.intellij.refactoring.util.NonCodeUsageInfo[] retargetUsages(com.intellij.usageView.UsageInfo[], java.util.Map&lt;com.intellij.psi.PsiElement,com.intellij.psi.PsiElement&gt;)'>
+        <annotation name='org.jetbrains.annotations.NotNull'/>
+    </item>
+</root>

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

@@ -29,25 +29,28 @@ import org.jetbrains.jet.lang.psi.JetCallExpression
 import org.jetbrains.jet.lang.psi.JetSimpleNameExpression
 import java.util.Collections
 import org.jetbrains.jet.plugin.caches.resolve.getLazyResolveSession
+import org.jetbrains.jet.lang.psi.JetQualifiedExpression
+import org.jetbrains.jet.lang.descriptors.ClassDescriptor
+import org.jetbrains.jet.lang.descriptors.PackageViewDescriptor
 
 object JetFileReferencesResolver {
     fun resolve(
             element: JetElement,
-            visitReceivers: Boolean = true,
+            resolveQualifiers: Boolean = true,
-            visitShortNames: Boolean = true
+            resolveShortNames: Boolean = true
     ): Map<JetReferenceExpression, BindingContext> {
         return (element.getContainingFile() as? JetFile)?.let { file ->
-            resolve(file, listOf(element), visitReceivers, visitShortNames)
+            resolve(file, listOf(element), resolveQualifiers, resolveShortNames)
         } ?: Collections.emptyMap()
     }
 
     fun resolve(
             file: JetFile,
             elements: Iterable<JetElement>? = null,
-            visitReceivers: Boolean = true,
+            resolveQualifiers: Boolean = true,
-            visitShortNames: Boolean = true
+            resolveShortNames: Boolean = true
     ): Map<JetReferenceExpression, BindingContext> {
-        val visitor = ResolveAllReferencesVisitor(file, visitReceivers, visitShortNames)
+        val visitor = ResolveAllReferencesVisitor(file, resolveQualifiers, resolveShortNames)
         if (elements != null) {
             elements.forEach { it.accept(visitor) }
         }
@@ -57,18 +60,18 @@ object JetFileReferencesResolver {
         return visitor.result
     }
 
-    private class ResolveAllReferencesVisitor(file: JetFile, val visitReceivers: Boolean, val visitShortNames: Boolean) : JetTreeVisitorVoid() {
+    private class ResolveAllReferencesVisitor(file: JetFile, val resolveQualifiers: Boolean, val resolveShortNames: Boolean) : JetTreeVisitorVoid() {
         private val resolveSession = file.getLazyResolveSession()
         private val resolveMap = HashMap<JetReferenceExpression, BindingContext>()
 
         public val result: Map<JetReferenceExpression, BindingContext> = resolveMap
 
         override fun visitUserType(userType: JetUserType) {
-            if (visitReceivers) {
+            if (resolveQualifiers) {
                 userType.acceptChildren(this)
             }
 
-            if (visitShortNames || userType.getQualifier() != null) {
+            if (resolveShortNames || userType.getQualifier() != null) {
                 val referenceExpression = userType.getReferenceExpression()
                 if (referenceExpression != null) {
                     resolveMap[referenceExpression] = resolveSession.resolveToElement(referenceExpression)
@@ -76,21 +79,24 @@ object JetFileReferencesResolver {
             }
         }
 
-        override fun visitDotQualifiedExpression(expression: JetDotQualifiedExpression) {
+        private fun JetExpression.isReceiver(): Boolean {
-            if (visitReceivers) {
+            val parent = getParent()
-                expression.acceptChildren(this)
+            if (parent !is JetQualifiedExpression) return false
-            }
+            if (parent.getReceiverExpression() == this) return true
 
-            val referenceExpression = expression.getSelectorExpression()?.referenceExpression()
+            val parentParent = parent.getParent()
-            if (referenceExpression != null) {
+            return parentParent is JetQualifiedExpression && parentParent.getReceiverExpression() == parent
-                resolveMap[referenceExpression] = resolveSession.resolveToElement(referenceExpression)
-            }
         }
 
         override fun visitSimpleNameExpression(expression: JetSimpleNameExpression) {
-            if (visitShortNames) {
+            val context = resolveSession.resolveToElement(expression)
-                resolveMap[expression] = resolveSession.resolveToElement(expression)
+            val descriptor = context[BindingContext.REFERENCE_TARGET, expression]
+            if ((descriptor is ClassDescriptor || descriptor is PackageViewDescriptor) && expression.isReceiver()) {
+                if (!resolveQualifiers) return
             }
+            if (!resolveShortNames) return
+
+            resolveMap[expression] = context
         }
     }
 }

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

@@ -101,7 +101,7 @@ public object ShortenReferences {
             ImportInsertHelper.optimizeImportsIfNeeded(file)
 
             // first resolve all qualified references - optimization
-            val referenceToContext = JetFileReferencesResolver.resolve(file, fileElements, visitShortNames = false)
+            val referenceToContext = JetFileReferencesResolver.resolve(file, fileElements, resolveShortNames = false)
 
             val shortenTypesVisitor = ShortenTypesVisitor(file, elementFilter, referenceToContext)
             processElements(fileElements, shortenTypesVisitor)

idea/src/org/jetbrains/jet/plugin/refactoring/move/moveFilesOrDirectories/MoveKotlinFileHandler.kt

@@ -25,7 +25,6 @@ import com.intellij.openapi.roots.JavaProjectRootsUtil
 import com.intellij.psi.PsiCompiledElement
 import org.jetbrains.jet.lang.psi.JetFile
 import java.util.ArrayList
-import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.refactoring.util.MoveRenameUsageInfo
 import org.jetbrains.jet.lang.psi.JetNamedDeclaration
@@ -39,7 +38,6 @@ import org.jetbrains.jet.lang.psi.psiUtil.getPackage
 import org.jetbrains.jet.plugin.references.JetReference
 import org.jetbrains.jet.asJava.toLightElements
 import org.jetbrains.jet.lang.psi.JetDeclaration
-import org.jetbrains.jet.lang.psi.JetPsiUtil
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.diagnostic.Logger
 import org.jetbrains.jet.lang.psi.JetClassOrObject
@@ -48,12 +46,12 @@ import org.jetbrains.jet.lang.resolve.DescriptorUtils
 import com.intellij.refactoring.util.TextOccurrencesUtil
 import java.util.Collections
 import org.jetbrains.jet.plugin.refactoring.move.PackageNameInfo
-import org.jetbrains.jet.plugin.refactoring.move.updateInternalReferencesOnPackageNameChange
 import org.jetbrains.jet.plugin.search.projectScope
 import org.jetbrains.jet.plugin.search.fileScope
-import org.jetbrains.jet.plugin.search.not
-import org.jetbrains.jet.plugin.search.and
 import org.jetbrains.jet.plugin.search.minus
+import org.jetbrains.jet.plugin.refactoring.move.getInternalReferencesToUpdateOnPackageNameChange
+import org.jetbrains.jet.plugin.refactoring.move.postProcessMoveUsages
+import org.jetbrains.jet.plugin.codeInsight.addToShorteningWaitSet
 
 public class MoveKotlinFileHandler : MoveFileHandler() {
     class object {
@@ -181,7 +179,8 @@ public class MoveKotlinFileHandler : MoveFileHandler() {
         val packageNameInfo = file.getAndRemoveCopyableUserData(PACKAGE_NAME_INFO_KEY)
         if (packageNameInfo == null) return
 
-        file.updateInternalReferencesOnPackageNameChange(packageNameInfo, updateImportedReferences = false)
+        val usages = file.getInternalReferencesToUpdateOnPackageNameChange(packageNameInfo)
+        postProcessMoveUsages(usages)
 
         val packageRef = file.getPackageDirective()?.getLastReferenceExpression()?.getReference() as? JetSimpleNameReference
         packageRef?.bindToFqName(packageNameInfo.newPackageName)

idea/src/org/jetbrains/jet/plugin/refactoring/move/moveTopLevelDeclarations/MoveKotlinTopLevelDeclarationsProcessor.kt

@@ -27,7 +27,6 @@ import com.intellij.usageView.UsageViewUtil
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.psi.PsiElement
 import com.intellij.refactoring.util.NonCodeUsageInfo
-import com.intellij.refactoring.move.moveClassesOrPackages.CommonMoveUtil
 import com.intellij.util.IncorrectOperationException
 import com.intellij.refactoring.util.RefactoringUIUtil
 import org.jetbrains.jet.utils.keysToMap
@@ -39,7 +38,6 @@ import org.jetbrains.jet.lang.psi.psiUtil.getPackage
 import org.jetbrains.jet.lang.psi.JetFile
 import org.jetbrains.jet.plugin.refactoring.move.PackageNameInfo
 import org.jetbrains.jet.plugin.refactoring.createKotlinFile
-import org.jetbrains.jet.plugin.refactoring.move.updateInternalReferencesOnPackageNameChange
 import org.jetbrains.jet.plugin.codeInsight.addToShorteningWaitSet
 import org.jetbrains.jet.plugin.refactoring.move.getFileNameAfterMove
 import org.jetbrains.jet.lang.psi.JetNamedDeclaration
@@ -67,13 +65,13 @@ import com.intellij.psi.PsiModifier
 import com.intellij.util.VisibilityUtil
 import com.intellij.openapi.util.Ref
 import org.jetbrains.jet.plugin.refactoring.getKotlinFqName
-import org.jetbrains.jet.plugin.references.JetSimpleNameReference.ShorteningMode
 import org.jetbrains.jet.plugin.search.projectScope
-import org.jetbrains.jet.lang.psi.psiUtil.isExtensionDeclaration
+import org.jetbrains.jet.plugin.refactoring.move.getInternalReferencesToUpdateOnPackageNameChange
-import org.jetbrains.jet.plugin.references.JetReference
+import org.jetbrains.jet.plugin.refactoring.move.createMoveUsageInfo
-import org.jetbrains.jet.plugin.quickfix.ImportInsertHelper
+import org.jetbrains.jet.plugin.refactoring.move.postProcessMoveUsages
-import org.jetbrains.jet.lang.psi.psiUtil.getParentByType
+import org.jetbrains.jet.plugin.references.JetSimpleNameReference.ShorteningMode
-import org.jetbrains.jet.lang.psi.JetImportDirective
+import org.jetbrains.jet.lang.psi.psiUtil.isAncestor
+import org.jetbrains.jet.plugin.refactoring.move.MoveRenameUsageInfoForExtension
 
 public class MoveKotlinTopLevelDeclarationsOptions(
         val elementsToMove: Collection<JetNamedDeclaration>,
@@ -83,14 +81,6 @@ public class MoveKotlinTopLevelDeclarationsOptions(
         val moveCallback: MoveCallback? = null
 )
 
-class MoveRenameUsageInfoForExtension(
-        element: PsiElement,
-        reference: PsiReference,
-        startOffset: Int,
-        endOffset: Int,
-        referencedElement: PsiElement
-): MoveRenameUsageInfo(element, reference, startOffset, endOffset, referencedElement, false)
-
 public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val options: MoveKotlinTopLevelDeclarationsOptions) : BaseRefactoringProcessor(project) {
     class object {
         private val LOG: Logger = Logger.getInstance(javaClass<MoveKotlinTopLevelDeclarationsProcessor>())
@@ -111,21 +101,6 @@ public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val optio
     }
 
     override fun findUsages(): Array<UsageInfo> {
-        fun createUsageInfo(
-                element: PsiElement,
-                reference: PsiReference,
-                startOffset: Int,
-                endOffset: Int,
-                referencedElement: PsiElement
-        ): UsageInfo {
-            if (reference is JetReference
-                    && referencedElement.namedUnwrappedElement!!.isExtensionDeclaration()
-                    && element.getParentByType(javaClass<JetImportDirective>()) == null) {
-                return MoveRenameUsageInfoForExtension(element, reference, startOffset, endOffset, referencedElement)
-            }
-            return MoveRenameUsageInfo(element, reference, startOffset, endOffset, referencedElement, false)
-        }
-
         val newPackageName = options.moveTarget.packageWrapper?.getQualifiedName() ?: ""
 
         fun collectUsages(): List<UsageInfo> {
@@ -137,9 +112,8 @@ public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val optio
                 val results = ReferencesSearch
                         .search(lightElement, projectScope, false)
                         .mapTo(ArrayList<UsageInfo?>()) { ref ->
-                            if (foundReferences.add(ref)) {
+                            if (foundReferences.add(ref) && elementsToMove.all { !it.isAncestor(ref.getElement())}) {
-                                val range = ref.getRangeInElement()!!
+                                createMoveUsageInfo(ref, lightElement, true)
-                                createUsageInfo(ref.getElement(), ref, range.getStartOffset(), range.getEndOffset(), lightElement)
                             }
                             else null
                         }
@@ -201,7 +175,7 @@ public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val optio
         fun collectConflictsInDeclarations() {
             val declarationToReferenceTargets = HashMap<JetNamedDeclaration, MutableSet<PsiElement>>()
             for (declaration in elementsToMove) {
-                val referenceToContext = JetFileReferencesResolver.resolve(element = declaration, visitReceivers = false)
+                val referenceToContext = JetFileReferencesResolver.resolve(element = declaration, resolveQualifiers = false)
                 for ((refExpr, bindingContext) in referenceToContext) {
                     val refTarget = bindingContext[BindingContext.REFERENCE_TARGET, refExpr]?.let { descriptor ->
                         DescriptorToDeclarationUtil.getDeclaration(declaration.getProject(), descriptor)
@@ -246,7 +220,11 @@ public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val optio
     }
 
     override fun performRefactoring(usages: Array<out UsageInfo>?) {
-        fun moveDeclaration(declaration: JetNamedDeclaration, moveTarget: KotlinMoveTarget): JetNamedDeclaration? {
+        fun moveDeclaration(
+                declaration: JetNamedDeclaration,
+                moveTarget: KotlinMoveTarget,
+                usagesToProcessAfterMove: MutableList<UsageInfo>
+        ): JetNamedDeclaration? {
             val file = declaration.getContainingFile() as? JetFile
             assert (file != null, "${declaration.javaClass}: ${declaration.getText()}")
 
@@ -265,14 +243,12 @@ public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val optio
 
             assert(targetFile is JetFile, "Couldn't create Koltin file for: ${declaration.javaClass}: ${declaration.getText()}")
 
-            val newPackageFqName = (targetFile as JetFile).getPackageFqName()
+            val packageNameInfo = PackageNameInfo(file!!.getPackageFqName(), (targetFile as JetFile).getPackageFqName())
-
+            val (usagesToProcessLater, usagesToProcessNow) = declaration
-            val packageNameInfo = PackageNameInfo(file!!.getPackageFqName(), newPackageFqName)
+                    .getInternalReferencesToUpdateOnPackageNameChange(packageNameInfo)
-            declaration.updateInternalReferencesOnPackageNameChange(
+                    .partition { it is MoveRenameUsageInfoForExtension }
-                    packageNameInfo,
+            postProcessMoveUsages(usagesToProcessNow, shorteningMode = ShorteningMode.NO_SHORTENING)
-                    updateImportedReferences = true,
+            usagesToProcessAfterMove.addAll(usagesToProcessLater)
-                    shorteningMode = ShorteningMode.NO_SHORTENING
-            )
 
             val newElement = targetFile.add(declaration) as JetNamedDeclaration
             declaration.delete()
@@ -283,9 +259,13 @@ public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val optio
         }
 
         try {
+            val usageList = usages.toArrayList()
+
             val oldToNewElementsMapping = HashMap<PsiElement, PsiElement>()
             for ((oldDeclaration, oldLightElements) in kotlinToLightElements) {
-                val newDeclaration = moveDeclaration(oldDeclaration, options.moveTarget)
+                val oldFile = oldDeclaration.getContainingJetFile()
+
+                val newDeclaration = moveDeclaration(oldDeclaration, options.moveTarget, usageList)
                 if (newDeclaration == null) {
                     for (oldElement in oldLightElements) {
                         oldToNewElementsMapping[oldElement] = oldElement
@@ -293,26 +273,15 @@ public class MoveKotlinTopLevelDeclarationsProcessor(project: Project, val optio
                     continue
                 }
 
+                oldToNewElementsMapping[oldFile] = newDeclaration.getContainingJetFile()
+
                 getTransaction()!!.getElementListener(oldDeclaration).elementMoved(newDeclaration)
                 for ((oldElement, newElement) in oldLightElements.stream() zip newDeclaration.toLightElements().stream()) {
                     oldToNewElementsMapping[oldElement] = newElement
                 }
             }
 
-            if (usages == null) return
+            nonCodeUsages = postProcessMoveUsages(usageList, oldToNewElementsMapping).copyToArray()
-
-            nonCodeUsages = CommonMoveUtil.retargetUsages(
-                    usages.filter { it !is MoveRenameUsageInfoForExtension }.copyToArray(),
-                    oldToNewElementsMapping
-            )
-            for (usage in usages.filterIsInstance(javaClass<MoveRenameUsageInfoForExtension>())) {
-                val newElement = oldToNewElementsMapping[usage.getReferencedElement()]
-
-                ImportInsertHelper.addImportDirectiveIfNeeded(
-                        newElement!!.getKotlinFqName()!!,
-                        usage.getElement()!!.getContainingFile() as JetFile
-                )
-            }
         }
         catch (e: IncorrectOperationException) {
             nonCodeUsages = null

idea/src/org/jetbrains/jet/plugin/refactoring/move/moveUtils.kt

@@ -26,7 +26,6 @@ import org.jetbrains.jet.plugin.references.JetSimpleNameReference
 import org.jetbrains.jet.lang.resolve.name.FqName
 import org.jetbrains.jet.lang.psi.JetFile
 import org.jetbrains.jet.lang.psi.JetElement
-import org.jetbrains.jet.plugin.references.JetSimpleNameReference.ShorteningMode
 import org.jetbrains.jet.plugin.JetFileType
 import org.jetbrains.jet.lang.psi.JetNamedDeclaration
 import org.jetbrains.jet.plugin.imports.canBeReferencedViaImport
@@ -34,30 +33,54 @@ import org.jetbrains.jet.plugin.codeInsight.DescriptorToDeclarationUtil
 import org.jetbrains.jet.lang.psi.psiUtil.isAncestor
 import java.util.Collections
 import org.jetbrains.jet.lang.resolve.name.isImported
+import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor
+import org.jetbrains.jet.lang.descriptors.ClassDescriptor
+import org.jetbrains.jet.lang.descriptors.PackageViewDescriptor
+import com.intellij.usageView.UsageInfo
+import org.jetbrains.jet.lang.psi.psiUtil.isExtensionDeclaration
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiReference
+import com.intellij.refactoring.util.MoveRenameUsageInfo
+import org.jetbrains.jet.plugin.references.JetReference
+import org.jetbrains.jet.asJava.namedUnwrappedElement
+import org.jetbrains.jet.lang.psi.psiUtil.getParentByType
+import org.jetbrains.jet.lang.psi.JetImportDirective
+import java.util.ArrayList
+import com.intellij.refactoring.util.NonCodeUsageInfo
+import org.jetbrains.jet.plugin.quickfix.ImportInsertHelper
+import org.jetbrains.jet.plugin.refactoring.getKotlinFqName
+import org.jetbrains.jet.lang.psi.JetThisExpression
+import org.jetbrains.jet.plugin.references.JetSimpleNameReference.ShorteningMode
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.openapi.util.Comparing
+import java.util.Comparator
+import com.intellij.util.IncorrectOperationException
+import com.intellij.psi.PsiFile
 
 public class PackageNameInfo(val oldPackageName: FqName, val newPackageName: FqName)
 
-public fun JetElement.updateInternalReferencesOnPackageNameChange(
+public fun JetElement.getInternalReferencesToUpdateOnPackageNameChange(packageNameInfo: PackageNameInfo): List<UsageInfo> {
-        packageNameInfo: PackageNameInfo,
-        updateImportedReferences: Boolean,
-        shorteningMode: ShorteningMode = ShorteningMode.DELAYED_SHORTENING
-) {
     val file = getContainingFile() as? JetFile
-    if (file == null) return
+    if (file == null) return Collections.emptyList()
 
-    val importPaths =
+    val importPaths = file.getImportDirectives().map { it.getImportPath() }.filterNotNull()
-            if (updateImportedReferences) {
-                file.getImportDirectives().map { it.getImportPath() }.filterNotNull()
-            }
-            else Collections.emptyList()
 
-    fun isImportedName(fqName: FqName): Boolean =
+    [tailRecursive] fun isImported(descriptor: DeclarationDescriptor): Boolean {
-            updateImportedReferences && importPaths.any { fqName.isImported(it, false) }
+        val fqName = DescriptorUtils.getFqName(descriptor).let { if (it.isSafe()) it.toSafe() else return@isImported false }
+        if (importPaths.any { fqName.isImported(it, false) }) return true
 
-    val referenceToContext = JetFileReferencesResolver.resolve(file = file, elements = listOf(this), visitReceivers = false)
+        val containingDescriptor = descriptor.getContainingDeclaration()
+        return when (containingDescriptor) {
+            is ClassDescriptor, is PackageViewDescriptor -> isImported(containingDescriptor)
+            else -> false
+        }
+    }
 
+    val referenceToContext = JetFileReferencesResolver.resolve(file = file, elements = listOf(this), resolveQualifiers = false)
+
+    val usages = ArrayList<UsageInfo>()
     for ((refExpr, bindingContext) in referenceToContext) {
-        if (refExpr !is JetSimpleNameExpression) continue
+        if (refExpr !is JetSimpleNameExpression || refExpr.getParent() is JetThisExpression) continue
 
         val descriptor = bindingContext[BindingContext.REFERENCE_TARGET, refExpr]?.let { descriptor ->
             if (descriptor is ConstructorDescriptor) descriptor.getContainingDeclaration() else descriptor
@@ -70,23 +93,122 @@ public fun JetElement.updateInternalReferencesOnPackageNameChange(
         val fqName = DescriptorUtils.getFqName(descriptor)
         if (!fqName.isSafe()) continue
 
-        val fqNameSafe = fqName.toSafe()
+        val packageName = DescriptorUtils.getParentOfType(descriptor, javaClass<PackageFragmentDescriptor>(), false)?.let {
-        val packageName = DescriptorUtils.getParentOfType(
+            DescriptorUtils.getFqName(it).toSafe()
-                descriptor, javaClass<PackageFragmentDescriptor>(), false
+        }
-        )?.let { DescriptorUtils.getFqName(it).toSafe() }
 
         when {
+            declaration.isExtensionDeclaration(),
             packageName == packageNameInfo.oldPackageName,
             packageName == packageNameInfo.newPackageName,
-            isImportedName(fqNameSafe) -> {
+            isImported(descriptor) -> {
-                (refExpr.getReference() as? JetSimpleNameReference)?.bindToFqName(fqNameSafe, shorteningMode)
+                (refExpr.getReference() as? JetSimpleNameReference)?.let { usages.add(createMoveUsageInfo(it, declaration, false)) }
             }
         }
     }
+
+    return usages
+}
+
+class MoveRenameUsageInfoForExtension(
+        element: PsiElement,
+        reference: PsiReference,
+        startOffset: Int,
+        endOffset: Int,
+        referencedElement: PsiElement,
+        val originalFile: PsiFile,
+        val addImportToOriginalFile: Boolean
+): MoveRenameUsageInfo(element, reference, startOffset, endOffset, referencedElement, false)
+
+fun createMoveUsageInfo(
+        reference: PsiReference,
+        referencedElement: PsiElement,
+        addImportToOriginalFile: Boolean
+): UsageInfo {
+    val range = reference.getRangeInElement()!!
+    val element = reference.getElement()
+    val startOffset = range.getStartOffset()
+    val endOffset = range.getEndOffset()
+
+    if (reference is JetReference
+        && referencedElement.namedUnwrappedElement!!.isExtensionDeclaration()
+        && element.getParentByType(javaClass<JetImportDirective>()) == null) {
+        return MoveRenameUsageInfoForExtension(
+                element, reference, startOffset, endOffset, referencedElement, element.getContainingFile()!!, addImportToOriginalFile
+        )
+    }
+    return MoveRenameUsageInfo(element, reference, startOffset, endOffset, referencedElement, false)
 }
 
 public fun JetNamedDeclaration.getFileNameAfterMove(): String? {
     return (getContainingFile() as? JetFile)?.let { file ->
         if (file.getDeclarations().size > 1) "${getName()}.${JetFileType.INSTANCE.getDefaultExtension()}" else file.getName()
     }
+}
+
+/**
+ * Perform usage postprocessing and return non-code usages
+ */
+fun postProcessMoveUsages(usages: List<UsageInfo>,
+                          oldToNewElementsMapping: Map<PsiElement, PsiElement> = Collections.emptyMap(),
+                          shorteningMode: ShorteningMode = ShorteningMode.DELAYED_SHORTENING
+): List<NonCodeUsageInfo> {
+    fun counterpart(e: PsiElement) = oldToNewElementsMapping[e] ?: e
+
+    val sortedUsages = usages.sortBy(
+            object : Comparator<UsageInfo> {
+                override fun compare(o1: UsageInfo, o2: UsageInfo): Int {
+                    val file1 = o1.getVirtualFile()
+                    val file2 = o2.getVirtualFile()
+                    if (Comparing.equal<VirtualFile>(file1, file2)) {
+                        val rangeInElement1 = o1.getRangeInElement()
+                        val rangeInElement2 = o2.getRangeInElement()
+                        if (rangeInElement1 != null && rangeInElement2 != null) {
+                            return rangeInElement2.getStartOffset() - rangeInElement1.getStartOffset()
+                        }
+                        return 0
+                    }
+                    if (file1 == null) return -1
+                    if (file2 == null) return 1
+                    return Comparing.compare<String>(file1.getPath(), file2.getPath())
+                }
+            }
+    )
+
+    val nonCodeUsages = ArrayList<NonCodeUsageInfo>()
+
+    for (usage in sortedUsages) {
+        when (usage) {
+            is NonCodeUsageInfo -> {
+                nonCodeUsages.add(usage as NonCodeUsageInfo)
+            }
+
+            is MoveRenameUsageInfoForExtension -> {
+                val element = counterpart(usage.getReferencedElement()!!)
+                val file = with(usage) { if (addImportToOriginalFile) originalFile else counterpart(originalFile) } as JetFile
+                ImportInsertHelper.addImportDirectiveIfNeeded(element.getKotlinFqName()!!, file)
+            }
+
+            is MoveRenameUsageInfo -> {
+                val moveRenameUsage = usage as MoveRenameUsageInfo
+                val oldElement = moveRenameUsage.getReferencedElement()!!
+                val newElement = counterpart(oldElement)
+                moveRenameUsage.getReference()?.let {
+                    try {
+                        if (it is JetSimpleNameReference) {
+                            it.bindToElement(newElement, shorteningMode)
+                        }
+                        else {
+                            it.bindToElement(newElement)
+                        }
+                    }
+                    catch (e: IncorrectOperationException) {
+                        // Suppress exception if bindToElement is not implemented
+                    }
+                }
+            }
+        }
+    }
+
+    return nonCodeUsages
 }

idea/src/org/jetbrains/jet/plugin/references/JetSimpleNameReference.kt

@@ -105,9 +105,11 @@ public class JetSimpleNameReference(
     }
 
     // By default reference binding is delayed
-    override fun bindToElement(element: PsiElement): PsiElement {
+    override fun bindToElement(element: PsiElement): PsiElement =
-        return element.getKotlinFqName()?.let { fqName -> bindToFqName(fqName) } ?: expression
+            bindToElement(element, ShorteningMode.DELAYED_SHORTENING)
-    }
+
+    fun bindToElement(element: PsiElement, shorteningMode: ShorteningMode): PsiElement =
+            element.getKotlinFqName()?.let { fqName -> bindToFqName(fqName, shorteningMode) } ?: expression
 
     public fun bindToFqName(fqName: FqName, shorteningMode: ShorteningMode = ShorteningMode.DELAYED_SHORTENING): PsiElement {
         if (fqName.isRoot()) return expression

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/after/library/misc.kt

@@ -0,0 +1,7 @@
+package library
+
+open public class B {
+    fun foo(): B = this
+}
+
+fun B.bar(): B = this

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/after/source/some.kt

@@ -0,0 +1,5 @@
+package source
+
+import target.A
+
+fun test1() = target.test(A())

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/after/target/A.kt

@@ -0,0 +1,21 @@
+package target
+
+import library.B
+import library.bar
+
+public class A: B() {
+    fun test() {
+        bar()
+        foo()
+        this.bar()
+        this.foo()
+    }
+}
+
+public fun test(a: A) {
+    a.bar()
+    a.foo()
+    B().foo().bar()
+    B().foo().bar()
+    B().bar().foo()
+}

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/after/target/some.kt

@@ -0,0 +1,3 @@
+package target
+
+fun test2() = test(A())

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/before/library/misc.kt

@@ -0,0 +1,7 @@
+package library
+
+open public class B {
+    fun foo(): B = this
+}
+
+fun B.bar(): B = this

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/before/source/A.kt

@@ -0,0 +1,21 @@
+package source
+
+import library.B
+import library.bar
+
+public class A: B() {
+    fun test() {
+        bar()
+        foo()
+        this.bar()
+        this.foo()
+    }
+}
+
+public fun test(a: A) {
+    a.bar()
+    a.foo()
+    B().foo().bar()
+    library.B().foo().bar()
+    B().bar().foo()
+}

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/before/source/some.kt

@@ -0,0 +1,3 @@
+package source
+
+fun test1() = test(A())

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/before/target/some.kt

@@ -0,0 +1,3 @@
+package target
+
+fun test2() = source.test(source.A())

idea/testData/refactoring/move/kotlin/moveFile/internalReferences/internalReferences.test

@@ -0,0 +1,5 @@
+{
+  "mainFile": "source/A.kt",
+  "type": "MOVE_FILES",
+  "targetPackage": "target"
+}

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/after/library/some.kt

@@ -0,0 +1,7 @@
+package library
+
+open public class B {
+    fun foo(): B = this
+}
+
+fun B.bar(): B = this

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/after/source/some.kt

@@ -0,0 +1,5 @@
+package source
+
+import target.A
+
+fun test1() = A()

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/after/target/A.kt

@@ -0,0 +1,16 @@
+package target
+
+import library.bar
+import library.B
+
+public class A: B() {
+    fun test() {
+        bar()
+        foo()
+        this.bar()
+        this.foo()
+        B().foo().bar()
+        B().foo().bar()
+        B().bar().foo()
+    }
+}

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/after/target/some.kt

@@ -0,0 +1,3 @@
+package target
+
+fun test2() = A()

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/before/library/some.kt

@@ -0,0 +1,7 @@
+package library
+
+open public class B {
+    fun foo(): B = this
+}
+
+fun B.bar(): B = this

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/before/source/A.kt

@@ -0,0 +1,16 @@
+package source
+
+import library.B
+import library.bar
+
+public class <caret>A: B() {
+    fun test() {
+        bar()
+        foo()
+        this.bar()
+        this.foo()
+        B().foo().bar()
+        library.B().foo().bar()
+        B().bar().foo()
+    }
+}

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/before/source/some.kt

@@ -0,0 +1,3 @@
+package source
+
+fun test1() = A()

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/before/target/some.kt

@@ -0,0 +1,3 @@
+package target
+
+fun test2() = source.A()

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/internalReferences.test

@@ -0,0 +1,5 @@
+{
+  "mainFile": "source/A.kt",
+  "type": "MOVE_KOTLIN_TOP_LEVEL_DECLARATIONS",
+  "targetPackage": "target"
+}

idea/tests/org/jetbrains/jet/plugin/refactoring/move/JetMoveTestGenerated.java

@@ -212,6 +212,12 @@ public class JetMoveTestGenerated extends AbstractJetMoveTest {
         doTest(fileName);
     }
     
+    @TestMetadata("kotlin/moveFile/internalReferences/internalReferences.test")
+    public void testKotlin_moveFile_internalReferences_InternalReferences() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/move/kotlin/moveFile/internalReferences/internalReferences.test");
+        doTest(fileName);
+    }
+    
     @TestMetadata("kotlin/moveFile/moveFileToFile/moveFileToFile.test")
     public void testKotlin_moveFile_moveFileToFile_MoveFileToFile() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/move/kotlin/moveFile/moveFileToFile/moveFileToFile.test");
@@ -236,6 +242,12 @@ public class JetMoveTestGenerated extends AbstractJetMoveTest {
         doTest(fileName);
     }
     
+    @TestMetadata("kotlin/moveTopLevelDeclarations/misc/internalReferences/internalReferences.test")
+    public void testKotlin_moveTopLevelDeclarations_misc_internalReferences_InternalReferences() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/internalReferences/internalReferences.test");
+        doTest(fileName);
+    }
+    
     @TestMetadata("kotlin/moveTopLevelDeclarations/misc/moveClassWithImportsRetained/moveClassWithImportsRetained.test")
     public void testKotlin_moveTopLevelDeclarations_misc_moveClassWithImportsRetained_MoveClassWithImportsRetained() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/moveClassWithImportsRetained/moveClassWithImportsRetained.test");