More correct and easy comment restoring range management

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/baseTransformations.kt

@@ -27,9 +27,7 @@ abstract class ReplaceLoopResultTransformation(override val loop: KtForExpressio
 
     override val commentSavingRange = PsiChildRange.singleElement(loop.unwrapIfLabeled())
 
-    override fun commentRestoringRange(convertLoopResult: KtExpression) = PsiChildRange.singleElement(convertLoopResult)
+    override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
-
-    override fun convertLoop(resultCallChain: KtExpression): KtExpression {
         return loop.unwrapIfLabeled().replaced(resultCallChain)
     }
 }
@@ -41,17 +39,13 @@ abstract class AssignToVariableResultTransformation(
 
     override val commentSavingRange = PsiChildRange(initialization.initializationStatement, loop.unwrapIfLabeled())
 
-    //TODO: does not work in case of move
+    override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
-    private val commentRestoringRange = commentSavingRange.withoutLastStatement()
-
-    override fun commentRestoringRange(convertLoopResult: KtExpression) = commentRestoringRange
-
-    override fun convertLoop(resultCallChain: KtExpression): KtExpression {
         initialization.initializer.replace(resultCallChain)
 
-        val previousStatement = loop.unwrapIfLabeled().previousStatement()
+        val loopUnwrapped = loop.unwrapIfLabeled()
+        val previousStatement = loopUnwrapped.previousStatement()
 
-        loop.deleteWithLabels()
+        loopUnwrapped.delete()
 
         if (initialization.variable.isVar && !initialization.variable.hasWriteUsages()) { // change variable to 'val' if possible
             initialization.variable.valOrVarKeyword.replace(KtPsiFactory(initialization.variable).createValKeyword())
@@ -64,6 +58,9 @@ abstract class AssignToVariableResultTransformation(
             assert(block is KtBlockExpression)
             val movedInitializationStatement = block.addAfter(initializationStatement, previousStatement) as KtExpression
             block.addAfter(KtPsiFactory(block).createNewLine(), previousStatement)
+
+            commentSavingRangeHolder.remove(initializationStatement)
+
             initializationStatement.delete()
             initializationStatement = movedInitializationStatement
         }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt

@@ -16,10 +16,14 @@
 
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.psi.KtCallableDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.psi.psiUtil.siblings
 
 /**
  * An abstraction for generating a chained call that knows about receiver expression and handles proper formatting
@@ -73,9 +77,12 @@ interface ResultTransformation : Transformation {
     fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? = null
 
     val commentSavingRange: PsiChildRange
-    fun commentRestoringRange(convertLoopResult: KtExpression): PsiChildRange
 
-    fun convertLoop(resultCallChain: KtExpression): KtExpression
+    /**
+     * Implementations of this method are obliged to update [commentSavingRangeHolder] when deleting or adding any element into the tree
+     * except for the loop itself and the result element returned from this method
+     */
+    fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression
 }
 
 /**
@@ -124,3 +131,68 @@ class ResultTransformationMatch(
     constructor(resultTransformation: ResultTransformation, vararg sequenceTransformations: SequenceTransformation)
     : this(resultTransformation, sequenceTransformations.asList())
 }
+
+/**
+ * Helper class for holding and updating PsiChildRange to be used for [CommentSaver.restore] call
+ */
+class CommentSavingRangeHolder(range: PsiChildRange) {
+    var range = range
+        private set
+
+    /**
+     * Call this method when a new element to be included into the range is added into the tree
+     */
+    fun add(element: PsiElement) {
+        if (range.isEmpty) {
+            range = PsiChildRange.singleElement(element)
+            return
+        }
+
+        val rangeParent = range.first!!.parent
+        val elementToAdd = element.parentsWithSelf.takeWhile { it != rangeParent }.last()
+        when (elementToAdd) {
+            in range -> return
+
+            in range.first!!.siblingsBefore() -> range = PsiChildRange(elementToAdd, range.last)
+
+            else -> range = PsiChildRange(range.first, elementToAdd)
+        }
+    }
+
+    /**
+     * Call this method before deletion of any element which can belong to the range
+     */
+    fun remove(element: PsiElement) {
+        when {
+            range.isEmpty -> return
+
+            element == range.first -> {
+                val newFirst = element
+                        .siblings(forward = true, withItself = false)
+                        .takeWhile { it != range.last!!.nextSibling }
+                        .firstOrNull { it !is PsiWhiteSpace }
+                if (newFirst != null) {
+                    range = PsiChildRange(newFirst, range.last)
+                }
+                else {
+                    range = PsiChildRange.EMPTY
+                }
+            }
+
+            element == range.last -> {
+                val newLast = element
+                        .siblings(forward = false, withItself = false)
+                        .takeWhile { it != range.first!!.prevSibling }
+                        .firstOrNull { it !is PsiWhiteSpace }
+                if (newLast != null) {
+                    range = PsiChildRange(range.first, newLast)
+                }
+                else {
+                    range = PsiChildRange.EMPTY
+                }
+            }
+        }
+    }
+
+    private fun PsiElement.siblingsBefore() = if (prevSibling != null) PsiChildRange(parent.firstChild, prevSibling) else PsiChildRange.EMPTY
+}

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt

@@ -17,8 +17,6 @@
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
 import com.intellij.openapi.util.Key
-import com.intellij.psi.PsiDocumentManager
-import com.intellij.psi.codeStyle.CodeStyleManager
 import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
@@ -140,24 +138,26 @@ fun match(loop: KtForExpression): MatchResult? {
 //TODO: offer to use of .asSequence() as an option
 fun convertLoop(loop: KtForExpression, matchResult: MatchResult): KtExpression {
     val resultTransformation = matchResult.transformationMatch.resultTransformation
-    val commentSaver = CommentSaver(resultTransformation.commentSavingRange)
+
+    val commentSavingRange = resultTransformation.commentSavingRange
+    val commentSaver = CommentSaver(commentSavingRange)
+    val commentSavingRangeHolder = CommentSavingRangeHolder(commentSavingRange)
+
+    matchResult.initializationStatementsToDelete.forEach { commentSavingRangeHolder.add(it) }
 
     val callChain = matchResult.generateCallChain(loop)
 
-    val result = resultTransformation.convertLoop(callChain)
+    commentSavingRangeHolder.remove(loop.unwrapIfLabeled()) // loop will be deleted in all cases
+    val result = resultTransformation.convertLoop(callChain, commentSavingRangeHolder)
+    commentSavingRangeHolder.add(result)
 
-    //TODO: preserve comments?
+    for (statement in matchResult.initializationStatementsToDelete) {
-    matchResult.initializationStatementsToDelete.forEach { it.delete() }
+        commentSavingRangeHolder.remove(statement)
+        statement.delete()
+    }
 
-    commentSaver.restore(resultTransformation.commentRestoringRange(result))
+    // we need to adjust indent of the result because in some cases it's made incorrect when moving statement closer to the loop
-
+    commentSaver.restore(commentSavingRangeHolder.range, forceAdjustIndent = true)
-    // need to manually adjust indent of the result because in some cases it's made incorrect when moving closer to the loop
-    // TODO: use forceAdjustIndent = true of CommentSaver.restore
-    val file = result.containingFile
-    val psiDocumentManager = PsiDocumentManager.getInstance(file.project)
-    psiDocumentManager.doPostponedOperationsAndUnblockDocument(psiDocumentManager.getDocument(file)!!)
-    val codeStyleManager = CodeStyleManager.getInstance(file.project)
-    codeStyleManager.adjustLineIndent(file, result.textRange)
 
     return result
 }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndReturnTransformation.kt

@@ -38,10 +38,6 @@ class FindAndReturnTransformation(
 
     override val commentSavingRange = PsiChildRange(loop.unwrapIfLabeled(), endReturn)
 
-    private val commentRestoringRange = commentSavingRange.withoutFirstStatement()
-
-    override fun commentRestoringRange(convertLoopResult: KtExpression) = commentRestoringRange
-
     override val presentation: String
         get() = generator.functionName + (if (filter != null) "{}" else "()")
 
@@ -55,7 +51,7 @@ class FindAndReturnTransformation(
         return generator.generate(chainedCallGenerator, filter)
     }
 
-    override fun convertLoop(resultCallChain: KtExpression): KtExpression {
+    override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
         endReturn.returnedExpression!!.replace(resultCallChain)
         loop.deleteWithLabels()
         return endReturn