Merge pull request #705 from ronshapiro/CodeWriter_resolve

Remove n^2 algorithm in CodeWriter.resolve() by precomputing all of the nested simple names of a TypeSpec

2 files changed, 8 insertions, 4 deletions

diff --git a/src/main/java/com/squareup/javapoet/CodeWriter.java b/src/main/java/com/squareup/javapoet/CodeWriter.java
index ac8178f..b2b088b 100644
--- a/src/main/java/com/squareup/javapoet/CodeWriter.java
+++ b/src/main/java/com/squareup/javapoet/CodeWriter.java
@@ -420,10 +420,8 @@ final class CodeWriter {
     // Match a child of the current (potentially nested) class.
     for (int i = typeSpecStack.size() - 1; i >= 0; i--) {
       TypeSpec typeSpec = typeSpecStack.get(i);
-      for (TypeSpec visibleChild : typeSpec.typeSpecs) {
-        if (Objects.equals(visibleChild.name, simpleName)) {
-          return stackClassName(i, simpleName);
-        }
+      if (typeSpec.nestedTypesSimpleNames.contains(simpleName)) {
+        return stackClassName(i, simpleName);
       }
     }
 
diff --git a/src/main/java/com/squareup/javapoet/TypeSpec.java b/src/main/java/com/squareup/javapoet/TypeSpec.java
index c315579..3346763 100644
--- a/src/main/java/com/squareup/javapoet/TypeSpec.java
+++ b/src/main/java/com/squareup/javapoet/TypeSpec.java
@@ -21,6 +21,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumSet;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -53,6 +54,7 @@ public final class TypeSpec {
   public final CodeBlock initializerBlock;
   public final List<MethodSpec> methodSpecs;
   public final List<TypeSpec> typeSpecs;
+  final Set<String> nestedTypesSimpleNames;
   public final List<Element> originatingElements;
 
   private TypeSpec(Builder builder) {
@@ -72,11 +74,14 @@ public final class TypeSpec {
     this.methodSpecs = Util.immutableList(builder.methodSpecs);
     this.typeSpecs = Util.immutableList(builder.typeSpecs);
 
+    nestedTypesSimpleNames = new HashSet<>(builder.typeSpecs.size());
     List<Element> originatingElementsMutable = new ArrayList<>();
     originatingElementsMutable.addAll(builder.originatingElements);
     for (TypeSpec typeSpec : builder.typeSpecs) {
+      nestedTypesSimpleNames.add(typeSpec.name);
       originatingElementsMutable.addAll(typeSpec.originatingElements);
     }
+
     this.originatingElements = Util.immutableList(originatingElementsMutable);
   }
 
@@ -102,6 +107,7 @@ public final class TypeSpec {
     this.methodSpecs = Collections.emptyList();
     this.typeSpecs = Collections.emptyList();
     this.originatingElements = Collections.emptyList();
+    this.nestedTypesSimpleNames = Collections.emptySet();
   }
 
   public boolean hasModifier(Modifier modifier) {