Formatter, NJ2K: fix indent for comments inside function literals

#KT-4194 Fixed
#KT-31881 Fixed
#KT-34673 Fixed
#KT-35152 Fixed

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.idea.formatter.KotlinSpacingBuilder.CustomSpacingBui
 import org.jetbrains.kotlin.idea.util.requireNode
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.children
 import org.jetbrains.kotlin.psi.psiUtil.isObjectLiteral
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.psi.psiUtil.textRangeWithoutComments
@@ -90,6 +91,14 @@ fun createSpacingBuilder(settings: CodeStyleSettings, builderUtil: KotlinSpacing
 
             inPosition(right = BLOCK_COMMENT).spacing(commentSpacing(0))
             inPosition(right = EOL_COMMENT).spacing(commentSpacing(1))
+            inPosition(parent = FUNCTION_LITERAL, right = BLOCK).customRule { _, _, right ->
+                when (right.node?.children()?.firstOrNull()?.elementType) {
+                    BLOCK_COMMENT -> commentSpacing(0)
+                    EOL_COMMENT -> commentSpacing(1)
+                    else -> null
+                }
+            }
+
             inPosition(leftSet = DECLARATIONS, rightSet = DECLARATIONS).customRule(fun(
                 _: ASTBlock,
                 _: ASTBlock,
@@ -97,9 +106,12 @@ fun createSpacingBuilder(settings: CodeStyleSettings, builderUtil: KotlinSpacing
             ): Spacing? {
                 val node = right.node ?: return null
                 val elementStart = node.startOfDeclaration() ?: return null
-                return if (StringUtil.containsLineBreak(node.text.subSequence(0, elementStart.startOffset - node.startOffset).trimStart())) {
+                return if (StringUtil.containsLineBreak(
+                        node.text.subSequence(0, elementStart.startOffset - node.startOffset).trimStart()
+                    )
+                )
                     createSpacing(0, minLineFeeds = 2)
-                } else
+                else
                     null
             })
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt

@@ -22,6 +22,7 @@ import com.intellij.psi.*
 import com.intellij.psi.codeStyle.CodeStyleManager
 import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.psiUtil.*
 import java.util.*
@@ -347,6 +348,11 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
                     if (commentTreeElement.spaceAfter.isNotEmpty()) {
                         parent.addBefore(psiFactory.createWhiteSpace(commentTreeElement.spaceAfter), anchorElement)
                     }
+
+                    val functionLiteral = restored.parent as? KtFunctionLiteral
+                    if (functionLiteral != null && commentTreeElement.spaceBefore.isNotEmpty()) {
+                        functionLiteral.addBefore(psiFactory.createWhiteSpace(commentTreeElement.spaceBefore), restored)
+                    }
                 }
             } else {
                 restored = putAbandonedCommentsAfter.parent.addBefore(comment, putAbandonedCommentsAfter) as PsiComment

idea/testData/formatter/CommentInFunctionLiteral.after.kt

@@ -6,4 +6,67 @@ fun foo() {
         // Some comment
         it
     }
-}
+
+    test() {
+//        val a = 42
+    }
+
+    test() {
+/*
+        val a = 42
+*/
+    }
+
+    test() {
+        /*
+            val a = 42
+        */
+    }
+
+    test() {
+//        val a = 42
+        val b = 44
+    }
+}
+
+val s = Shadow { -> // wdwd
+    val a = 42
+}
+
+val s2 = Shadow { // wdwd
+    val a = 42
+}
+
+val s3 = Shadow(fun() {  // wdwd
+    val a = 42
+})
+
+val s4 = Shadow(fun() { /* s */
+    val a = 42
+})
+
+val s5 = Shadow { ->
+    // wdwd
+    val a = 42
+}
+
+val s6 = Shadow {
+    // wdwd
+    val a = 42
+}
+
+val s7 = Shadow(fun() {
+    // wdwd
+    val a = 42
+})
+
+val s8 = Shadow(fun() {
+    // wdwd
+    val a = 42
+})
+
+val s9 = Shadow { ->   /* s */
+    val a = 42
+}
+
+class Shadow(callback: () -> Unit)

idea/testData/formatter/CommentInFunctionLiteral.kt

@@ -6,4 +6,67 @@ fun foo() {
     // Some comment
         it
     }
-}
+
+    test() {
+//        val a = 42
+    }
+
+    test() {
+/*
+        val a = 42
+*/
+    }
+
+    test() {
+    /*
+        val a = 42
+    */
+    }
+
+    test() {
+//        val a = 42
+        val b = 44
+    }
+}
+
+val s = Shadow { -> // wdwd
+    val a = 42
+}
+
+val s2 = Shadow { // wdwd
+    val a = 42
+}
+
+val s3 = Shadow(fun() {  // wdwd
+    val a = 42
+})
+
+val s4 = Shadow(fun() { /* s */
+    val a = 42
+})
+
+val s5 = Shadow { ->
+    // wdwd
+    val a = 42
+}
+
+val s6 = Shadow {
+    // wdwd
+    val a = 42
+}
+
+val s7 = Shadow(fun() {
+    // wdwd
+    val a = 42
+})
+
+val s8 = Shadow(fun() {
+    // wdwd
+    val a = 42
+})
+
+val s9 = Shadow { ->   /* s */
+    val a = 42
+}
+
+class Shadow(callback: () -> Unit)

idea/testData/intentions/loopToCallChain/firstOrNull/ifAssign_preserveComments.kt.after

@@ -4,8 +4,7 @@
 fun foo(list: List<String>) {
     // string should be non-empty
     // save it into result
-    val result: String? = list.firstOrNull {
-        // search for first non-empty string in the list
+    val result: String? = list.firstOrNull {  // search for first non-empty string in the list
         it.length > 0
     }
 }