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 {
+    val result: String? = list.firstOrNull {  // search for first non-empty string in the list
-        // search for first non-empty string in the list
         it.length > 0
     }
 }

nj2k/src/org/jetbrains/kotlin/nj2k/FormattingCollector.kt

@@ -5,14 +5,13 @@
 
 package org.jetbrains.kotlin.nj2k
 
-import com.intellij.psi.PsiComment
+import com.intellij.openapi.util.text.StringUtil
-import com.intellij.psi.PsiElement
+import com.intellij.psi.*
-import com.intellij.psi.PsiKeyword
-import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.javadoc.PsiDocComment
 import org.jetbrains.kotlin.idea.j2k.IdeaDocCommentConverter
 import org.jetbrains.kotlin.nj2k.tree.JKComment
 import org.jetbrains.kotlin.nj2k.tree.JKFormattingOwner
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class FormattingCollector {
     private val commentCache = mutableMapOf<PsiElement, JKComment>()
@@ -48,13 +47,21 @@ class FormattingCollector {
                     this
                 )
             )
-            is PsiComment -> JKComment(text)
+            is PsiComment -> JKComment(text, indent())
             else -> null
         } ?: return null
         commentCache[this] = token
         return token
     }
 
+    private fun PsiComment.indent(): String? = takeIf { parent is PsiCodeBlock }?.prevSibling?.safeAs<PsiWhiteSpace>()?.let { space ->
+        val text = space.text
+        if (space.prevSibling is PsiStatement)
+            text.indexOfFirst(StringUtil::isLineBreak).takeIf { it != -1 }?.let { text.substring(it + 1) } ?: text
+        else
+            text
+    }
+
     private fun Sequence<PsiElement>.toComments(): List<JKComment> =
         takeWhile { it is PsiComment || it is PsiWhiteSpace || it.text == ";" }
             .mapNotNull { it.asComment() }

nj2k/src/org/jetbrains/kotlin/nj2k/printing/JKCommentPrinter.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.nj2k.printing
 
+import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.nj2k.tree.JKComment
 import org.jetbrains.kotlin.nj2k.tree.JKDeclaration
 import org.jetbrains.kotlin.nj2k.tree.JKFormattingOwner
