1 files changed, 15 insertions, 16 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorSupport.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorSupport.java
index 379f6f6..0563691 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorSupport.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorSupport.java
@@ -1,42 +1,41 @@
 package io.trygvis.persistence;
 
-import javax.lang.model.element.Element;
-import javax.lang.model.type.TypeMirror;
+import io.trygvis.container.compiler.model.TypeRef;
+
 import java.util.Date;
-import java.util.HashMap;
 import java.util.Map;
+import java.util.TreeMap;
 
 public class GeneratorSupport {
 
-    private final Map<String, TypeHandler> primitiveTypeHandlers = new HashMap<>();
-    private final Map<String, TypeHandler> typeHandlers = new HashMap<>();
+    private final Map<TypeRef, TypeHandler> primitiveTypeHandlers = new TreeMap<>();
+    private final Map<TypeRef, TypeHandler> typeHandlers = new TreeMap<>();
 
     {
-        typeHandlers.put(Integer.class.getName(), new TypeHandler.IntTypeHandler());
-        typeHandlers.put(Long.class.getName(), new TypeHandler.LongTypeHandler());
-        typeHandlers.put(String.class.getName(), new TypeHandler.StringTypeHandler());
-        typeHandlers.put(Date.class.getName(), new TypeHandler.DateTypeHandler());
+        typeHandlers.put(new TypeRef(Integer.class), new TypeHandler.IntTypeHandler());
+        typeHandlers.put(new TypeRef(Long.class), new TypeHandler.LongTypeHandler());
+        typeHandlers.put(new TypeRef(String.class), new TypeHandler.StringTypeHandler());
+        typeHandlers.put(new TypeRef(Date.class), new TypeHandler.DateTypeHandler());
 
         primitiveTypeHandlers.putAll(typeHandlers);
     }
 
-    public void addTypeHandler(String type, TypeHandler typeHandler) {
+    public void addTypeHandler(TypeRef type, TypeHandler typeHandler) {
         typeHandlers.put(type, typeHandler);
     }
 
-    public TypeHandler typeHandler(Element element) {
-        String type = element.asType().toString();
+    public TypeHandler typeHandler(TypeRef type) {
         TypeHandler typeHandler = typeHandlers.get(type);
         if (typeHandler == null)
             throw new RuntimeException("Unsupported field type: " + type);
         return typeHandler;
     }
 
-    public boolean isPrimitive(TypeMirror type) {
-        return primitiveTypeHandlers.containsKey(type.toString());
+    public boolean isPrimitive(TypeRef type) {
+        return primitiveTypeHandlers.containsKey(type);
     }
 
-    public boolean hasTypeHandler(TypeMirror type) {
-        return typeHandlers.containsKey(type.toString());
+    public boolean hasTypeHandler(TypeRef type) {
+        return typeHandlers.containsKey(type);
     }
 }