1 files changed, 23 insertions, 23 deletions

diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java
index f0f4ae7..5d6defb 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java
@@ -160,31 +160,31 @@ public class EntityMirror implements Comparable<EntityMirror> {
 
             i++;
 
+            TypeHandler typeHandler;
+            String accessor;
+            String setter;
             if (field.fieldType == PRIMITIVE) {
-                TypeHandler typeHandler = generatorConfiguration.typeHandler(field.type);
-                String access;
-                if (field.accessorType == FIELD) {
-                    access = o.name + "." + field.javaName;
-                } else {
-                    access = o.name + "." + toGetterName(field.javaName) + "()";
-                }
-                String setter = "    stmt." + typeHandler.setter(i, access) + ";";
-
-                if (field.notNull) {
-                    body.add(setter);
-                } else {
-                    body.add("    " + field.type + " " + field.javaName + " = " + access + ";");
-                    body.add("    if(" + field.javaName + " == null) {");
-                    body.add("        stmt.setNull(" + i + ", " + typesType + "." + typeHandler.typeName() + ");");
-                    body.add("    } else {");
-                    body.add("    " + setter);
-                    body.add("    }");
-                }
-            } else if (field.fieldType == REFERENCE) {
+                typeHandler = generatorConfiguration.typeHandler(field.type);
+                accessor = field.fieldAccessor(o);
+                setter = "    stmt." + typeHandler.setter(i, accessor) + ";";
+            } else {
                 EntityMirror referenced = unit.get(field.type);
                 FieldMirror idField = referenced.getIdField();
-                TypeHandler typeHandler = generatorConfiguration.typeHandler(idField.type);
-                body.add("    stmt." + typeHandler.setter(i, o.name + "." + field.javaName + "." + idField.javaName) + ";");
+                typeHandler = generatorConfiguration.typeHandler(idField.type);
+                accessor = field.referenceAccessor(o, idField);
+                setter = "    stmt." + typeHandler.setter(i, accessor) + ";";
+                accessor = "null";
+            }
+
+            if (field.notNull) {
+                body.add(setter);
+            } else {
+                body.add("    " + field.type + " " + field.javaName + " = " + accessor + ";");
+                body.add("    if(" + field.javaName + " == null) {");
+                body.add("        stmt.setNull(" + i + ", " + typesType + "." + typeHandler.typeName() + ");");
+                body.add("    } else {");
+                body.add("    " + setter);
+                body.add("    }");
             }
         }
         body.add("    stmt.executeUpdate();");
@@ -202,7 +202,7 @@ public class EntityMirror implements Comparable<EntityMirror> {
         List<String> arguments = new ArrayList<>();
         arguments.add(con.name);
         for (FieldMirror field : idFields) {
-            if(field.accessorType == FIELD) {
+            if (field.accessorType == FIELD) {
                 arguments.add(o.name + "." + field.javaName);
             } else {
                 arguments.add(o.name + "." + toGetterName(field.javaName) + "()");