Fix TypeName.get() on top-level error types (#646)

* Fix TypeName.get() on top-level error types

Discussion: https://github.com/square/javapoet/commit/1c897e44ddd6b50e22180b6d32e4a3c69223614a#commitcomment-28905786

* Use an ElementVisitor instead of instanceof checks

* Remove unused import

2 files changed, 59 insertions, 62 deletions

diff --git a/src/main/java/com/squareup/javapoet/ClassName.java b/src/main/java/com/squareup/javapoet/ClassName.java
index 766d9e2..99c4ed2 100644
--- a/src/main/java/com/squareup/javapoet/ClassName.java
+++ b/src/main/java/com/squareup/javapoet/ClassName.java
@@ -20,9 +20,10 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import javax.lang.model.element.Name;
+import javax.lang.model.element.Element;
 import javax.lang.model.element.PackageElement;
 import javax.lang.model.element.TypeElement;
+import javax.lang.model.util.SimpleElementVisitor8;
 
 import static com.squareup.javapoet.Util.checkArgument;
 import static com.squareup.javapoet.Util.checkNotNull;
@@ -206,18 +207,23 @@ public final class ClassName extends TypeName implements Comparable<ClassName> {
     checkNotNull(element, "element == null");
     String simpleName = element.getSimpleName().toString();
 
-    switch (element.getNestingKind()) {
-      case TOP_LEVEL:
-        Name packageName = ((PackageElement) element.getEnclosingElement()).getQualifiedName();
-        return new ClassName(packageName.toString(), null, simpleName);
+    return element.getEnclosingElement().accept(new SimpleElementVisitor8<ClassName, Void>() {
+      @Override public ClassName visitPackage(PackageElement packageElement, Void p) {
+        return new ClassName(packageElement.getQualifiedName().toString(), null, simpleName);
+      }
 
-      case MEMBER:
-        ClassName enclosingClass = get((TypeElement) element.getEnclosingElement());
-        return enclosingClass.nestedClass(simpleName);
+      @Override public ClassName visitType(TypeElement enclosingClass, Void p) {
+        return ClassName.get(enclosingClass).nestedClass(simpleName);
+      }
 
-      default:
-        throw new IllegalArgumentException("unexpected type nesting");
-    }
+      @Override public ClassName visitUnknown(Element unknown, Void p) {
+        return get("", simpleName);
+      }
+
+      @Override public ClassName defaultAction(Element enclosingElement, Void p) {
+        throw new IllegalArgumentException("Unexpected type nesting: " + element);
+      }
+    }, null);
   }
 
   @Override public int compareTo(ClassName o) {
diff --git a/src/test/java/com/squareup/javapoet/AbstractTypesTest.java b/src/test/java/com/squareup/javapoet/AbstractTypesTest.java
index 6d1f5a2..86d9cbc 100644
--- a/src/test/java/com/squareup/javapoet/AbstractTypesTest.java
+++ b/src/test/java/com/squareup/javapoet/AbstractTypesTest.java
@@ -16,19 +16,27 @@
 package com.squareup.javapoet;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.testing.compile.CompilationSubject.assertThat;
+import static com.google.testing.compile.Compiler.javac;
+import static javax.lang.model.util.ElementFilter.fieldsIn;
 import static org.junit.Assert.*;
 
+import com.google.testing.compile.Compilation;
+import com.google.testing.compile.JavaFileObjects;
 import java.io.Serializable;
 import java.lang.annotation.Annotation;
 import java.nio.charset.Charset;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-
+import javax.annotation.processing.AbstractProcessor;
+import javax.annotation.processing.RoundEnvironment;
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.element.TypeParameterElement;
+import javax.lang.model.element.VariableElement;
 import javax.lang.model.type.DeclaredType;
 import javax.lang.model.type.ErrorType;
 import javax.lang.model.type.TypeKind;
@@ -37,6 +45,7 @@ import javax.lang.model.type.TypeVisitor;
 import javax.lang.model.type.WildcardType;
 import javax.lang.model.util.Elements;
 import javax.lang.model.util.Types;
+import javax.tools.JavaFileObject;
 
 import org.junit.Test;
 
@@ -68,10 +77,38 @@ public abstract class AbstractTypesTest {
         .isEqualTo(ParameterizedTypeName.get(ClassName.get(Set.class), ClassName.OBJECT));
   }
 
-  @Test public void getErrorType() {
-    ErrorType errorType =
-        new DeclaredTypeAsErrorType(getTypes().getDeclaredType(getElement(Set.class)));
-    assertThat(TypeName.get(errorType)).isEqualTo(ClassName.get(Set.class));
+  @Test public void errorTypes() {
+    JavaFileObject hasErrorTypes =
+        JavaFileObjects.forSourceLines(
+            "com.squareup.tacos.ErrorTypes",
+            "package com.squareup.tacos;",
+            "",
+            "@SuppressWarnings(\"hook-into-compiler\")",
+            "class ErrorTypes {",
+            "  Tacos tacos;",
+            "  Ingredients.Guacamole guacamole;",
+            "}");
+    Compilation compilation = javac().withProcessors(new AbstractProcessor() {
+      @Override
+      public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnvironment) {
+        TypeElement classFile =
+            processingEnv.getElementUtils().getTypeElement("com.squareup.tacos.ErrorTypes");
+        List<VariableElement> fields = fieldsIn(classFile.getEnclosedElements());
+        ErrorType topLevel = (ErrorType) fields.get(0).asType();
+        ErrorType member = (ErrorType) fields.get(1).asType();
+
+        assertThat(TypeName.get(topLevel)).isEqualTo(ClassName.get("", "Tacos"));
+        assertThat(TypeName.get(member)).isEqualTo(ClassName.get("Ingredients", "Guacamole"));
+        return false;
+      }
+
+      @Override
+      public Set<String> getSupportedAnnotationTypes() {
+        return Collections.singleton("*");
+      }
+    }).compile(hasErrorTypes);
+
+    assertThat(compilation).failed();
   }
 
   static class Parameterized<
@@ -245,50 +282,4 @@ public abstract class AbstractTypesTest {
     } catch (UnsupportedOperationException expected) {
     }
   }
-
-  private static class DeclaredTypeAsErrorType implements ErrorType {
-    private final DeclaredType declaredType;
-
-    public DeclaredTypeAsErrorType(DeclaredType declaredType) {
-      this.declaredType = declaredType;
-    }
-
-    @Override
-    public Element asElement() {
-      return declaredType.asElement();
-    }
-
-    @Override
-    public TypeMirror getEnclosingType() {
-      return declaredType.getEnclosingType();
-    }
-
-    @Override
-    public List<? extends TypeMirror> getTypeArguments() {
-      return declaredType.getTypeArguments();
-    }
-
-    @Override
-    public TypeKind getKind() {
-      return declaredType.getKind();
-    }
-
-    @Override
-    public <R, P> R accept(TypeVisitor<R, P> typeVisitor, P p) {
-      return typeVisitor.visitError(this, p);
-    }
-
-    // JDK8 Compatibility:
-    public <A extends Annotation> A[] getAnnotationsByType(Class<A> annotationType) {
-      throw new UnsupportedOperationException();
-    }
-
-    public <A extends Annotation> A getAnnotation(Class<A> annotationType) {
-      throw new UnsupportedOperationException();
-    }
-
-    public List<? extends AnnotationMirror> getAnnotationMirrors() {
-      throw new UnsupportedOperationException();
-    }
-  }
 }