1 files changed, 32 insertions, 73 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 e71bbb7..c99feae 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
@@ -1,13 +1,11 @@
 package io.trygvis.container.compiler.model;
 
-import io.trygvis.container.compiler.Utils;
 import org.apache.commons.lang.StringUtils;
 
 import javax.annotation.Generated;
 import javax.lang.model.type.TypeMirror;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 import java.util.Set;
 import java.util.TreeSet;
@@ -23,18 +21,17 @@ public class ClassG {
     public final TypeRef type;
     private TypeRef extendsType;
     private List<TypeRef> implementsTypes = new ArrayList<>();
-    private final Set<TypeRef> imports = new TreeSet<>();
+    public final Imports imports = new Imports();
     private final Set<FieldRef> fields = new TreeSet<>();
     private final List<MethodRef> methods = new ArrayList<>();
     private final List<Constructor> constructors = new ArrayList<>();
     private final List<InnerClassG> innerClasses = new ArrayList<>();
 
-    public static class InnerClassG extends ClassG {
-        public final ClassG parent;
+    public static class InnerClassG {
+        public final ClassG inner;
 
-        public InnerClassG(ClassG parent, int modifiers, TypeRef type) {
-            super(modifiers, type);
-            this.parent = parent;
+        public InnerClassG(ClassG inner) {
+            this.inner = inner;
         }
     }
 
@@ -57,44 +54,15 @@ public class ClassG {
         return this;
     }
 
-    public TypeRef addImport(final TypeMirror type) {
-        return addImport(new TypeRef(type));
-    }
-
-    public TypeRef addImport(Class<?> c) {
-        return addImport(new TypeRef(c));
-    }
-
-    public TypeRef addImport(TypeRef klass) {
-        String fqName = klass.fqName;
-        String name = klass.toString();
-        for (TypeRef i : imports) {
-            if (i.fqName.equals(fqName)) {
-                return i;
-            }
-
-            // If we've already found an import with the same name, use the fq version
-            if (i.plainName.equals(name)) {
-                name = fqName;
-                break;
-            }
-        }
-
-        TypeRef ref = new TypeRef(fqName, name);
-        imports.add(ref);
-
-        return ref;
-    }
-
     public FieldRef addField(TypeMirror klass, String name) {
-        TypeRef type = addImport(klass);
+        TypeRef type = imports.add(klass);
         FieldRef ref = new FieldRef(PRIVATE | FINAL, type, name);
         fields.add(ref);
         return ref;
     }
 
     public FieldRef addField(int modifiers, TypeRef type, String name) {
-        TypeRef t = addImport(type);
+        TypeRef t = imports.add(type);
         FieldRef ref = new FieldRef(modifiers, t, name);
         fields.add(ref);
         return ref;
@@ -118,44 +86,43 @@ public class ClassG {
         return constructor;
     }
 
-    public MethodRef addMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
-        MethodRef ref = new MethodRef(PUBLIC, returnType, name, parameters, body);
-        methods.add(ref);
-        return ref;
+    public ClassG addMethod(MethodRef methodRef) {
+        this.methods.add(methodRef);
+        return this;
     }
 
-    public MethodRef addStaticMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
-        MethodRef ref = new MethodRef(PUBLIC | STATIC, returnType, name, parameters, body);
-        methods.add(ref);
-        return ref;
-    }
+//    public MethodRef addMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
+//        MethodRef ref = new MethodRef(PUBLIC, returnType, name, parameters, body);
+//        methods.add(ref);
+//        return ref;
+//    }
+//
+//    public MethodRef addStaticMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
+//        MethodRef ref = new MethodRef(PUBLIC | STATIC, returnType, name, parameters, body);
+//        methods.add(ref);
+//        return ref;
+//    }
 
-    public InnerClassG addInnerClass(int modifiers, TypeRef type) {
-        InnerClassG inner = new InnerClassG(this, modifiers, type);
-        innerClasses.add(inner);
-        return inner;
+    public InnerClassG addInnerClass(ClassG inner) {
+        InnerClassG i = new InnerClassG(inner);
+        innerClasses.add(i);
+        return i;
     }
 
     public final List<String> generate() {
-        TypeRef generatedType = addImport(Generated.class);
+        return generate(false);
+    }
+
+    public final List<String> generate(boolean isInner) {
+        TypeRef generatedType = imports.add(Generated.class);
 
         List<String> body = new ArrayList<>();
-        boolean isInner = this instanceof InnerClassG;
         if (!type.inUnnamedPackage() && !isInner) {
             body.add("package " + type.packageName() + ";");
             body.add("");
         }
         if (!isInner) {
-            // TODO: Add imports from inner classes
-            for (TypeRef i : getImports()) {
-                if (i.isPrimitive() || i.inUnnamedPackage()) {
-                    continue;
-                }
-                body.add("import " + i.fqName + ";");
-            }
-            if (!imports.isEmpty()) {
-                body.add("");
-            }
+            body.addAll(imports.generate());
         }
 
         String extendsString = extendsType == null ? "" : " extends " + extendsType;
@@ -173,7 +140,7 @@ public class ClassG {
 
         for (InnerClassG innerClass : innerClasses) {
             body.add("");
-            addAll(1, body, innerClass.generate());
+            addAll(1, body, innerClass.inner.generate(true));
         }
 
         for (Constructor constructor : constructors) {
@@ -197,14 +164,6 @@ public class ClassG {
         }
     }
 
-    protected Collection<? extends TypeRef> getImports() {
-        List<TypeRef> imports = new ArrayList<>(this.imports);
-        for (InnerClassG c : innerClasses) {
-            imports.addAll(c.getImports());
-        }
-        return imports;
-    }
-
     private List<String> write(MethodRef method) {
         List<String> body = new ArrayList<>();
         String returnString;