Fix shortening of nested names in generated hierarchies

.gitignore

@@ -138,3 +138,4 @@ gradle-app.setting
 
 run
 /.idea
+.mise.toml

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

@@ -69,7 +69,7 @@ internal class CodeWriter(
   }
 
   fun popModule() = apply {
-    val lastModuleName = this.moduleStack.removeLast()
+    val lastModuleName = this.moduleStack.removeLastOrNull()
     require(lastModuleName !== NO_MODULE) { "module stack imbalance" }
   }
 
@@ -299,6 +299,17 @@ internal class CodeWriter(
     var currentTypeName: DeclaredTypeName? = typeName
     val currentNestedSimpleNames = mutableListOf<String>()
     while (currentTypeName != null) {
+
+      // Check if the type stack is in an external type that matches what we're looking
+      // for. This cannot be done using `resolve` because we do not know the contents of
+      // external types, and assume that any type given is nested in an external type.
+      typeSpecStack.filterIsInstance<ExternalTypeSpec>().forEach {
+        val external = listOf(it.name)
+        if ((external + currentNestedSimpleNames) == typeName.simpleNames) {
+          return currentNestedSimpleNames.ifEmpty { external }.joinToString(".")
+        }
+      }
+
       val simpleName = currentTypeName.simpleName
       val resolved = resolve(simpleName)?.nestedType(currentNestedSimpleNames)
 
@@ -307,8 +318,9 @@ internal class CodeWriter(
         if (currentNestedSimpleNames.isEmpty()) {
           return simpleName
         }
-        // Otherwise, we need to use all the nested names that didn't match
-        return currentNestedSimpleNames.joinToString(".")
+        // Otherwise, we need to use the current matching name as the resolved base,
+        // and all the nested names that didn't match.
+        return (listOf(simpleName) + currentNestedSimpleNames).joinToString(".")
       }
 
       currentNestedSimpleNames.add(0, simpleName)
@@ -353,11 +365,6 @@ internal class CodeWriter(
     // Match a child of the current (potentially nested) type.
     for (i in typeSpecStack.indices.reversed()) {
       val typeSpec = typeSpecStack[i]
-      if (typeSpec is ExternalTypeSpec) {
-        if (typeSpec.name == simpleName) {
-          return stackTypeName(i)
-        }
-      }
       for (visibleChild in typeSpec.typeSpecs) {
         if (visibleChild.name == simpleName) {
           return stackTypeName(i).nestedType(simpleName)

src/test/java/io/outfoxx/swiftpoet/test/FileSpecTests.kt

@@ -124,15 +124,110 @@ class FileSpecTests {
   }
 
   @Test
-  @DisplayName("Generates correct imports for extension type names")
-  fun testImportsForSameExtensionTypes() {
+  @DisplayName("Generates correct names for generated nested type names")
+  fun testGeneratesCorrectNamesForGeneratedNestedTypeNames() {
+
+    val typeSpec =
+      TypeSpec.classBuilder("Root")
+        .addType(
+          TypeSpec.classBuilder("Node")
+            .addType(
+              TypeSpec.classBuilder("Leaf")
+                .addType(
+                  TypeSpec.classBuilder("Iterator")
+                    .build()
+                )
+                .addFunction(
+                  FunctionSpec.builder("test")
+                    .addStatement("let iter = %T()", typeName(".Root.Node.Leaf.Iterator"))
+                    .build()
+                )
+                .build()
+            )
+            .addFunction(
+              FunctionSpec.builder("test")
+                .addStatement("let leaf = %T()", typeName(".Root.Node.Leaf"))
+                .addStatement("let leafIter = %T()", typeName(".Root.Node.Leaf.Iterator"))
+                .build()
+            )
+            .build()
+        )
+        .addFunction(
+          FunctionSpec.builder("test")
+            .addStatement("let node = %T()", typeName(".Root.Node"))
+            .addStatement("let leaf = %T()", typeName(".Root.Node.Leaf"))
+            .addStatement("let leafIter = %T()", typeName(".Root.Node.Leaf.Iterator"))
+            .build()
+        )
+        .build()
+
+    val testFile = FileSpec.builder("", "Root")
+      .addType(typeSpec)
+      .build()
+
+    val out = StringWriter()
+    testFile.writeTo(out)
+
+    assertThat(
+      out.toString(),
+      equalTo(
+        """
+            class Root {
+
+              func test() {
+                let node = Node()
+                let leaf = Node.Leaf()
+                let leafIter = Node.Leaf.Iterator()
+              }
+
+              class Node {
+
+                func test() {
+                  let leaf = Leaf()
+                  let leafIter = Leaf.Iterator()
+                }
+
+                class Leaf {
+
+                  func test() {
+                    let iter = Iterator()
+                  }
+
+                  class Iterator {
+                  }
+
+                }
+
+              }
+
+            }
+
+        """.trimIndent()
+      )
+    )
+  }
+
+  @Test
+  @DisplayName("Generates correct imports for extensions on declared type names")
+  fun testImportsForSameExtensionDeclaredTypes() {
 
     val obs = typeName("RxSwift.Observable")
     val obsElement = typeName("RxSwift.Observable.Element")
     val obsElementSub = typeName("RxSwift.Observable.Element.SubSequence")
 
     val extension =
       ExtensionSpec.builder(obsElement.enclosingTypeName()!!)
+        .addType(
+          TypeSpec.classBuilder("Sub")
+            .addFunction(
+              FunctionSpec.builder("test")
+                .addStatement("let obs = %T()", obs)
+                .addStatement("let obsElement = %T()", obsElement)
+                .addStatement("let obsElementSub = %T()", obsElementSub)
+                .build()
+            )
+            .build()
+        )
         .addFunction(
           FunctionSpec.builder("test")
             .returns(obs)
@@ -165,6 +260,105 @@ class FileSpecTests {
 
             extension Observable {
             
+              class Sub {
+
+                func test() {
+                  let obs = Observable()
+                  let obsElement = Element()
+                  let obsElementSub = Element.SubSequence()
+                }
+
+              }
+
+              func test() -> Observable {
+              }
+            
+              func test2() -> Element {
+              }
+            
+              func test3() -> Element.SubSequence {
+              }
+            
+            }
+
+        """.trimIndent()
+      )
+    )
+  }
+
+  @Test
+  @DisplayName("Generates correct imports for extensions on type specs")
+  fun testImportsForSameExtensionTypeSpecs() {
+
+    val typeSpec =
+      TypeSpec.classBuilder("Observable")
+        .addType(
+          TypeSpec.classBuilder("Element")
+            .addType(
+              TypeSpec.classBuilder("SubSequence")
+                .build()
+            )
+            .build()
+        )
+        .build()
+
+    val obs = typeName(".Observable")
+    val obsElement = typeName(".Observable.Element")
+    val obsElementSub = typeName(".Observable.Element.SubSequence")
+
+    val extension =
+      ExtensionSpec.builder(typeSpec)
+        .addType(
+          TypeSpec.classBuilder("Sub")
+            .addFunction(
+              FunctionSpec.builder("test")
+                .addStatement("let obs = %T()", obs)
+                .addStatement("let obsElement = %T()", obsElement)
+                .addStatement("let obsElementSub = %T()", obsElementSub)
+                .build()
+            )
+            .build()
+        )
+        .addFunction(
+          FunctionSpec.builder("test")
+            .returns(obs)
+            .build()
+        )
+        .addFunction(
+          FunctionSpec.builder("test2")
+            .returns(obsElement)
+            .build()
+        )
+        .addFunction(
+          FunctionSpec.builder("test3")
+            .returns(obsElementSub)
+            .build()
+        )
+        .build()
+
+    val testFile = FileSpec.builder("Test", "Test")
+      .addExtension(extension)
+      .build()
+
+    val out = StringWriter()
+    testFile.writeTo(out)
+
+    assertThat(
+      out.toString(),
+      equalTo(
+        """
+            extension Observable {
+            
+              class Sub {
+
+                func test() {
+                  let obs = Observable()
+                  let obsElement = Element()
+                  let obsElementSub = Element.SubSequence()
+                }
+
+              }
+
               func test() -> Observable {
               }