Fix outstanding cases with importing

api/swiftpoet.api

@@ -107,7 +107,9 @@ public final class io/outfoxx/swiftpoet/DeclaredTypeName : io/outfoxx/swiftpoet/
 	public final fun getSimpleName ()Ljava/lang/String;
 	public final fun getSimpleNames ()Ljava/util/List;
 	public final fun nestedType (Ljava/lang/String;Z)Lio/outfoxx/swiftpoet/DeclaredTypeName;
+	public final fun nestedType (Ljava/util/List;Z)Lio/outfoxx/swiftpoet/DeclaredTypeName;
 	public static synthetic fun nestedType$default (Lio/outfoxx/swiftpoet/DeclaredTypeName;Ljava/lang/String;ZILjava/lang/Object;)Lio/outfoxx/swiftpoet/DeclaredTypeName;
+	public static synthetic fun nestedType$default (Lio/outfoxx/swiftpoet/DeclaredTypeName;Ljava/util/List;ZILjava/lang/Object;)Lio/outfoxx/swiftpoet/DeclaredTypeName;
 	public final fun peerType (Ljava/lang/String;Z)Lio/outfoxx/swiftpoet/DeclaredTypeName;
 	public static synthetic fun peerType$default (Lio/outfoxx/swiftpoet/DeclaredTypeName;Ljava/lang/String;ZILjava/lang/Object;)Lio/outfoxx/swiftpoet/DeclaredTypeName;
 	public static final fun qualifiedTypeName (Ljava/lang/String;)Lio/outfoxx/swiftpoet/DeclaredTypeName;

src/main/java/io/outfoxx/swiftpoet/CodeWriter.kt

@@ -26,12 +26,13 @@ private val NO_MODULE = String()
  * Converts a [FileSpec] to a string suitable to both human- and swiftc-consumption. This honors
  * imports, indentation, and variable names.
  */
