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
square · Jan 12, 2024 · 848e6cf · 848e6cf
1 parent 291a300
commit 848e6cf
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)