Rewritten AddForLoopIndicesIntention to not use by text generation & checks

core/descriptors/src/org/jetbrains/kotlin/resolve/DescriptorUtils.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind.ENUM_ENTRY
 import org.jetbrains.kotlin.descriptors.ClassKind.OBJECT
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.FqNameUnsafe
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.types.JetType
 
@@ -36,6 +37,12 @@ public fun DeclarationDescriptor.getImportableDescriptor(): DeclarationDescripto
     }
 }
 
+public val DeclarationDescriptor.fqNameUnsafe: FqNameUnsafe
+    get() = DescriptorUtils.getFqName(this)
+
+public val DeclarationDescriptor.fqNameSafe: FqName
+    get() = DescriptorUtils.getFqNameSafe(this)
+
 public val DeclarationDescriptor.isExtension: Boolean
     get() = this is CallableDescriptor && getExtensionReceiverParameter() != null
 

idea/src/org/jetbrains/kotlin/idea/intentions/AddForLoopIndicesIntention.kt

@@ -18,67 +18,67 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.codeInsight.template.TemplateBuilderImpl
 import com.intellij.codeInsight.template.TemplateManager
-import com.intellij.codeInsight.template.impl.TemplateManagerImpl
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiDocumentManager
-import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.analyzer.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.idea.core.replaced
+import org.jetbrains.kotlin.psi.JetExpression
+import org.jetbrains.kotlin.psi.JetForExpression
+import org.jetbrains.kotlin.psi.JetPsiFactory
+import org.jetbrains.kotlin.psi.createExpressionByPattern
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
-public class AddForLoopIndicesIntention : JetSelfTargetingIntention<JetForExpression>(
+public class AddForLoopIndicesIntention : JetSelfTargetingRangeIntention<JetForExpression>(javaClass(), "Add indices to 'for' loop") {
-        javaClass(), "Add indices to 'for' loop") {
     override fun applyTo(element: JetForExpression, editor: Editor) {
-        val loopRange = element.getLoopRange()!!
+        val loopRange = element.loopRange!!
-        val newRangeText = "${loopRange.getText()}.withIndex()"
+        val loopParameter = element.loopParameter!!
-        val project = editor.getProject()!!
+        val project = element.project
-        val psiFactory = JetPsiFactory(project)
+        val psiFactory = JetPsiFactory(element)
-        val newRange = psiFactory.createExpression(newRangeText)
 
-        //Roundabout way to create new multiparameter element so as not to incorrectly trigger syntax error highlighting
+        loopRange.replace(createWithIndexExpression(loopRange))
-        val loopParameter = element.getLoopParameter()!!
-        val parenthesizedParam = psiFactory.createExpression("(index)") as JetParenthesizedExpression
-        val indexElement = parenthesizedParam.getExpression()!!
-        val comma = psiFactory.createComma()
-        val newParamElement = psiFactory.createExpression(loopParameter.getText())
-        parenthesizedParam.addAfter(newParamElement, indexElement)
-        parenthesizedParam.addAfter(comma, indexElement)
 
-        loopParameter.replace(parenthesizedParam)
+        var multiParameter = (psiFactory.createExpressionByPattern("for((index, $0) in x){}", loopParameter.text) as JetForExpression).multiParameter!!
-        loopRange.replace(newRange)
 
-        val multiParameter = PsiTreeUtil.findChildOfType(element, indexElement.javaClass)!!
+        multiParameter = loopParameter.replaced(multiParameter)
 
-        editor.getCaretModel().moveToOffset(multiParameter.getTextOffset())
+        val indexVariable = multiParameter.entries[0]
-        val templateBuilder = TemplateBuilderImpl(multiParameter)
+        editor.caretModel.moveToOffset(indexVariable.startOffset)
-        templateBuilder.replaceElement(multiParameter, "index")
+
-        PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(editor.getDocument())
+        val templateBuilder = TemplateBuilderImpl(indexVariable)
+        templateBuilder.replaceElement(indexVariable, "index")
+        PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(editor.document)
         TemplateManager.getInstance(project).startTemplate(editor, templateBuilder.buildInlineTemplate()!!)
     }
 
-    override fun isApplicableTo(element: JetForExpression, caretOffset: Int): Boolean {
+    private val WITH_INDEX_FQ_NAME = "kotlin.withIndex"
-        if (element.getLoopParameter() == null) return false
-        val body = element.getBody()
-        if (body != null && caretOffset >= body.getTextRange().getStartOffset()) return false
 
-        val range = element.getLoopRange() ?: return false
+    override fun applicabilityRange(element: JetForExpression): TextRange? {
-        if (range is JetDotQualifiedExpression) {
+        if (element.loopParameter == null) return null
-            val selector = range.getSelectorExpression() ?: return true
+        val loopRange = element.loopRange ?: return null
-            if (selector.getText() == "withIndex()") return false
-        }
 
-        val psiFactory = JetPsiFactory(element.getProject())
+        val bindingContext = element.analyze(BodyResolveMode.PARTIAL)
-        val potentialExpression = psiFactory.createExpression("${range.getText()}.withIndex()") as JetDotQualifiedExpression
 
-        val bindingContext = element.analyze()
+        val resolvedCall = loopRange.getResolvedCall(bindingContext)
-        val scope = bindingContext[BindingContext.RESOLUTION_SCOPE, element] ?: return false
+        if (resolvedCall?.resultingDescriptor?.fqNameUnsafe?.asString() == WITH_INDEX_FQ_NAME) return null // already withIndex() call
-        val functionSelector = potentialExpression.getSelectorExpression() as JetCallExpression
+
-        val updatedContext = potentialExpression.analyzeInContext(scope)
+        val resolutionScope = bindingContext[BindingContext.RESOLUTION_SCOPE, element] ?: return null
-        val call = updatedContext[BindingContext.CALL, functionSelector.getCalleeExpression()] ?: return false
+        val potentialExpression = createWithIndexExpression(loopRange)
-        val callScope = updatedContext[BindingContext.RESOLVED_CALL, call] ?: return false
+
-        val callFqName = DescriptorUtils.getFqNameSafe(callScope.getCandidateDescriptor())
+        val newBindingContext = potentialExpression.analyzeInContext(resolutionScope)
-        return callFqName.toString() == "kotlin.withIndex"
+        val newResolvedCall = potentialExpression.getResolvedCall(newBindingContext) ?: return null
+        if (newResolvedCall.resultingDescriptor.fqNameUnsafe.asString() != WITH_INDEX_FQ_NAME) return null
+
+        return TextRange(element.startOffset, element.body?.startOffset ?: element.endOffset)
+    }
+
+    private fun createWithIndexExpression(originalExpression: JetExpression): JetExpression {
+        return JetPsiFactory(originalExpression).createExpressionByPattern("$0.withIndex()", originalExpression)
     }
 }