1 files changed, 7 insertions, 5 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
index f944b1d..b5db075 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
@@ -2,6 +2,7 @@ package io.trygvis.container.compiler;
 
 import io.trygvis.container.compiler.model.ClassG;
 import io.trygvis.container.compiler.model.FieldRef;
+import io.trygvis.container.compiler.model.MethodRef;
 import io.trygvis.container.compiler.model.Parameters;
 import io.trygvis.container.compiler.model.TypeRef;
 import io.trygvis.container.tx.PlatformTransactionManager;
@@ -77,13 +78,13 @@ public class TransactionalHandler extends AbstractHandler {
             TypeKind kind = returnTypeMirror.getKind();
 
             boolean isVoid = kind == TypeKind.VOID;
-            TypeRef returnType = isVoid ? TypeRef.VOID : g.addImport(returnTypeMirror);
+            TypeRef returnType = isVoid ? TypeRef.VOID : g.imports.add(returnTypeMirror);
 
             Parameters parameters = new Parameters();
             List<String> arguments = new ArrayList<>();
             for (VariableElement ve : ee.getParameters()) {
 //                    parameters.add("final " + ve.asType().toString() + " " + ve.getSimpleName().toString());
-                TypeRef k = g.addImport(ve.asType());
+                TypeRef k = g.imports.add(ve.asType());
                 parameters.addParameter(k, ve.getSimpleName().toString());
                 arguments.add(ve.getSimpleName().toString());
             }
@@ -107,7 +108,8 @@ public class TransactionalHandler extends AbstractHandler {
             body.add("            }");
             body.add("        });");
 
-            g.addMethod(body, returnType, ee.getSimpleName().toString(), parameters);
+            MethodRef m = new MethodRef(PUBLIC, returnType, ee.getSimpleName().toString(), parameters, body);
+            g.addMethod(m);
         }
 
         JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile((p.length() == 0 ? "" : p + ".") + className, element);
@@ -122,11 +124,11 @@ public class TransactionalHandler extends AbstractHandler {
         Parameters parameters = new Parameters();
         List<String> goesToSuper = new ArrayList<>();
 
-        TypeRef k = g.addImport(PlatformTransactionManager.class);
+        TypeRef k = g.imports.add(PlatformTransactionManager.class);
         Parameters.ParameterRef transactionManager = parameters.addParameter(k, "transactionManager");
 
         for (VariableElement p : constructor.getParameters()) {
-            k = g.addImport(p.asType());
+            k = g.imports.add(p.asType());
             String name = p.getSimpleName().toString();
             parameters.addParameter(k, name);
             goesToSuper.add(name);