-internal class CodeWriter constructor(
+internal class CodeWriter(
   out: Appendable,
   private val indent: String = DEFAULT_INDENT,
   internal val importedTypes: Map<String, DeclaredTypeName> = emptyMap(),
   private val importedModules: Set<String> = emptySet()
 ) : Closeable {
+
   private val out = LineWrapper(out, indent, 100)
   private var indentLevel = 0
 
@@ -40,6 +41,7 @@ internal class CodeWriter constructor(
   private var moduleStack = mutableListOf(NO_MODULE)
   private val typeSpecStack = mutableListOf<AnyTypeSpec>()
   private val importableTypes = mutableMapOf<String, DeclaredTypeName>()
+  private val referencedTypes = mutableMapOf<String, DeclaredTypeName>()
   private var trailingNewline = false
 
   /**
@@ -205,8 +207,7 @@ internal class CodeWriter constructor(
     var a = 0
     val partIterator = codeBlock.formatParts.listIterator()
     while (partIterator.hasNext()) {
-      val part = partIterator.next()
-      when (part) {
+      when (val part = partIterator.next()) {
         "%L" -> emitLiteral(codeBlock.args[a++], isConstantContext)
 
         "%N" -> emit(codeBlock.args[a++] as String)
@@ -275,27 +276,29 @@ internal class CodeWriter constructor(
    * names visible due to inheritance.
    */
   fun lookupName(typeName: DeclaredTypeName): String {
+
+    // Track all referenced type names, Swift needs to import the module for each type
+    referencedTypes[typeName.canonicalName] = typeName
+
     // Find the shortest suffix of typeName that resolves to typeName. This uses both local type
     // names (so `Entry` in `Map` refers to `Map.Entry`). Also uses imports.
-    var nameResolved = false
-    var c: DeclaredTypeName? = typeName
-    while (c != null) {
-      val simpleName = c.simpleName
-      val resolved = resolve(simpleName)
-      nameResolved = resolved != null
-
-      if (resolved == c.unwrapOptional()) {
-        val suffixOffset = c.simpleNames.size - 1
-        return typeName.simpleNames.subList(suffixOffset, typeName.simpleNames.size)
-          .joinToString(".")
+    var currentTypeName: DeclaredTypeName? = typeName
+    val currentNestedSimpleNames = mutableListOf<String>()
+    while (currentTypeName != null) {
+      val simpleName = currentTypeName.simpleName
+      val resolved = resolve(simpleName)?.nestedType(currentNestedSimpleNames)
+
+      if (resolved == typeName.unwrapOptional()) {
+        // If the type is the same as the type we're resolving for, we must use at least that name.
+        if (currentNestedSimpleNames.isEmpty()) {
+          return simpleName
+        }
+        // Otherwise, we need to use all the nested names that didn't match
+        return currentNestedSimpleNames.joinToString(".")
       }
 
-      c = c.enclosingTypeName()
-    }
-
-    // If the name resolved but wasn't a match, we're stuck with the fully qualified name.
-    if (nameResolved) {
-      return typeName.canonicalName
+      currentNestedSimpleNames.add(0, simpleName)
+      currentTypeName = currentTypeName.enclosingTypeName()
     }
 
     // If the type is in the same module, we're done.
@@ -304,7 +307,7 @@ internal class CodeWriter constructor(
     }
 
     // If the type is in a manually imported module and doesn't clash, use an unqualified type
-    if (importedModules.contains(typeName.moduleName) && !importableTypes.containsKey(typeName.simpleName)) {
+    if (importedModules.contains(typeName.moduleName) && !importedTypes.containsKey(typeName.simpleName)) {
       return typeName.simpleName
     }
 
@@ -313,7 +316,7 @@ internal class CodeWriter constructor(
       importableType(typeName)
     }
 
-    return typeName.canonicalName
+    return resolveImport(typeName)
   }
 
   private fun importableType(typeName: DeclaredTypeName) {
@@ -322,10 +325,7 @@ internal class CodeWriter constructor(
     }
     val topLevelTypeName = typeName.topLevelTypeName()
     val simpleName = topLevelTypeName.simpleName
-    val replaced = importableTypes.put(simpleName, topLevelTypeName)
-    if (replaced != null) {
-      importableTypes[simpleName] = replaced // On collision, prefer the first inserted.
-    }
+    importableTypes.putIfAbsent(simpleName, topLevelTypeName)
   }
 
   /**
@@ -337,12 +337,12 @@ internal class CodeWriter constructor(
       val typeSpec = typeSpecStack[i]
       if (typeSpec is ExternalTypeSpec) {
         if (typeSpec.name == simpleName) {
-          return stackTypeName(i, simpleName)
+          return stackTypeName(i)
         }
       }
       for (visibleChild in typeSpec.typeSpecs) {
         if (visibleChild.name == simpleName) {
-          return stackTypeName(i, simpleName)
+          return stackTypeName(i).nestedType(simpleName)
         }
       }
     }
@@ -352,21 +352,29 @@ internal class CodeWriter constructor(
       return DeclaredTypeName(moduleStack.last(), simpleName)
     }
 
-    // Match an imported type.
-    val importedType = importedTypes[simpleName]
-    if (importedType != null) return importedType
-
     // No match.
     return null
   }
 
+  /**
+   * Looks up `typeName` in the imports and returns the shortest name possible for that type name.
+   */
+  private fun resolveImport(typeName: DeclaredTypeName): String {
+    val topLevelTypeName = typeName.topLevelTypeName()
+    return if (importedTypes.values.any { it == topLevelTypeName }) {
+      typeName.simpleNames.joinToString(".")
+    } else {
+      typeName.canonicalName
+    }
+  }
+
   /** Returns the type named `simpleName` when nested in the type at `stackDepth`.  */
-  private fun stackTypeName(stackDepth: Int, simpleName: String): DeclaredTypeName {
+  private fun stackTypeName(stackDepth: Int): DeclaredTypeName {
     var typeName = DeclaredTypeName(moduleStack.last(), typeSpecStack[0].name)
     for (i in 1..stackDepth) {
       typeName = typeName.nestedType(typeSpecStack[i].name)
     }
-    return typeName.nestedType(simpleName)
+    return typeName
   }
 
   /**
@@ -422,54 +430,28 @@ internal class CodeWriter constructor(
   }
 
   /**
-   * Returns the modules that should have been imported for this code.
+   * Returns the non-colliding importable types and module names for all referenced types.
    */
-  private fun suggestedImports(): Map<String, DeclaredTypeName> {
-    return importableTypes
+  private fun generateImports(): Pair<Map<String, DeclaredTypeName>, Set<String>> {
+    return importableTypes to referencedTypes.values.map { it.moduleName }.toSet()
   }
 
   companion object {
+
     /**
-     * Makes a pass to collect imports by executing [emitStep], and returns an instance of
-     * [CodeWriter] pre-initialized with collected imports.
+     * Collect imports by executing [emitStep], and returns the non-colliding imported types
+     * and referenced modules.
      */
-    fun withCollectedImports(
-      out: Appendable,
+    fun collectImports(
       indent: String,
       emitStep: (importsCollector: CodeWriter) -> Unit,
-    ): CodeWriter {
-      // First pass: emit the entire class, just to collect the types we'll need to import.
-      val suggestedImports = CodeWriter(
-        NullAppendable,
-        indent,
-      ).use { importsCollector ->
-        emitStep(importsCollector)
-
-        val generatedImports = mutableMapOf<String, String>()
-        importsCollector.suggestedImports()
-          .generateImports(
-            generatedImports,
-            canonicalName = DeclaredTypeName::canonicalName,
-          )
-      }
+    ): Pair<Map<String, DeclaredTypeName>, Set<String>> =
+      CodeWriter(NullAppendable, indent)
+        .use { importsCollector ->
 
-      return CodeWriter(
-        out,
-        indent,
-        suggestedImports,
-      )
-    }
+          emitStep(importsCollector)
 
-    private fun <T> Map<String, T>.generateImports(
-      generatedImports: MutableMap<String, String>,
-      canonicalName: T.() -> String,
-    ): Map<String, T> {
-      return flatMap { (simpleName, qualifiedName) ->
-        listOf(simpleName to qualifiedName).also {
-          val canonicalName = qualifiedName.canonicalName()
-          generatedImports[canonicalName] = canonicalName
+          importsCollector.generateImports()
         }
-      }.toMap()
-    }
   }
 }

src/main/java/io/outfoxx/swiftpoet/DeclaredTypeName.kt

@@ -68,6 +68,13 @@ class DeclaredTypeName internal constructor(
   fun nestedType(name: String, alwaysQualify: Boolean = this.alwaysQualify) =
     DeclaredTypeName(names + name, alwaysQualify)
 
+  /**
+   * Returns a new [DeclaredTypeName] instance for the specified `names` as nested inside this
+   * type.
+   */
+  fun nestedType(names: List<String>, alwaysQualify: Boolean = this.alwaysQualify) =
+    DeclaredTypeName(this.names + names, alwaysQualify)
+
   /**
    * Returns a type that shares the same enclosing package or type. If this type is enclosed by
    * another type, this is equivalent to `enclosingTypeName().nestedType(name)`. Otherwise

src/main/java/io/outfoxx/swiftpoet/FileSpec.kt

@@ -47,12 +47,15 @@ class FileSpec private constructor(
 
   @Throws(IOException::class)
   fun writeTo(out: Appendable) {
-    val codeWriter = CodeWriter.withCollectedImports(
-      out = out,
-      indent = indent,
-      emitStep = { importsCollector -> emit(importsCollector) },
-    )
-    codeWriter.use(::emit)
+
+    val (importedTypes, referencedModules) =
+      CodeWriter.collectImports(
+        indent = indent,
+        emitStep = { importsCollector -> emit(importsCollector) },
+      )
+
+    val codeWriter = CodeWriter(out, indent = indent, importedTypes = importedTypes)
+    emit(codeWriter, referencedModules = referencedModules)
   }
 
   /** Writes this to `directory` as UTF-8 using the standard directory structure.  */
@@ -70,19 +73,20 @@ class FileSpec private constructor(
   @Throws(IOException::class)
   fun writeTo(directory: File) = writeTo(directory.toPath())
 
-  private fun emit(codeWriter: CodeWriter) {
+  private fun emit(codeWriter: CodeWriter, referencedModules: Set<String> = setOf()) {
     if (comment.isNotEmpty()) {
       codeWriter.emitComment(comment)
     }
 
     codeWriter.pushModule(moduleName)
 
-    val importedTypeImports = codeWriter.importedTypes.map { ImportSpec.builder(it.value.moduleName).build() }
-    val allImports = moduleImports + importedTypeImports
-    val imports = allImports.filter { it.name != "Swift" }
+    val implicitModuleImports = referencedModules.map { ImportSpec.builder(it).build() }
+    val allModuleImports = moduleImports + implicitModuleImports
+    val nonImportedModules = NON_IMPORTED_MODULES + moduleName
+    val moduleImports = allModuleImports.filterNot { nonImportedModules.contains(it.name) }
 
-    if (imports.isNotEmpty()) {
-      for (import in imports.toSortedSet()) {
+    if (moduleImports.isNotEmpty()) {
+      for (import in moduleImports.toSortedSet()) {
         import.emit(codeWriter)
         codeWriter.emit("\n")
       }
@@ -173,6 +177,9 @@ class FileSpec private constructor(
   }
 
   companion object {
+
+    private val NON_IMPORTED_MODULES = setOf("Swift")
+
     @JvmStatic fun get(moduleName: String, typeSpec: AnyTypeSpec): FileSpec {
       return builder(moduleName, typeSpec.name).addType(typeSpec).build()
     }