From 1675a2902c4e0ed15c26da9231b13ecd41c54819 Mon Sep 17 00:00:00 2001
From: Christopher-Marcel Esser <chris.esser.91@gmail.com>
Date: Sun, 30 Apr 2023 12:44:16 +0200
Subject: [PATCH] Remove blank lines after context receiver

---
 .../java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt  | 5 +++++
 .../src/test/java/com/facebook/ktfmt/format/FormatterTest.kt | 3 +++
 .../src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt | 5 +++--
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
index e7d31a06..c60578e5 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -95,6 +95,7 @@ import org.jetbrains.kotlin.psi.KtQualifiedExpression
 import org.jetbrains.kotlin.psi.KtReferenceExpression
 import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.KtScript
+import org.jetbrains.kotlin.psi.KtScriptInitializer
 import org.jetbrains.kotlin.psi.KtSecondaryConstructor
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
@@ -2432,6 +2433,7 @@ class KotlinInputAstVisitor(
   override fun visitScript(script: KtScript) {
     markForPartialFormat()
     var lastChildHadBlankLineBefore = false
+    var lastChildIsContextReceiver = false
     var first = true
     for (child in script.blockExpression.children) {
       if (child.text.isBlank()) {
@@ -2441,6 +2443,8 @@ class KotlinInputAstVisitor(
       val childGetsBlankLineBefore = child !is KtProperty
       if (first) {
         builder.blankLineWanted(OpsBuilder.BlankLineWanted.PRESERVE)
+      } else if (lastChildIsContextReceiver) {
+        builder.blankLineWanted(OpsBuilder.BlankLineWanted.NO)
       } else if (child !is PsiComment &&
           (childGetsBlankLineBefore || lastChildHadBlankLineBefore)) {
         builder.blankLineWanted(OpsBuilder.BlankLineWanted.YES)
@@ -2448,6 +2452,7 @@ class KotlinInputAstVisitor(
       visit(child)
       builder.guessToken(";")
       lastChildHadBlankLineBefore = childGetsBlankLineBefore
+      lastChildIsContextReceiver = child is KtScriptInitializer && child.firstChild?.firstChild?.firstChild?.text == "context"
       first = false
     }
     markForPartialFormat()
diff --git a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
index 3b6d10d1..a70cae7f 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -6635,6 +6635,8 @@ class FormatterTest {
   fun `context receivers`() {
     val code =
         """
+      |context(Something)
+      |
       |class A {
       |  context(Logger, Raise<Error>)
       |
@@ -6651,6 +6653,7 @@ class FormatterTest {
 
     val expected =
         """
+      |context(Something)
       |class A {
       |  context(Logger, Raise<Error>)
       |  @SomeAnnotation
diff --git a/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt b/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt
index b3cfb55d..06ffe510 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt
@@ -114,6 +114,7 @@ class TokenizerTest {
   @Test
   fun `Context receivers are parsed correctly`() {
     val code = """
+      |context(Something)
       |class A {
       |  context(Logger, Raise<Error>)
       |  fun test() {}
@@ -125,10 +126,10 @@ class TokenizerTest {
     file.accept(tokenizer)
 
     assertThat(tokenizer.toks.map { it.originalText })
-            .containsExactly("class", " ", "A", " ", "{", "\n", "  ", "context", "(", "Logger", ",", " ", "Raise", "<", "Error", ">", ")", "\n", "  ", "fun", " ", "test", "(", ")", " ", "{", "}", "\n", "}")
+            .containsExactly("context", "(", "Something", ")", "\n", "class", " ", "A", " ", "{", "\n", "  ", "context", "(", "Logger", ",", " ", "Raise", "<", "Error", ">", ")", "\n", "  ", "fun", " ", "test", "(", ")", " ", "{", "}", "\n", "}")
             .inOrder()
     assertThat(tokenizer.toks.map { it.index })
-            .containsExactly(0, -1, 1, -1, 2, -1, -1, 3, 4, 5, 6, -1, 7, 8, 9, 10, 11, -1, -1, 12, -1, 13, 14, 15, -1, 16, 17, -1, 18)
+            .containsExactly(0, 1, 2, 3, -1, 4, -1, 5, -1, 6, -1, -1, 7, 8, 9, 10, -1, 11, 12, 13, 14, 15, -1, -1, 16, -1, 17, 18, 19, -1, 20, 21, -1, 22)
             .inOrder()
   }
 }