1 files changed, 10 insertions, 3 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java
index 717715e..1c0f538 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java
@@ -10,7 +10,8 @@ import java.util.TreeMap;
 
 public class SqlUnitModel {
     private Map<TypeRef, EntityMirror> entities = new TreeMap<>();
-    private Map<EntityMirror, Element> entityElements = new TreeMap<>();
+    private Map<Element, EntityMirror> elementToEntity = new TreeMap<>(new ElementComparator());
+    private Map<EntityMirror, Element> entityToElement = new TreeMap<>();
     private Map<String, SequenceMirror> sequences = new TreeMap<>();
     private Map<SequenceMirror, Element> sequenceElements = new TreeMap<>();
     private String packageName;
@@ -44,7 +45,8 @@ public class SqlUnitModel {
 
     public SqlUnitModel add(EntityMirror entity, Element element) {
         entities.put(entity.type, entity);
-        entityElements.put(entity, element);
+        entityToElement.put(entity, element);
+        elementToEntity.put(element, entity);
 
         return this;
     }
@@ -70,7 +72,11 @@ public class SqlUnitModel {
     }
 
     public Element element(EntityMirror entity) {
-        return entityElements.get(entity);
+        return entityToElement.get(entity);
+    }
+
+    public EntityMirror entity(Element element) {
+        return elementToEntity.get(element);
     }
 
     // -----------------------------------------------------------------------
@@ -120,4 +126,5 @@ public class SqlUnitModel {
     public Element element(SequenceMirror sequence) {
         return sequenceElements.get(sequence);
     }
+
 }