Support generating a single import for overloaded MemberNames (#1909)

* Support generating a single import for overloaded MemberNames

MemberName imports were modeled after TypeName imports, where an
import can only represent a single fully-qualified name. However
with MemberNames, a single import can represent multiple
overloaded members that share a fully-qualified name. While
overloaded members can be represented with a single MemberName, we
distinguish between extension and non-extension members, and their
MemberNames are not equal. We'll currently try to generate import
aliases in this case, which is incorrect, and actually impossible,
since they have the exact same fully-qualified name (see #1907).
This change handles this scenario and adds support for having a
single import generated for multiple MemberNames.

* Update changelog

docs/changelog.md

@@ -20,6 +20,7 @@ Change Log
 * Fix: Throw if primary constructor delegates to other constructors (#1859).
 * Fix: Aliased imports with nested class (#1876).
 * Fix: Check for error types in `KSType.toClassName()` (#1890).
+* Fix: Support generating a single import for overloaded `MemberName`s (#1909).
 
 ## Version 1.16.0
 

kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/CodeWriter.kt

@@ -59,7 +59,7 @@ internal class CodeWriter(
   private val indent: String = DEFAULT_INDENT,
   imports: Map<String, Import> = emptyMap(),
   private val importedTypes: Map<String, ClassName> = emptyMap(),
-  private val importedMembers: Map<String, MemberName> = emptyMap(),
+  private val importedMembers: Map<String, Set<MemberName>> = emptyMap(),
   columnLimit: Int = 100,
 ) : Closeable {
   private var out = LineWrapper(out, indent, columnLimit)
@@ -463,11 +463,11 @@ internal class CodeWriter(
   fun lookupName(memberName: MemberName): String {
     val simpleName = imports[memberName.canonicalName]?.alias ?: memberName.simpleName
     // Match an imported member.
-    val importedMember = importedMembers[simpleName]
-    val found = importedMember == memberName
+    val importedMembers = importedMembers[simpleName] ?: emptySet()
+    val found = memberName in importedMembers
     if (found && !isMethodNameUsedInCurrentContext(simpleName)) {
       return simpleName
-    } else if (importedMember != null && memberName.enclosingClassName != null) {
+    } else if (importedMembers.isNotEmpty() && memberName.enclosingClassName != null) {
       val enclosingClassName = lookupName(memberName.enclosingClassName)
       return "$enclosingClassName.$simpleName"
     } else if (found) {
@@ -717,14 +717,14 @@ internal class CodeWriter(
       )
       emitStep(importsCollector)
       val generatedImports = mutableMapOf<String, Import>()
-      val suggestedTypeImports = importsCollector.suggestedTypeImports()
+      val importedTypes = importsCollector.suggestedTypeImports()
         .generateImports(
           generatedImports,
           computeCanonicalName = ClassName::canonicalName,
           capitalizeAliases = true,
           referencedNames = importsCollector.referencedNames,
         )
-      val suggestedMemberImports = importsCollector.suggestedMemberImports()
+      val importedMembers = importsCollector.suggestedMemberImports()
         .generateImports(
           generatedImports,
           computeCanonicalName = MemberName::canonicalName,
@@ -737,8 +737,8 @@ internal class CodeWriter(
         out = out,
         indent = indent,
         imports = memberImports + generatedImports.filterKeys { it !in memberImports },
-        importedTypes = suggestedTypeImports,
-        importedMembers = suggestedMemberImports,
+        importedTypes = importedTypes.mapValues { it.value.single() },
+        importedMembers = importedMembers,
       )
     }
 
@@ -747,42 +747,49 @@ internal class CodeWriter(
       computeCanonicalName: T.() -> String,
       capitalizeAliases: Boolean,
       referencedNames: Set<String>,
-    ): Map<String, T> {
-      return flatMap { (simpleName, qualifiedNames) ->
-        if (qualifiedNames.size == 1 && simpleName !in referencedNames) {
-          listOf(simpleName to qualifiedNames.first()).also {
-            val canonicalName = qualifiedNames.first().computeCanonicalName()
-            generatedImports[canonicalName] = Import(canonicalName)
-          }
+    ): Map<String, Set<T>> {
+      val imported = mutableMapOf<String, Set<T>>()
+      forEach { (simpleName, qualifiedNames) ->
+        val canonicalNamesToQualifiedNames = qualifiedNames.associateBy { it.computeCanonicalName() }
+        if (canonicalNamesToQualifiedNames.size == 1 && simpleName !in referencedNames) {
+          val canonicalName = canonicalNamesToQualifiedNames.keys.single()
+          generatedImports[canonicalName] = Import(canonicalName)
+
+          // For types, qualifiedNames should consist of a single name, for which an import will be generated. For
+          // members, there can be more than one qualified name mapping to a single simple name, e.g. overloaded
+          // functions declared in the same package. In these cases, a single import will suffice for all of them.
+          imported[simpleName] = qualifiedNames
         } else {
-          generateImportAliases(simpleName, qualifiedNames, computeCanonicalName, capitalizeAliases)
+          generateImportAliases(simpleName, canonicalNamesToQualifiedNames, capitalizeAliases)
             .onEach { (alias, qualifiedName) ->
               val canonicalName = qualifiedName.computeCanonicalName()
               generatedImports[canonicalName] = Import(canonicalName, alias)
+
+              imported[alias] = setOf(qualifiedName)
             }
         }
-      }.toMap()
+      }
+      return imported
     }
 
     private fun <T> generateImportAliases(
       simpleName: String,
-      qualifiedNames: Set<T>,
-      canonicalName: T.() -> String,
+      canonicalNamesToQualifiedNames: Map<String, T>,
       capitalizeAliases: Boolean,
     ): List<Pair<String, T>> {
-      val canonicalNameSegments = qualifiedNames.associateWith { qualifiedName ->
-        qualifiedName.canonicalName().split('.')
+      val canonicalNameSegmentsToQualifiedNames = canonicalNamesToQualifiedNames.mapKeys { (canonicalName, _) ->
+        canonicalName.split('.')
           .dropLast(1) // Last segment of the canonical name is the simple name, drop it to avoid repetition.
           .filter { it != "Companion" }
           .map { it.replaceFirstChar(Char::uppercaseChar) }
       }
       val aliasNames = mutableMapOf<String, T>()
       var segmentsToUse = 0
       // Iterate until we have unique aliases for all names.
-      while (aliasNames.size != qualifiedNames.size) {
+      while (aliasNames.size != canonicalNamesToQualifiedNames.size) {
         segmentsToUse += 1
         aliasNames.clear()
-        for ((qualifiedName, segments) in canonicalNameSegments) {
+        for ((segments, qualifiedName) in canonicalNameSegmentsToQualifiedNames) {
           val aliasPrefix = segments.takeLast(min(segmentsToUse, segments.size))
             .joinToString(separator = "")
             .replaceFirstChar { if (!capitalizeAliases) it.lowercaseChar() else it }

kotlinpoet/src/commonTest/kotlin/com/squareup/kotlinpoet/MemberNameTest.kt

@@ -657,4 +657,29 @@ class MemberNameTest {
       """.trimIndent(),
     )
   }
+
+  // https://github.com/square/kotlinpoet/issues/1907
+  @Test fun `extension and non-extension MemberName clash`() {
+    val file = FileSpec.builder("com.squareup.tacos", "Tacos")
+      .addFunction(
+        FunSpec.builder("main")
+          .addStatement("println(%M(Taco()))", MemberName("com.squareup.wrappers", "wrap"))
+          .addStatement("println(Taco().%M())", MemberName("com.squareup.wrappers", "wrap", isExtension = true))
+          .build(),
+      )
+      .build()
+    assertThat(file.toString()).isEqualTo(
+      """
+      package com.squareup.tacos
+
+      import com.squareup.wrappers.wrap
+
+      public fun main() {
+        println(wrap(Taco()))
+        println(Taco().wrap())
+      }
+
+      """.trimIndent(),
+    )
+  }
 }