Support context receivers (#400)

Summary: Resolves #397, #314 and #374. Pull Request resolved: #400 Reviewed By: davidtorosyan Differential Revision: D48169986 Pulled By: hick209 fbshipit-source-id: df4ffe4d939635b3db481ba478d52bffa4c78ec1
facebook · Sep 13, 2023 · 17d80bf · 17d80bf
1 parent 46ac818
commit 17d80bf
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 1 deletion.
diff --git a/core/pom.xml b/core/pom.xml
@@ -16,7 +16,7 @@
 
     <properties>
         <dokka.version>0.10.1</dokka.version>
-        <kotlin.version>1.6.10</kotlin.version>
+        <kotlin.version>1.6.20-M1</kotlin.version>
         <kotlin.compiler.incremental>true</kotlin.compiler.incremental>
         <kotlin.compiler.jvmTarget>1.8</kotlin.compiler.jvmTarget>
         <main.class>com.facebook.ktfmt.cli.Main</main.class>

diff --git a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -52,6 +52,7 @@ import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
 import org.jetbrains.kotlin.psi.KtConstantExpression
 import org.jetbrains.kotlin.psi.KtConstructorDelegationCall
 import org.jetbrains.kotlin.psi.KtContainerNode
+import org.jetbrains.kotlin.psi.KtContextReceiverList
 import org.jetbrains.kotlin.psi.KtContinueExpression
 import org.jetbrains.kotlin.psi.KtDelegatedSuperTypeEntry
 import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
@@ -94,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
@@ -124,6 +126,7 @@ import org.jetbrains.kotlin.psi.psiUtil.children
 import org.jetbrains.kotlin.psi.psiUtil.getPrevSiblingIgnoringWhitespace
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.psi.psiUtil.startsWithComment
+import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 
 /** An AST visitor that builds a stream of {@link Op}s to format. */
 class KotlinInputAstVisitor(
@@ -162,6 +165,7 @@ class KotlinInputAstVisitor(
     builder.sync(function)
     builder.block(ZERO) {
       visitFunctionLikeExpression(
+          function.getStubOrPsiChild(KtStubElementTypes.CONTEXT_RECEIVER_LIST),
           function.modifierList,
           "fun",
           function.typeParameterList,
@@ -282,6 +286,7 @@ class KotlinInputAstVisitor(
    *   list of supertypes.
    */
   private fun visitFunctionLikeExpression(
+      contextReceiverList: KtContextReceiverList?,
       modifierList: KtModifierList?,
       keyword: String,
       typeParameters: KtTypeParameterList?,
@@ -294,6 +299,9 @@ class KotlinInputAstVisitor(
       typeOrDelegationCall: KtElement?,
   ) {
     builder.block(ZERO) {
+      if (contextReceiverList != null) {
+        visitContextReceiverList(contextReceiverList)
+      }
       if (modifierList != null) {
         visitModifierList(modifierList)
       }
@@ -1372,6 +1380,7 @@ class KotlinInputAstVisitor(
 
           builder.block(ZERO) {
             visitFunctionLikeExpression(
+                null,
                 accessor.modifierList,
                 accessor.namePlaceholder.text,
                 null,
@@ -1461,8 +1470,12 @@ class KotlinInputAstVisitor(
 
   override fun visitClassOrObject(classOrObject: KtClassOrObject) {
     builder.sync(classOrObject)
+    val contextReceiverList = classOrObject.getStubOrPsiChild(KtStubElementTypes.CONTEXT_RECEIVER_LIST)
     val modifierList = classOrObject.modifierList
     builder.block(ZERO) {
+      if (contextReceiverList != null) {
+        visitContextReceiverList(contextReceiverList)
+      }
       if (modifierList != null) {
         visitModifierList(modifierList)
       }
@@ -1533,6 +1546,7 @@ class KotlinInputAstVisitor(
 
     val delegationCall = constructor.getDelegationCall()
     visitFunctionLikeExpression(
+        constructor.getStubOrPsiChild(KtStubElementTypes.CONTEXT_RECEIVER_LIST),
         constructor.modifierList,
         "constructor",
         null,
@@ -1638,6 +1652,20 @@ class KotlinInputAstVisitor(
     builder.forcedBreak()
   }
 
+  /** Example `context(Logger, Raise<Error>)` */
+  override fun visitContextReceiverList(contextReceiverList: KtContextReceiverList) {
+    builder.sync(contextReceiverList)
+    builder.token("context")
+    visitEachCommaSeparated(
+            contextReceiverList.contextReceivers(),
+            prefix = "(",
+            postfix = ")",
+            breakAfterPrefix = false,
+            breakBeforePostfix = false
+    )
+    builder.forcedBreak()
+  }
+
   /** For example `@Magic private final` */
   override fun visitModifierList(list: KtModifierList) {
     builder.sync(list)
@@ -2427,6 +2455,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()) {
@@ -2436,13 +2465,16 @@ 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)
       }
       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
@@ -6764,6 +6764,48 @@ class FormatterTest {
     assertThatFormatting(code).isEqualTo(expected)
   }
 
+  @Test
+  fun `context receivers`() {
+    val code =
+        """
+      |context(Something)
+      |
+      |class A {
+      |  context(
+      |  // Test comment.
+      |  Logger, Raise<Error>)
+      |
+      |  @SomeAnnotation
+      |
+      |  fun doNothing() {}
+      |
+      |  context(SomethingElse)
+      |
+      |  private class NestedClass {}
+      |}
+      |"""
+            .trimMargin()
+
+    val expected =
+        """
+      |context(Something)
+      |class A {
+      |  context(
+      |  // Test comment.
+      |  Logger,
+      |  Raise<Error>)
+      |  @SomeAnnotation
+      |  fun doNothing() {}
+      |
+      |  context(SomethingElse)
+      |  private class NestedClass {}
+      |}
+      |"""
+            .trimMargin()
+
+    assertThatFormatting(code).isEqualTo(expected)
+  }
+
   companion object {
     /** Triple quotes, useful to use within triple-quoted strings. */
     private const val TQ = "\"\"\""

diff --git a/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt b/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt
@@ -110,4 +110,28 @@ class TokenizerTest {
         .containsExactly(0, -1, 1, 2, 3, -1, 4, -1, 5, 6, 7)
         .inOrder()
   }
+
+  @Test
+  fun `Context receivers are parsed correctly`() {
+    val code = """
+      |context(Something)
+      |class A {
+      |  context(
+      |  // Test comment.
+      |  Logger, Raise<Error>)
+      |  fun test() {}
+      |}
+      |""".trimMargin().trimMargin()
+
+    val file = Parser.parse(code)
+    val tokenizer = Tokenizer(code, file)
+    file.accept(tokenizer)
+
+    assertThat(tokenizer.toks.map { it.originalText })
+            .containsExactly("context", "(", "Something", ")", "\n", "class", " ", "A", " ", "{", "\n", "  ", "context", "(", "\n", "  ", "// Test comment.", "\n", "  ", "Logger", ",", " ", "Raise", "<", "Error", ">", ")", "\n", "  ", "fun", " ", "test", "(", ")", " ", "{", "}", "\n", "}")
+            .inOrder()
+    assertThat(tokenizer.toks.map { it.index })
+            .containsExactly(0, 1, 2, 3, -1, 4, -1, 5, -1, 6, -1, -1, 7, 8, -1, -1, 9, -1, -1, 10, 11, -1, 12, 13, 14, 15, 16, -1, -1, 17, -1, 18, 19, 20, -1, 21, 22, -1, 23)
+            .inOrder()
+  }
 }