Add support for context-receivers (#1233)

* Add API for context receivers

* Use existing opt-in annotation and make context-receivers not nullable

* Remove collection overloads

* Revert unwanted code style changes

* Add wrongly remove @JvmOverloads annotation

* Add code generator and tests

* Apply requested changes
- Use proper experimental annotation
- Remove not needed overloads
- Move contextReceivers parameter to the end

* Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt

Co-authored-by: Zac Sweers <pandanomic@gmail.com>

* Apply requested changes
- Add more tests
- Add new-line after context()

* Update test for suggestion

* Apply suggestions from code review

Co-authored-by: Egor Andreevich <github@egorand.dev>

* Apply requested changes

* Rename contextReceiver to contextReceivers

* Fix compiler error

* Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt

Co-authored-by: Egor Andreevich <github@egorand.dev>

* Fix compiler errors in tests

* Run spotless

* Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt

Co-authored-by: Zac Sweers <pandanomic@gmail.com>

* Apply requested changes

* Don't emit context receiver spacing if there are no context receivers

* Apply suggestions from code review

Co-authored-by: Egor Andreevich <github@egorand.dev>

Co-authored-by: Zac Sweers <pandanomic@gmail.com>
Co-authored-by: Egor Andreevich <github@egorand.dev>

6 files changed, 205 insertions, 9 deletions

diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/CodeWriter.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/CodeWriter.kt
index c208ad2e..9b68d949 100644
--- a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/CodeWriter.kt
+++ b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/CodeWriter.kt
@@ -168,6 +168,19 @@ internal class CodeWriter constructor(
   }
 
   /**
+   * Emits the `context` block for [contextReceivers].
+   */
+  fun emitContextReceivers(contextReceivers: List<TypeName>, suffix: String = "") {
+    if (contextReceivers.isNotEmpty()) {
+      val receivers = contextReceivers
+        .map { CodeBlock.of("%T", it) }
+        .joinToCode(prefix = "context(", suffix = ")")
+      emitCode(receivers)
+      emit(suffix)
+    }
+  }
+
+  /**
    * Emit type variables with their bounds. If a type variable has more than a single bound - call
    * [emitWhereBlock] with same input to produce an additional `where` block.
    *
diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/ExperimentalKotlinPoetApi.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/ExperimentalKotlinPoetApi.kt
new file mode 100644
index 00000000..2e58556d
--- /dev/null
+++ b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/ExperimentalKotlinPoetApi.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2022 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.squareup.kotlinpoet
+
+import kotlin.annotation.AnnotationTarget.CLASS
+import kotlin.annotation.AnnotationTarget.FUNCTION
+import kotlin.annotation.AnnotationTarget.PROPERTY
+import kotlin.annotation.AnnotationTarget.TYPEALIAS
+
+/**
+ * Indicates that a given API is experimental and subject to change.
+ */
+@RequiresOptIn
+@Retention(AnnotationRetention.BINARY)
+@Target(CLASS, FUNCTION, PROPERTY, TYPEALIAS)
+public annotation class ExperimentalKotlinPoetApi
diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
index 671bcfde..4d5f2428 100644
--- a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
+++ b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
@@ -32,6 +32,7 @@ import kotlin.DeprecationLevel.WARNING
 import kotlin.reflect.KClass
 
 /** A generated function declaration. */
+@OptIn(ExperimentalKotlinPoetApi::class)
 public class FunSpec private constructor(
   builder: Builder,
   private val tagMap: TagMap = builder.buildTagMap(),
@@ -45,6 +46,9 @@ public class FunSpec private constructor(
   public val modifiers: Set<KModifier> = builder.modifiers.toImmutableSet()
   public val typeVariables: List<TypeVariableName> = builder.typeVariables.toImmutableList()
   public val receiverType: TypeName? = builder.receiverType
+
+  @ExperimentalKotlinPoetApi
+  public val contextReceiverTypes: List<TypeName> = builder.contextReceiverTypes.toImmutableList()
   public val returnType: TypeName? = builder.returnType
   public val parameters: List<ParameterSpec> = builder.parameters.toImmutableList()
   public val delegateConstructor: String? = builder.delegateConstructor
@@ -84,6 +88,7 @@ public class FunSpec private constructor(
       codeWriter.emitKdoc(kdoc.ensureEndsWithNewLine())
     }
     codeWriter.emitAnnotations(annotations, false)
+    codeWriter.emitContextReceivers(contextReceiverTypes, suffix = "\n")
     codeWriter.emitModifiers(modifiers, implicitModifiers)
 
     if (!isConstructor && !name.isAccessor) {
@@ -285,6 +290,7 @@ public class FunSpec private constructor(
     internal var returnKdoc = CodeBlock.EMPTY
     internal var receiverKdoc = CodeBlock.EMPTY
     internal var receiverType: TypeName? = null
+    internal val contextReceiverTypes: MutableList<TypeName> = mutableListOf()
     internal var returnType: TypeName? = null
     internal var delegateConstructor: String? = null
     internal var delegateConstructorArguments = listOf<CodeBlock>()
@@ -359,6 +365,15 @@ public class FunSpec private constructor(
       typeVariables += typeVariable
     }
 
+    @ExperimentalKotlinPoetApi
+    public fun contextReceivers(receiverTypes: Iterable<TypeName>): Builder = apply {
+      check(!name.isConstructor) { "constructors cannot have context receivers" }
+      contextReceiverTypes += receiverTypes
+    }
+
+    @ExperimentalKotlinPoetApi
+    public fun contextReceivers(vararg receiverType: TypeName): Builder = contextReceivers(receiverType.toList())
+
     @JvmOverloads public fun receiver(
       receiverType: TypeName,
       kdoc: CodeBlock = CodeBlock.EMPTY
diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt
index 3f8514a3..2baa8d24 100644
--- a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt
+++ b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt
@@ -17,8 +17,11 @@ package com.squareup.kotlinpoet
 
 import kotlin.reflect.KClass
 
+@OptIn(ExperimentalKotlinPoetApi::class)
 public class LambdaTypeName private constructor(
   public val receiver: TypeName? = null,
+  @property:ExperimentalKotlinPoetApi
+  public val contextReceivers: List<TypeName> = emptyList(),
   parameters: List<ParameterSpec> = emptyList(),
   public val returnType: TypeName = UNIT,
   nullable: Boolean = false,
@@ -50,10 +53,11 @@ public class LambdaTypeName private constructor(
     suspending: Boolean = this.isSuspending,
     tags: Map<KClass<*>, Any> = this.tags.toMap()
   ): LambdaTypeName {
-    return LambdaTypeName(receiver, parameters, returnType, nullable, suspending, annotations, tags)
+    return LambdaTypeName(receiver, contextReceivers, parameters, returnType, nullable, suspending, annotations, tags)
   }
 
   override fun emit(out: CodeWriter): CodeWriter {
+    out.emitContextReceivers(contextReceivers, suffix = "·")
     if (isNullable) {
       out.emit("(")
     }
@@ -81,11 +85,19 @@ public class LambdaTypeName private constructor(
 
   public companion object {
     /** Returns a lambda type with `returnType` and parameters listed in `parameters`. */
+    @ExperimentalKotlinPoetApi @JvmStatic public fun get(
+      receiver: TypeName? = null,
+      parameters: List<ParameterSpec> = emptyList(),
+      returnType: TypeName,
+      contextReceivers: List<TypeName> = emptyList()
+    ): LambdaTypeName = LambdaTypeName(receiver, contextReceivers, parameters, returnType)
+
+    /** Returns a lambda type with `returnType` and parameters listed in `parameters`. */
     @JvmStatic public fun get(
       receiver: TypeName? = null,
       parameters: List<ParameterSpec> = emptyList(),
       returnType: TypeName
-    ): LambdaTypeName = LambdaTypeName(receiver, parameters, returnType)
+    ): LambdaTypeName = LambdaTypeName(receiver, emptyList(), parameters, returnType)
 
     /** Returns a lambda type with `returnType` and parameters listed in `parameters`. */
     @JvmStatic public fun get(
@@ -95,6 +107,7 @@ public class LambdaTypeName private constructor(
     ): LambdaTypeName {
       return LambdaTypeName(
         receiver,
+        emptyList(),
         parameters.toList().map { ParameterSpec.unnamed(it) },
         returnType
       )
@@ -105,6 +118,6 @@ public class LambdaTypeName private constructor(
       receiver: TypeName? = null,
       vararg parameters: ParameterSpec = emptyArray(),
       returnType: TypeName
-    ): LambdaTypeName = LambdaTypeName(receiver, parameters.toList(), returnType)
+    ): LambdaTypeName = LambdaTypeName(receiver, emptyList(), parameters.toList(), returnType)
   }
 }
diff --git a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
index 13857a6f..8804e7b3 100644
--- a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
+++ b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
@@ -33,6 +33,7 @@ import javax.lang.model.util.Types
 import kotlin.test.BeforeTest
 import kotlin.test.Test
 
+@OptIn(ExperimentalKotlinPoetApi::class)
 class FunSpecTest {
   @Rule @JvmField val compilation = CompilationRule()
 
@@ -69,6 +70,8 @@ class FunSpecTest {
 
   internal interface ExtendsOthers : Callable<Int>, Comparable<Long>
 
+  annotation class TestAnnotation
+
   abstract class InvalidOverrideMethods {
     fun finalMethod() {
     }
@@ -429,6 +432,76 @@ class FunSpecTest {
     )
   }
 
+  @Test fun functionWithContextReceiver() {
+    val stringType = STRING
+    val funSpec = FunSpec.builder("foo")
+      .contextReceivers(stringType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String)
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun functionWithMultipleContextReceivers() {
+    val stringType = STRING
+    val intType = INT
+    val booleanType = BOOLEAN
+    val funSpec = FunSpec.builder("foo")
+      .contextReceivers(stringType, intType, booleanType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String, kotlin.Int, kotlin.Boolean)
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun functionWithGenericContextReceiver() {
+    val genericType = TypeVariableName("T")
+    val funSpec = FunSpec.builder("foo")
+      .addTypeVariable(genericType)
+      .contextReceivers(genericType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(T)
+      |public fun <T> foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun functionWithAnnotatedContextReceiver() {
+    val genericType = STRING.copy(annotations = listOf(AnnotationSpec.get(TestAnnotation())))
+    val funSpec = FunSpec.builder("foo")
+      .contextReceivers(genericType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(@com.squareup.kotlinpoet.FunSpecTest.TestAnnotation kotlin.String)
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun constructorWithContextReceiver() {
+    assertThrows<IllegalStateException> {
+      FunSpec.constructorBuilder()
+        .contextReceivers(STRING)
+    }.hasMessageThat().isEqualTo("constructors cannot have context receivers")
+  }
+
   @Test fun functionParamSingleLambdaParam() {
     val unitType = UNIT
     val booleanType = BOOLEAN
diff --git a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/LambdaTypeNameTest.kt b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/LambdaTypeNameTest.kt
index e86d1f2b..dd44d719 100644
--- a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/LambdaTypeNameTest.kt
+++ b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/LambdaTypeNameTest.kt
@@ -20,6 +20,7 @@ import com.squareup.kotlinpoet.KModifier.VARARG
 import javax.annotation.Nullable
 import kotlin.test.Test
 
+@OptIn(ExperimentalKotlinPoetApi::class)
 class LambdaTypeNameTest {
 
   @Retention(AnnotationRetention.RUNTIME)
@@ -30,9 +31,9 @@ class LambdaTypeNameTest {
 
   @Test fun receiverWithoutAnnotationHasNoParens() {
     val typeName = LambdaTypeName.get(
-      Int::class.asClassName(),
-      listOf(),
-      Unit::class.asTypeName()
+      receiver = Int::class.asClassName(),
+      parameters = listOf(),
+      returnType = Unit::class.asTypeName()
     )
     assertThat(typeName.toString()).isEqualTo("kotlin.Int.() -> kotlin.Unit")
   }
@@ -40,17 +41,69 @@ class LambdaTypeNameTest {
   @Test fun receiverWithAnnotationHasParens() {
     val annotation = IsAnnotated::class.java.getAnnotation(HasSomeAnnotation::class.java)
     val typeName = LambdaTypeName.get(
-      Int::class.asClassName().copy(
+      receiver = Int::class.asClassName().copy(
         annotations = listOf(AnnotationSpec.get(annotation, includeDefaultValues = true))
       ),
-      listOf(),
-      Unit::class.asTypeName()
+      parameters = listOf(),
+      returnType = Unit::class.asTypeName()
     )
     assertThat(typeName.toString()).isEqualTo(
       "(@com.squareup.kotlinpoet.LambdaTypeNameTest.HasSomeAnnotation kotlin.Int).() -> kotlin.Unit"
     )
   }
 
+  @Test fun contextReceiver() {
+    val typeName = LambdaTypeName.get(
+      receiver = Int::class.asTypeName(),
+      parameters = listOf(),
+      returnType = Unit::class.asTypeName(),
+      contextReceivers = listOf(STRING)
+    )
+    assertThat(typeName.toString()).isEqualTo(
+      "context(kotlin.String) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
+  @Test fun functionWithMultipleContextReceivers() {
+    val typeName = LambdaTypeName.get(
+      Int::class.asTypeName(),
+      listOf(),
+      Unit::class.asTypeName(),
+      listOf(STRING, BOOLEAN)
+    )
+    assertThat(typeName.toString()).isEqualTo(
+      "context(kotlin.String, kotlin.Boolean) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
+  @Test fun functionWithGenericContextReceiver() {
+    val genericType = TypeVariableName("T")
+    val typeName = LambdaTypeName.get(
+      Int::class.asTypeName(),
+      listOf(),
+      Unit::class.asTypeName(),
+      listOf(genericType)
+    )
+
+    assertThat(typeName.toString()).isEqualTo(
+      "context(T) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
+  @Test fun functionWithAnnotatedContextReceiver() {
+    val annotatedType = STRING.copy(annotations = listOf(AnnotationSpec.get(FunSpecTest.TestAnnotation())))
+    val typeName = LambdaTypeName.get(
+      Int::class.asTypeName(),
+      listOf(),
+      Unit::class.asTypeName(),
+      listOf(annotatedType)
+    )
+
+    assertThat(typeName.toString()).isEqualTo(
+      "context(@com.squareup.kotlinpoet.FunSpecTest.TestAnnotation kotlin.String) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
   @Test fun paramsWithAnnotationsForbidden() {
     assertThrows<IllegalArgumentException> {
       LambdaTypeName.get(