1 files changed, 24 insertions, 39 deletions

diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
index 9befa36..857311d 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
@@ -12,58 +12,46 @@ import static java.lang.reflect.Modifier.*;
 import static org.springframework.util.StringUtils.collectionToDelimitedString;
 
 public class ClassG {
-    private final String packageName;
-    public final String className;
-    private final String extendsClass;
+    public final TypeRef type;
+    private final TypeRef extendsClass;
     private final Set<TypeRef> imports = new TreeSet<>();
     private final Set<FieldRef> fields = new TreeSet<>();
     private final List<MethodRef> methods = new ArrayList<>();
     private final List<Constructor> constructors = new ArrayList<>();
 
-    public ClassG(String packageName, String className) {
-        this(packageName, className, null);
+    public ClassG(TypeRef type) {
+        this(type, null);
     }
 
-    public ClassG(String packageName, String className, String extendsClass) {
-        this.packageName = packageName;
-        this.className = className;
+    public ClassG(TypeRef type, TypeRef extendsClass) {
+        this.type = type;
         this.extendsClass = extendsClass;
     }
 
     public TypeRef addImport(final TypeMirror type) {
-        if (type.getKind().isPrimitive()) {
-            return TypeRef.find(type.getKind());
-        }
-
-        String canonicalName = type.toString();
-
-        for (TypeRef i : imports) {
-            if (i.canonicalName().equals(canonicalName)) {
-                return i;
-            }
-        }
+        return addImport(new TypeRef(type));
+    }
 
-        TypeRef ref = new TypeRef(canonicalName, canonicalName);
-        imports.add(ref);
-        return ref;
+    public TypeRef addImport(Class<?> c) {
+        return addImport(new TypeRef(c));
     }
 
-    public TypeRef addImport(Class<?> klass) {
-        String canonicalName = klass.getCanonicalName();
-        String simpleName = klass.getSimpleName();
-        String name = simpleName;
+    public TypeRef addImport(TypeRef klass) {
+        String fqName = klass.fqName;
+        String name = klass.name;
         for (TypeRef i : imports) {
-            if (i.canonicalName().equals(canonicalName)) {
+            if (i.fqName.equals(fqName)) {
                 return i;
             }
 
-            if (i.name.equals(simpleName)) {
-                name = canonicalName;
+            // If we've already found an import with the same name, use the fq version
+            if (i.name.equals(name)) {
+                name = fqName;
                 break;
             }
         }
 
-        TypeRef ref = new TypeRef(name, canonicalName);
+        TypeRef ref = new TypeRef(name, fqName);
         imports.add(ref);
 
         return ref;
@@ -114,26 +102,23 @@ public class ClassG {
     }
 
     public void write(PrintWriter writer) {
-        if (packageName.length() > 0) {
-            writer.println("package " + packageName + ";");
+        if (!type.inUnnamedPackage()) {
+            writer.println("package " + type.packageName() + ";");
             writer.println();
         }
         for (TypeRef i : imports) {
-            if (i.isPrimitive()) {
-                continue;
-            }
-            if (i.canonicalName().indexOf('.') == -1) {
+            if (i.isPrimitive() || i.inUnnamedPackage()) {
                 continue;
             }
-            writer.println("import " + i.canonicalName() + ";");
+            writer.println("import " + i.fqName + ";");
         }
         if (!imports.isEmpty()) {
             writer.println();
         }
 
-        String extendsString = extendsClass == null ? "" : " extends " + extendsClass;
+        String extendsString = extendsClass == null ? "" : " extends " + extendsClass.name;
 
-        writer.println("public class " + className + extendsString + " {");
+        writer.println("public class " + type.className + extendsString + " {");
         for (FieldRef field : fields) {
             writer.println();
             writer.println("    " + field.toJava() + ";");