From 848e6cf743e8327e28f64aa04cffa311a4d6eaa8 Mon Sep 17 00:00:00 2001
From: Egor Andreevich <egor@squareup.com>
Date: Fri, 12 Jan 2024 13:06:36 -0500
Subject: [PATCH] Add NameAllocator API to control keyword pre-allocation
 (#1803)

* Add NameAllocator API to control keyword pre-allocation

* Update API dump

* preAllocate to preallocate

* Clarify documentation
---
 kotlinpoet/api/kotlinpoet.api                 |  1 +
 .../com/squareup/kotlinpoet/NameAllocator.kt  | 36 +++++++++++++++++--
 .../kotlin/com/squareup/kotlinpoet/Util.kt    |  2 +-
 .../squareup/kotlinpoet/NameAllocatorTest.kt  |  6 ++++
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/kotlinpoet/api/kotlinpoet.api b/kotlinpoet/api/kotlinpoet.api
index 62b6afbecc..4e92873d47 100644
--- a/kotlinpoet/api/kotlinpoet.api
+++ b/kotlinpoet/api/kotlinpoet.api
@@ -561,6 +561,7 @@ public final class com/squareup/kotlinpoet/MemberName$Companion {
 
 public final class com/squareup/kotlinpoet/NameAllocator {
 	public fun <init> ()V
+	public fun <init> (Z)V
 	public final fun copy ()Lcom/squareup/kotlinpoet/NameAllocator;
 	public final fun get (Ljava/lang/Object;)Ljava/lang/String;
 	public final fun newName (Ljava/lang/String;)Ljava/lang/String;
diff --git a/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/NameAllocator.kt b/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/NameAllocator.kt
index cda1004266..a6c570b9e3 100644
--- a/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/NameAllocator.kt
+++ b/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/NameAllocator.kt
@@ -77,7 +77,39 @@ public class NameAllocator private constructor(
   private val allocatedNames: MutableSet<String>,
   private val tagToName: MutableMap<Any, String>,
 ) {
-  public constructor() : this(mutableSetOf(), mutableMapOf())
+  public constructor() : this(preallocateKeywords = true)
+
+  /**
+   * @param preallocateKeywords If true, all Kotlin keywords will be preallocated. Requested names which
+   * collide with keywords will be suffixed with underscores to avoid being used as identifiers:
+   *
+   * ```kotlin
+   * val nameAllocator = NameAllocator(preallocateKeywords = true)
+   * println(nameAllocator.newName("when")) // prints "when_"
+   * ```
+   *
+   * If false, keywords will not get any special treatment:
+   *
+   * ```kotlin
+   * val nameAllocator = NameAllocator(preallocateKeywords = false)
+   * println(nameAllocator.newName("when")) // prints "when"
+   * ```
+   *
+   * Note that you can use the `%N` placeholder when emitting a name produced by [NameAllocator] to
+   * ensure it's properly escaped for use as an identifier:
+   *
+   * ```kotlin
+   * val nameAllocator = NameAllocator(preallocateKeywords = false)
+   * println(CodeBlock.of("%N", nameAllocator.newName("when"))) // prints "`when`"
+   * ```
+   *
+   * The default behaviour of [NameAllocator] is to preallocate keywords - this is the behaviour you'll
+   * get when using the no-arg constructor.
+   */
+  public constructor(preallocateKeywords: Boolean) : this(
+    allocatedNames = if (preallocateKeywords) KEYWORDS.toMutableSet() else mutableSetOf(),
+    tagToName = mutableMapOf(),
+  )
 
   /**
    * Return a new name using `suggestion` that will not be a Java identifier or clash with other
@@ -89,7 +121,7 @@ public class NameAllocator private constructor(
     tag: Any = UUID.randomUUID().toString(),
   ): String {
     var result = toJavaIdentifier(suggestion)
-    while (result.isKeyword || !allocatedNames.add(result)) {
+    while (!allocatedNames.add(result)) {
       result += "_"
     }
 
diff --git a/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/Util.kt b/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/Util.kt
index 16a443dd18..c226f39707 100644
--- a/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/Util.kt
+++ b/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/Util.kt
@@ -169,7 +169,7 @@ private val IDENTIFIER_REGEX =
 internal val String.isIdentifier get() = IDENTIFIER_REGEX.matches(this)
 
 // https://kotlinlang.org/docs/reference/keyword-reference.html
-private val KEYWORDS = setOf(
+internal val KEYWORDS = setOf(
   // Hard keywords
   "as",
   "break",
diff --git a/kotlinpoet/src/commonTest/kotlin/com/squareup/kotlinpoet/NameAllocatorTest.kt b/kotlinpoet/src/commonTest/kotlin/com/squareup/kotlinpoet/NameAllocatorTest.kt
index 28c7d34d1b..a61877c0ee 100644
--- a/kotlinpoet/src/commonTest/kotlin/com/squareup/kotlinpoet/NameAllocatorTest.kt
+++ b/kotlinpoet/src/commonTest/kotlin/com/squareup/kotlinpoet/NameAllocatorTest.kt
@@ -70,6 +70,12 @@ class NameAllocatorTest {
     assertThat(nameAllocator[1]).isEqualTo("when_")
   }
 
+  @Test fun kotlinKeywordNotPreAllocated() {
+    val nameAllocator = NameAllocator(preallocateKeywords = false)
+    assertThat(nameAllocator.newName("when", 1)).isEqualTo("when")
+    assertThat(nameAllocator[1]).isEqualTo("when")
+  }
+
   @Test fun tagReuseForbidden() {
     val nameAllocator = NameAllocator()
     nameAllocator.newName("foo", 1)