Tighten up the code that handles emitting type annotations.

I found the previous code difficult to follow. Hopefully this is less so.

2 files changed, 155 insertions, 232 deletions

diff --git a/src/main/java/com/squareup/javapoet/ClassName.java b/src/main/java/com/squareup/javapoet/ClassName.java
index b99c6e0..e8af154 100644
--- a/src/main/java/com/squareup/javapoet/ClassName.java
+++ b/src/main/java/com/squareup/javapoet/ClassName.java
@@ -17,10 +17,9 @@ package com.squareup.javapoet;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import javax.lang.model.SourceVersion;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.PackageElement;
@@ -32,232 +31,109 @@ import static javax.lang.model.element.NestingKind.MEMBER;
 import static javax.lang.model.element.NestingKind.TOP_LEVEL;
 
 /** A fully-qualified class name for top-level and member classes. */
-public abstract class ClassName extends TypeName implements Comparable<ClassName> {
-  final String simpleName;
-  String canonicalName;
-
-  /** A fully-qualified class name for top-level classes. */
-  private static final class TopLevelClassName extends ClassName {
-    final String packageName;
-
-    private TopLevelClassName(String packageName, String simpleName) {
-      this(packageName, simpleName, new ArrayList<>());
-    }
-
-    private TopLevelClassName(
-        String packageName, String simpleName, List<AnnotationSpec> annotations) {
-      super(simpleName, annotations);
-      this.packageName = packageName == null ? "" : packageName;
-      this.canonicalName = isDefaultPackage(packageName)
-          ? simpleName : String.join(".", Arrays.asList(packageName, simpleName));
-      checkArgument(
-          isDefaultPackage(simpleName) || SourceVersion.isName(simpleName),
-          "part '%s' is keyword", simpleName);
-    }
-
-    @Override public TopLevelClassName annotated(List<AnnotationSpec> annotations) {
-      return new TopLevelClassName(packageName, simpleName, concatAnnotations(annotations));
-    }
-
-    @Override public TopLevelClassName withoutAnnotations() {
-      return new TopLevelClassName(packageName, simpleName);
-    }
-
-    public String packageName() {
-      return packageName;
-    }
-
-    @Override
-    public ClassName enclosingClassName() {
-      return null;
-    }
-
-    @Override
-    public TopLevelClassName topLevelClassName() {
-      return this;
-    }
-
-    @Override
-    public String reflectionName() {
-      return isDefaultPackage(packageName)
-          ? simpleName
-          : String.join(".", Arrays.asList(packageName, simpleName));
-    }
+public final class ClassName extends TypeName implements Comparable<ClassName> {
+  public static final ClassName OBJECT = ClassName.get(Object.class);
 
-    @Override
-    public List<String> simpleNames() {
-      return Arrays.asList(simpleName);
-    }
+  /** The package name of this class, or "" if this is in the default package. */
+  final String packageName;
 
-    @Override
-    public ClassName peerClass(String name) {
-      return new TopLevelClassName(packageName, name);
-    }
+  /** The enclosing class, or null if this is not enclosed in another class. */
+  final ClassName enclosingClassName;
 
-    @Override
-    ClassName prefixWithAtMostOneAnnotatedClass() {
-      return this;
-    }
+  /** This class name, like "Entry" for java.util.Map.Entry. */
+  final String simpleName;
 
-    @Override
-    boolean hasAnnotatedEnclosingClass() {
-      return false;
-    }
+  /** The full class name like "java.util.Map.Entry". */
+  final String canonicalName;
 
-    @Override
-    CodeWriter emitWithoutPrefix(CodeWriter out, ClassName unannotatedPrefix) {
-      return out;
-    }
+  private ClassName(String packageName, ClassName enclosingClassName, String simpleName) {
+    this(packageName, enclosingClassName, simpleName, Collections.emptyList());
   }
 
-  /** A fully-qualified class name for nested classes. */
-  private static final class NestedClassName extends ClassName {
-    /** From top to bottom. This will be ["java.util", "Map", "Entry"] for {@link Map.Entry}. */
-    final ClassName enclosingClassName;
-
-    private NestedClassName(ClassName enclosingClassName, String simpleName) {
-      this(enclosingClassName, simpleName, new ArrayList<>());
-    }
-
-    private NestedClassName(
-        ClassName enclosingClassName, String simpleName, List<AnnotationSpec> annotations) {
-      super(simpleName, annotations);
-      this.enclosingClassName = enclosingClassName;
-      this.canonicalName =
-          String.join(".", Arrays.asList(enclosingClassName.canonicalName, simpleName));
-    }
-
-    @Override public NestedClassName annotated(List<AnnotationSpec> annotations) {
-      return new NestedClassName(enclosingClassName, simpleName, concatAnnotations(annotations));
-    }
-
-    @Override public NestedClassName withoutAnnotations() {
-      return new NestedClassName(enclosingClassName.withoutAnnotations(), simpleName);
-    }
-
-    /** Returns the package name, like {@code "java.util"} for {@code Map.Entry}. */
-    public String packageName() {
-      return enclosingClassName.packageName();
-    }
-
-    @Override
-    public ClassName enclosingClassName() {
-      return enclosingClassName;
-    }
-
-    @Override
-    public ClassName topLevelClassName() {
-      return enclosingClassName.topLevelClassName();
-    }
-
-    @Override
-    public String reflectionName() {
-      return enclosingClassName.reflectionName() + "$" + simpleName;
-    }
-
-    @Override
-    public List<String> simpleNames() {
-      List<String> simpleNames = new ArrayList<>(enclosingClassName().simpleNames());
-      simpleNames.add(simpleName);
-      return simpleNames;
-    }
-
-    @Override
-    public ClassName peerClass(String name) {
-      return enclosingClassName.nestedClass(name);
-    }
-
-    @Override
-    ClassName prefixWithAtMostOneAnnotatedClass() {
-      if (hasAnnotatedEnclosingClass()) {
-        enclosingClassName.prefixWithAtMostOneAnnotatedClass();
-      }
-
-      return this;
-    }
-
-    @Override
-    CodeWriter emitWithoutPrefix(
-        CodeWriter out, ClassName unannotatedPrefix) throws IOException {
-
-      if (unannotatedPrefix.equals(this)) {
-        return out;
-      }
-
-      enclosingClassName.emitWithoutPrefix(out, unannotatedPrefix);
-      out.emit(".");
-      if (isAnnotated()) {
-        out.emit(" ");
-        emitAnnotations(out);
-      }
-      return out.emit(simpleName);
-    }
-
-    @Override
-    boolean hasAnnotatedEnclosingClass() {
-      return enclosingClassName.isAnnotated() || enclosingClassName.hasAnnotatedEnclosingClass();
-    }
+  private ClassName(String packageName, ClassName enclosingClassName, String simpleName,
+      List<AnnotationSpec> annotations) {
+    super(annotations);
+    this.packageName = packageName;
+    this.enclosingClassName = enclosingClassName;
+    this.simpleName = simpleName;
+    this.canonicalName = enclosingClassName != null
+        ? (enclosingClassName.canonicalName + '.' + simpleName)
+        : (packageName.isEmpty() ? simpleName : packageName + '.' + simpleName);
   }
 
-  public static final ClassName OBJECT = ClassName.get(Object.class);
+  @Override public ClassName annotated(List<AnnotationSpec> annotations) {
+    return new ClassName(packageName, enclosingClassName, simpleName,
+        concatAnnotations(annotations));
+  }
 
-  private ClassName(String simpleName, List<AnnotationSpec> annotations) {
-    super(annotations);
-    checkArgument(SourceVersion.isName(simpleName), "part '%s' is keyword", simpleName);
-    this.simpleName = simpleName;
+  @Override public ClassName withoutAnnotations() {
+    if (!isAnnotated()) return this;
+    ClassName resultEnclosingClassName = enclosingClassName != null
+        ? enclosingClassName.withoutAnnotations()
+        : null;
+    return new ClassName(packageName, resultEnclosingClassName, simpleName);
   }
 
+  @Override public boolean isAnnotated() {
+    return super.isAnnotated() || (enclosingClassName != null && enclosingClassName.isAnnotated());
+  }
 
-  /** Returns the package name, like {@code "java.util"} for {@code Map.Entry}. */
-  public abstract String packageName();
+  /**
+   * Returns the package name, like {@code "java.util"} for {@code Map.Entry}. Returns the empty
+   * string for the default package.
+   */
+  public String packageName() {
+    return packageName;
+  }
 
   /**
    * Returns the enclosing class, like {@link Map} for {@code Map.Entry}. Returns null if this class
    * is not nested in another class.
    */
-  public abstract ClassName enclosingClassName();
+  public ClassName enclosingClassName() {
+    return enclosingClassName;
+  }
 
   /**
    * Returns the top class in this nesting group. Equivalent to chained calls to {@link
    * #enclosingClassName()} until the result's enclosing class is null.
    */
-  public abstract ClassName topLevelClassName();
-
-  /**
-   * Return the binary name of a class.
-   */
-  public abstract String reflectionName();
-
-  public abstract ClassName withoutAnnotations();
-
-  /**
-   * Returns a new {@link ClassName} instance for the specified {@code name} as nested inside this
-   * class.
-   */
-  public ClassName nestedClass(String name) {
-    return new NestedClassName(this, name);
+  public ClassName topLevelClassName() {
+    return enclosingClassName != null ? enclosingClassName.topLevelClassName() : this;
   }
 
-  @Override
-  public ClassName annotated(List<AnnotationSpec> annotations) {
-    return (ClassName) super.annotated(annotations);
+  /** Return the binary name of a class. */
+  public String reflectionName() {
+    return enclosingClassName != null
+        ? (enclosingClassName.reflectionName() + '$' + simpleName)
+        : (packageName.isEmpty() ? simpleName : packageName + '.' + simpleName);
   }
 
-  public abstract List<String> simpleNames();
+  public List<String> simpleNames() {
+    List<String> simpleNames = new ArrayList<>();
+    if (enclosingClassName != null) {
+      simpleNames.addAll(enclosingClassName().simpleNames());
+    }
+    simpleNames.add(simpleName);
+    return simpleNames;
+  }
 
   /**
    * Returns a class that shares the same enclosing package or class. If this class is enclosed by
    * another class, this is equivalent to {@code enclosingClassName().nestedClass(name)}. Otherwise
    * it is equivalent to {@code get(packageName(), name)}.
    */
-  public abstract ClassName peerClass(String name);
-
-  abstract ClassName prefixWithAtMostOneAnnotatedClass();
-
-  abstract boolean hasAnnotatedEnclosingClass();
+  public ClassName peerClass(String name) {
+    return new ClassName(packageName, enclosingClassName, name);
+  }
 
-  abstract CodeWriter emitWithoutPrefix(
-      CodeWriter out, ClassName unannotatedPrefix) throws IOException;
+  /**
+   * Returns a new {@link ClassName} instance for the specified {@code name} as nested inside this
+   * class.
+   */
+  public ClassName nestedClass(String name) {
+    return new ClassName(packageName, this, name);
+  }
 
   /** Returns the simple name of this class, like {@code "Entry"} for {@link Map.Entry}. */
   public String simpleName() {
@@ -281,8 +157,8 @@ public abstract class ClassName extends TypeName implements Comparable<ClassName
     if (clazz.getEnclosingClass() == null) {
       // Avoid unreliable Class.getPackage(). https://github.com/square/javapoet/issues/295
       int lastDot = clazz.getName().lastIndexOf('.');
-      String packageName = (lastDot != -1)  ? clazz.getName().substring(0, lastDot) : null;
-      return new TopLevelClassName(packageName, name);
+      String packageName = (lastDot != -1) ? clazz.getName().substring(0, lastDot) : null;
+      return new ClassName(packageName, null, name);
     }
 
     return ClassName.get(clazz.getEnclosingClass()).nestedClass(name);
@@ -303,21 +179,14 @@ public abstract class ClassName extends TypeName implements Comparable<ClassName
       p = classNameString.indexOf('.', p) + 1;
       checkArgument(p != 0, "couldn't make a guess for %s", classNameString);
     }
-    String packageName = p == 0 ? null : classNameString.substring(0, p - 1);
-    String[] classNames = classNameString.substring(p).split("\\.", -1);
-
-    checkArgument(classNames.length >= 1, "couldn't make a guess for %s", classNameString);
+    String packageName = p == 0 ? "" : classNameString.substring(0, p - 1);
 
-    String simpleName = classNames[0];
-    checkArgument(!simpleName.isEmpty() && Character.isUpperCase(simpleName.codePointAt(0)),
-        "couldn't make a guess for %s", classNameString);
-    ClassName className = new TopLevelClassName(packageName, simpleName);
-
-    // Add the class names, like "Map" and "Entry".
-    for (String part : Arrays.asList(classNames).subList(1, classNames.length)) {
-      checkArgument(!part.isEmpty() && Character.isUpperCase(part.codePointAt(0)),
+    // Add class names like "Map" and "Entry".
+    ClassName className = null;
+    for (String simpleName : classNameString.substring(p).split("\\.", -1)) {
+      checkArgument(!simpleName.isEmpty() && Character.isUpperCase(simpleName.codePointAt(0)),
           "couldn't make a guess for %s", classNameString);
-      className = className.nestedClass(part);
+      className = new ClassName(packageName, className, simpleName);
     }
 
     return className;
@@ -328,7 +197,7 @@ public abstract class ClassName extends TypeName implements Comparable<ClassName
    * {@code "java.util"} and simple names {@code "Map"}, {@code "Entry"} yields {@link Map.Entry}.
    */
   public static ClassName get(String packageName, String simpleName, String... simpleNames) {
-    ClassName className = new TopLevelClassName(packageName, simpleName);
+    ClassName className = new ClassName(packageName, null, simpleName);
     for (String name : simpleNames) {
       className = className.nestedClass(name);
     }
@@ -347,7 +216,7 @@ public abstract class ClassName extends TypeName implements Comparable<ClassName
     }
 
     String packageName = getPackage(element.getEnclosingElement()).getQualifiedName().toString();
-    return new TopLevelClassName(packageName, simpleName);
+    return new ClassName(packageName, null, simpleName);
   }
 
   private static boolean isClassOrInterface(Element e) {
@@ -362,27 +231,54 @@ public abstract class ClassName extends TypeName implements Comparable<ClassName
   }
 
   @Override public int compareTo(ClassName o) {
-    return reflectionName().compareTo(o.reflectionName());
+    return canonicalName.compareTo(o.canonicalName);
   }
 
   @Override CodeWriter emit(CodeWriter out) throws IOException {
-    ClassName prefix = prefixWithAtMostOneAnnotatedClass();
-    String unqualifiedName = out.lookupName(prefix);
-    if (prefix.isAnnotated())  {
-      int dot = unqualifiedName.lastIndexOf(".");
-      out.emitAndIndent(unqualifiedName.substring(0, dot + 1));
-      if (dot != -1) {
-        out.emit(" ");
+    boolean charsEmitted = false;
+    for (ClassName className : enclosingClasses()) {
+      String simpleName;
+      if (charsEmitted) {
+        // We've already emitted an enclosing class. Emit as we go.
+        out.emit(".");
+        simpleName = className.simpleName;
+
+      } else if (className.isAnnotated() || className == this) {
+        // We encountered the first enclosing class that must be emitted.
+        String qualifiedName = out.lookupName(className);
+        int dot = qualifiedName.lastIndexOf('.');
+        if (dot != -1) {
+          out.emitAndIndent(qualifiedName.substring(0, dot + 1));
+          simpleName = qualifiedName.substring(dot + 1);
+          charsEmitted = true;
+        } else {
+          simpleName = qualifiedName;
+        }
+
+      } else {
+        // Don't emit this enclosing type. Keep going so we can be more precise.
+        continue;
+      }
+
+      if (className.isAnnotated()) {
+        if (charsEmitted) out.emit(" ");
+        className.emitAnnotations(out);
       }
-      prefix.emitAnnotations(out);
-      out.emit(unqualifiedName.substring(dot + 1));
-    } else {
-        out.emitAndIndent(unqualifiedName);
+
+      out.emit(simpleName);
+      charsEmitted = true;
     }
-    return emitWithoutPrefix(out, prefix);
+
+    return out;
   }
 
-  private static boolean isDefaultPackage(String packageName) {
-    return packageName == null || packageName.isEmpty();
+  /** Returns all enclosing classes in this, outermost first. */
+  private List<ClassName> enclosingClasses() {
+    List<ClassName> result = new ArrayList<>();
+    for (ClassName c = this; c != null; c = c.enclosingClassName) {
+      result.add(c);
+    }
+    Collections.reverse(result);
+    return result;
   }
 }
diff --git a/src/test/java/com/squareup/javapoet/AnnotatedTypeNameTest.java b/src/test/java/com/squareup/javapoet/AnnotatedTypeNameTest.java
index 7261568..42734ff 100644
--- a/src/test/java/com/squareup/javapoet/AnnotatedTypeNameTest.java
+++ b/src/test/java/com/squareup/javapoet/AnnotatedTypeNameTest.java
@@ -15,18 +15,18 @@
  */
 package com.squareup.javapoet;
 
-import static com.google.common.truth.Truth.assertThat;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertTrue;
-
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Target;
 import java.util.List;
 import java.util.Map;
 import org.junit.Test;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
+
 public class AnnotatedTypeNameTest {
 
   private final static String NN = NeverNull.class.getCanonicalName();
@@ -125,6 +125,12 @@ public class AnnotatedTypeNameTest {
     assertThat(type.toString()).isEqualTo("java.util.Map. @" + TUA + " Entry");
   }
 
+  @Test public void annotatedEnclosingAndNestedType() {
+    TypeName type = ((ClassName) TypeName.get(Map.class).annotated(TYPE_USE_ANNOTATION))
+        .nestedClass("Entry").annotated(TYPE_USE_ANNOTATION);
+    assertThat(type.toString()).isEqualTo("java.util. @" + TUA + " Map. @" + TUA + " Entry");
+  }
+
   // https://github.com/square/javapoet/issues/431
   @Test public void annotatedNestedParameterizedType() {
     TypeName type = ParameterizedTypeName.get(Map.Entry.class, Byte.class, Byte.class)
@@ -133,6 +139,27 @@ public class AnnotatedTypeNameTest {
         .isEqualTo("java.util.Map. @" + TUA + " Entry<java.lang.Byte, java.lang.Byte>");
   }
 
+  @Test public void withoutAnnotationsOnAnnotatedEnclosingAndNestedType() {
+    TypeName type = ((ClassName) TypeName.get(Map.class).annotated(TYPE_USE_ANNOTATION))
+        .nestedClass("Entry").annotated(TYPE_USE_ANNOTATION);
+    assertThat(type.isAnnotated()).isTrue();
+    assertThat(type.withoutAnnotations()).isEqualTo(TypeName.get(Map.Entry.class));
+  }
+
+  @Test public void withoutAnnotationsOnAnnotatedEnclosingType() {
+    TypeName type = ((ClassName) TypeName.get(Map.class).annotated(TYPE_USE_ANNOTATION))
+        .nestedClass("Entry");
+    assertThat(type.isAnnotated()).isTrue();
+    assertThat(type.withoutAnnotations()).isEqualTo(TypeName.get(Map.Entry.class));
+  }
+
+  @Test public void withoutAnnotationsOnAnnotatedNestedType() {
+    TypeName type = ((ClassName) TypeName.get(Map.class))
+        .nestedClass("Entry").annotated(TYPE_USE_ANNOTATION);
+    assertThat(type.isAnnotated()).isTrue();
+    assertThat(type.withoutAnnotations()).isEqualTo(TypeName.get(Map.Entry.class));
+  }
+
   // https://github.com/square/javapoet/issues/614
    @Test public void annotatedArrayType() {
     TypeName type = ArrayTypeName.of(ClassName.get(Object.class)).annotated(TYPE_USE_ANNOTATION);