Add NameAllocator API to control keyword pre-allocation

square · Jan 12, 2024 · 95006e7 · 95006e7
1 parent 98a345b
commit 95006e7
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 3 deletions.
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 pre-allocated and suffixed with
+   * underscores to avoid using them 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 pre-allocate 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)