@@ -32,7 +33,8 @@ internal class JKCommentPrinter(val printer: JKPrinter) {
         for (comment in this@createText) {
             if (comment.shouldBeDropped()) continue
             val text = comment.createText() ?: continue
-            if (needNewLine) appendln() else append(' ')
+            if (needNewLine && comment.indent?.let { StringUtil.containsLineBreak(it) } != true) appendln()
+            append(comment.indent ?: ' ')
             append(text)
             needNewLine = text.startsWith("//") || '\n' in text
         }

nj2k/src/org/jetbrains/kotlin/nj2k/tree/formatting.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.nj2k.tree
 import org.jetbrains.kotlin.utils.SmartList
 
 
-class JKComment(val text: String) {
+class JKComment(val text: String, val indent: String? = null) {
     val isSingleline
         get() = text.startsWith("//")
 }

nj2k/testData/newJ2k/annotations/jetbrainsNotNull.kt

@@ -4,7 +4,8 @@ package test
 
 class Test(str: String) {
     var myStr = "String2"
-    fun sout(str: String) { // UNNECESSARY_NOT_NULL_ASSERTION heuristic does not work any more, instead we can skip generating !! altogether
+    fun sout(str: String) {
+        // UNNECESSARY_NOT_NULL_ASSERTION heuristic does not work any more, instead we can skip generating !! altogether
         println(str)
     }
 

nj2k/testData/newJ2k/comments/commentInsideCall.kt

@@ -5,7 +5,7 @@ object Test {
     fun main(args: Array<String>) {
         println() // Comment
         foo() // Comment1
-// Comment2
+                // Comment2
                 .indexOf("s")
     }
 

nj2k/testData/newJ2k/comments/comments.kt

@@ -6,16 +6,19 @@ package foo
 // we need ArrayList
 // let's declare a class:
 internal class A /* just a sample name*/ : Runnable /* let's implement Runnable */ {
-    fun foo /* again a sample name */(p: Int /* parameter p */, c: Char /* parameter c */) { // let's print something:
+    fun foo /* again a sample name */(p: Int /* parameter p */, c: Char /* parameter c */) {
+        // let's print something:
         println("1") // print 1
         println("2") // print 2
         println("3") // print 3
+
         // end of printing
         if (p > 0) { // do this only when p > 0
-// we print 4 and return
+            // we print 4 and return
             println("3")
             return  // do not continue
         }
+
         // some code to be added
     }
 } // end of class A

nj2k/testData/newJ2k/comments/comments2.kt

@@ -3,8 +3,9 @@
 package foo
 
 internal class A {
-    fun  /* nothing to return */foo( /* no parameters at all */) { // let declare a variable
+    fun  /* nothing to return */foo( /* no parameters at all */) {
-// with 2 comments before
+        // let declare a variable
+        // with 2 comments before
         val   /*int*/a /* it's a */ = 2 /* it's 2 */ + 1 /* it's 1 */ // variable a declared
     } // end of foo
 

nj2k/testData/newJ2k/comments/comments3.java

@@ -0,0 +1,21 @@
+public class Foo {
+    public void foo() {
+        System.out.println();
+        // comment
+        // comment 2
+        // comment 3
+    }
+
+    public void setState(Integer state) {
+        /*some comment 1/*/
+        /*some comment 2/*/
+        /*some comment 3/*/
+        /*some comment 4/*/
+    }
+
+    public void setState2(Integer state) { /*some comment 1/*/
+        /*some comment 2/*/
+        /*some comment 3/*/
+        /*some comment 4/*/
+    }
+}

nj2k/testData/newJ2k/comments/comments3.kt

@@ -0,0 +1,21 @@
+class Foo {
+    fun foo() {
+        println()
+        // comment
+        // comment 2
+        // comment 3
+    }
+
+    fun setState(state: Int?) {
+        /*some comment 1/ */
+        /*some comment 2/ */
+        /*some comment 3/ */
+        /*some comment 4/ */
+    }
+
+    fun setState2(state: Int?) { /*some comment 1/ */
+        /*some comment 2/ */
+        /*some comment 3/ */
+        /*some comment 4/ */
+    }
+}

nj2k/testData/newJ2k/comments/fieldWithEndOfLineComment2.java

@@ -0,0 +1,13 @@
+public class Foo {
+    private Integer someInt;
+    public void setState(Integer state) {
+        //some comment 1
+        someInt = state;
+        //some comment 2
+        if (state == 2)
+            System.out.println("2");
+    }
+    public Integer getState() {
+        return someInt;
+    }
+}

nj2k/testData/newJ2k/comments/fieldWithEndOfLineComment2.kt

@@ -0,0 +1,12 @@
+class Foo {
+    var state: Int? = null
+        private set
+
+    fun setState(state: Int) {
+        //some comment 1
+        this.state = state
+        //some comment 2
+        if (state == 2) println("2")
+    }
+
+}

nj2k/testData/newJ2k/formatting/nonStaticMembersWithComments.kt

@@ -1,10 +1,10 @@
 internal class F {
     //c1
-/*c2*/
+    /*c2*/
     fun f1() {}
 
     //c3
-//c4
+    //c4
     fun f2() {}
     var i = 0
     fun f3() {} //c5

nj2k/testData/newJ2k/formatting/staticAndNonStaticMembersWithComments.kt

@@ -1,13 +1,13 @@
 internal class F {
     //c3
-//c4
+    //c4
     fun f2() {}
     fun f3() {}
     fun f4() {}
 
     companion object {
         //c1
-/*c2*/
+        /*c2*/
         fun f1() {}
         var i = 0
 

nj2k/testData/newJ2k/formatting/staticMembersWithComments.kt

@@ -1,10 +1,10 @@
 internal object F {
     //c1
-/*c2*/
+    /*c2*/
     fun f1() {}
 
     //c3
-//c4
+    //c4
     fun f2() {}
     var i = 0
     fun f3() {} //c5

nj2k/testData/newJ2k/function/externalKtFunctionalInterface.kt

@@ -8,7 +8,8 @@ internal class Test {
         return x.toDouble()
     }
 
-    fun nya(): Double { //TODO explicitlly call apply here
+    fun nya(): Double {
+        //TODO explicitlly call apply here
         return foo(1, FunctionalI<Int, Double> { x: Int -> this.toDouble(x) })
     }
 }

nj2k/testData/newJ2k/issues/comments.kt

@@ -15,7 +15,8 @@ internal class C {
     /**
      * This is a function doc comment.
      */
-    fun foo() { /* This is a function comment */
+    fun foo() {
+        /* This is a function comment */
     }
 
     //simple one line comment for function
@@ -25,29 +26,29 @@ internal class C {
     var j = 0
 
     //double c style
-//comment before function
+    //comment before function
     fun f2() {}
 
     //double c style
-//comment before field
+    //comment before field
     var k = 0
     //combination
     /** of
      */
-//
+    //
     /**
      * different
      */
-//comments
+    //comments
     fun f3() {}
     //combination
     /** of
      */
-//
+    //
     /**
      * different
      */
-//comments
+    //comments
     var l = 0
 
     /*two*/ /*comments*/ /*line*/

nj2k/testData/newJ2k/issues/kt-19348.kt

@@ -14,6 +14,6 @@ class TestMutltipleCtorsWithJavadoc
         this.y = y
     }
     // ---
-// Constructors
+    // Constructors
-//
+    //
 }

nj2k/testData/newJ2k/tryWithResource/WithFinally.kt

@@ -10,7 +10,8 @@ class C {
                 val c = stream.read()
                 println(c)
             }
-        } finally { // dispose something else
+        } finally {
+            // dispose something else
         }
     }
 }

nj2k/tests/org/jetbrains/kotlin/nj2k/NewJavaToKotlinConverterSingleFileTestGenerated.java

@@ -1030,6 +1030,11 @@ public class NewJavaToKotlinConverterSingleFileTestGenerated extends AbstractNew
             runTest("nj2k/testData/newJ2k/comments/comments2.java");
         }
 
+        @TestMetadata("comments3.java")
+        public void testComments3() throws Exception {
+            runTest("nj2k/testData/newJ2k/comments/comments3.java");
+        }
+
         @TestMetadata("commentsForConstructors.java")
         public void testCommentsForConstructors() throws Exception {
             runTest("nj2k/testData/newJ2k/comments/commentsForConstructors.java");
@@ -1040,6 +1045,11 @@ public class NewJavaToKotlinConverterSingleFileTestGenerated extends AbstractNew
             runTest("nj2k/testData/newJ2k/comments/fieldWithEndOfLineComment.java");
         }
 
+        @TestMetadata("fieldWithEndOfLineComment2.java")
+        public void testFieldWithEndOfLineComment2() throws Exception {
+            runTest("nj2k/testData/newJ2k/comments/fieldWithEndOfLineComment2.java");
+        }
+
         @TestMetadata("fieldsInitializedFromParams.java")
         public void testFieldsInitializedFromParams() throws Exception {
             runTest("nj2k/testData/newJ2k/comments/fieldsInitializedFromParams.java");