Refactoring: work with ASTNodes in functional style

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.psi.psiUtil
 
 import com.intellij.injected.editor.VirtualFileWindow
+import com.intellij.lang.ASTNode
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.*
 import com.intellij.psi.search.PsiSearchScopeUtil
@@ -389,3 +390,12 @@ fun KtModifierList.hasExpectModifier() = hasModifier(KtTokens.HEADER_KEYWORD) ||
 fun KtModifierListOwner.hasActualModifier() = hasModifier(KtTokens.IMPL_KEYWORD) || hasModifier(KtTokens.ACTUAL_KEYWORD)
 fun KtModifierList.hasActualModifier() = hasModifier(KtTokens.IMPL_KEYWORD) || hasModifier(KtTokens.ACTUAL_KEYWORD)
 
+fun ASTNode.children() = generateSequence(firstChildNode) { node -> node.treeNext }
+
+fun ASTNode.siblings(forward: Boolean = true): Sequence<ASTNode> {
+    if (forward) {
+        return generateSequence(treeNext) { it.treeNext }
+    } else {
+        return generateSequence(treePrev) { it.treePrev }
+    }
+}

idea/formatter/src/org/jetbrains/kotlin/idea/formatter/KotlinCommonBlock.kt

@@ -33,7 +33,8 @@ import org.jetbrains.kotlin.kdoc.parser.KDocElementTypes
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.psi.*
-import java.util.*
+import org.jetbrains.kotlin.psi.psiUtil.children
+import org.jetbrains.kotlin.psi.psiUtil.siblings
 
 private val QUALIFIED_OPERATION = TokenSet.create(DOT, SAFE_ACCESS)
 private val KDOC_COMMENT_INDENT = 1
@@ -104,8 +105,7 @@ abstract class KotlinCommonBlock(
                         nodeSubBlocks.subList(operationBlockIndex, nodeSubBlocks.size),
                         null, indent, null, spacingBuilder) { createSyntheticSpacingNodeBlock(it) }
 
-                nodeSubBlocks = ArrayList<Block>(nodeSubBlocks.subList(0, operationBlockIndex))
-                nodeSubBlocks.add(operationSyntheticBlock)
+                nodeSubBlocks = nodeSubBlocks.subList(0, operationBlockIndex) + operationSyntheticBlock
             }
         }
 
@@ -295,31 +295,14 @@ abstract class KotlinCommonBlock(
         return createBlock(child, alignmentStrategy, createChildIndent(child), childWrap, settings, spacingBuilder)
     }
 
-    private fun buildSubBlocks(): ArrayList<Block> {
-        val blocks = ArrayList<Block>()
-
+    private fun buildSubBlocks(): List<Block> {
         val childrenAlignmentStrategy = getChildrenAlignmentStrategy()
         val wrappingStrategy = getWrappingStrategy()
 
-        var child: ASTNode? = node.firstChildNode
-        while (child != null) {
-            val childType = child.elementType
-
-            if (child.textRange.length == 0) {
-                child = child.treeNext
-                continue
-            }
-
-            if (childType === TokenType.WHITE_SPACE) {
-                child = child.treeNext
-                continue
-            }
-
-            blocks.add(buildSubBlock(child, childrenAlignmentStrategy, wrappingStrategy))
-            child = child.treeNext
-        }
-
-        return blocks
+        return node.children()
+                .filter { it.textRange.length > 0 && it.elementType != TokenType.WHITE_SPACE }
+                .map { buildSubBlock(it, childrenAlignmentStrategy, wrappingStrategy ) }
+                .toMutableList()
     }
 
     private fun getWrappingStrategy(): WrappingStrategy {
@@ -605,24 +588,14 @@ private fun getAlignmentForChildInParenthesis(
     }
 }
 
-private fun getPrevWithoutWhitespace(pNode: ASTNode?): ASTNode? {
-    var node = pNode
-    node = node!!.treePrev
-    while (node != null && node.elementType === TokenType.WHITE_SPACE) {
-        node = node.treePrev
-    }
-
-    return node
+private fun getPrevWithoutWhitespace(pNode: ASTNode): ASTNode? {
+    return pNode.siblings(forward = false).firstOrNull { it.elementType != TokenType.WHITE_SPACE }
 }
 
-private fun getPrevWithoutWhitespaceAndComments(pNode: ASTNode?): ASTNode? {
-    var node = pNode
-    node = node!!.treePrev
-    while (node != null && (node.elementType === TokenType.WHITE_SPACE || COMMENTS.contains(node.elementType))) {
-        node = node.treePrev
+private fun getPrevWithoutWhitespaceAndComments(pNode: ASTNode): ASTNode? {
+    return pNode.siblings(forward = false).firstOrNull {
+        it.elementType != TokenType.WHITE_SPACE && it.elementType !in COMMENTS
     }
-
-    return node
 }
 
 private fun getWrappingStrategyForItemList(wrapType: Int, itemType: IElementType, wrapFirstElement: Boolean = false): WrappingStrategy {

idea/formatter/src/org/jetbrains/kotlin/idea/formatter/SyntheticKotlinBlock.kt

@@ -16,19 +16,13 @@
 
 package org.jetbrains.kotlin.idea.formatter
 
-import com.intellij.formatting.Block
-import com.intellij.formatting.Alignment
-import com.intellij.formatting.Indent
-import com.intellij.formatting.Wrap
-import com.intellij.openapi.util.TextRange
-import com.intellij.formatting.ChildAttributes
-import com.intellij.formatting.Spacing
+import com.intellij.formatting.*
 import com.intellij.lang.ASTNode
-import com.intellij.formatting.ASTBlock
+import com.intellij.openapi.util.TextRange
 
 class SyntheticKotlinBlock(
         private val node: ASTNode,
-        private val subBlocks: MutableList<Block>,
+        private val subBlocks: List<Block>,
         private val alignment: Alignment?,
         private val indent: Indent?,
         private val wrap: Wrap?,