Formatter, NJ2K: fix indent for comments inside function literals

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

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.psi.PsiElement
-import com.intellij.psi.PsiKeyword
-import com.intellij.psi.PsiWhiteSpace
+import com.intellij.openapi.util.text.StringUtil
+import com.intellij.psi.*
 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
-// with 2 comments before
+    fun  /* nothing to return */foo( /* no parameters at all */) {
+        // 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");