--- a/graal/com.oracle.truffle.dsl.processor/src/com/oracle/truffle/dsl/processor/AnnotationProcessor.java	Mon Aug 11 15:53:05 2014 +0200
+++ b/graal/com.oracle.truffle.dsl.processor/src/com/oracle/truffle/dsl/processor/AnnotationProcessor.java	Mon Aug 11 15:53:05 2014 +0200
@@ -33,6 +33,7 @@
 import com.oracle.truffle.dsl.processor.ast.*;
 import com.oracle.truffle.dsl.processor.codewriter.*;
 import com.oracle.truffle.dsl.processor.compiler.*;
+import com.oracle.truffle.dsl.processor.node.*;
 import com.oracle.truffle.dsl.processor.template.*;
 
 /**
@@ -42,12 +43,10 @@
 
     private final AbstractParser<M> parser;
     private final CompilationUnitFactory<M> factory;
-    private final ProcessorContext context;
 
     private final Set<String> processedElements = new HashSet<>();
 
-    public AnnotationProcessor(ProcessorContext context, AbstractParser<M> parser, CompilationUnitFactory<M> factory) {
-        this.context = context;
+    public AnnotationProcessor(AbstractParser<M> parser, CompilationUnitFactory<M> factory) {
         this.parser = parser;
         this.factory = factory;
     }
@@ -56,13 +55,8 @@
         return parser;
     }
 
-    public ProcessorContext getContext() {
-        return context;
-    }
-
     @SuppressWarnings({"unchecked"})
     public void process(RoundEnvironment env, Element element, boolean callback) {
-
         // since it is not guaranteed to be called only once by the compiler
         // we check for already processed elements to avoid errors when writing files.
         if (!callback && element instanceof TypeElement) {
@@ -73,6 +67,7 @@
             processedElements.add(qualifiedName);
         }
 
+        ProcessorContext context = ProcessorContext.getInstance();
         TypeElement type = (TypeElement) element;
 
         M model = (M) context.getTemplate(type.asType(), false);
@@ -85,6 +80,7 @@
 
             if (model != null) {
                 CodeCompilationUnit unit = factory.process(null, model);
+                patchGeneratedTypes(unit);
                 unit.setGeneratorAnnotationMirror(model.getTemplateTypeAnnotation());
                 unit.setGeneratorElement(model.getTemplateType());
 
@@ -100,6 +96,55 @@
         }
     }
 
+    private static void patchGeneratedTypes(CodeCompilationUnit unit) {
+        final Map<String, CodeTypeElement> classes = new HashMap<>();
+
+        unit.accept(new CodeElementScanner<Void, Void>() {
+            @Override
+            public Void visitType(CodeTypeElement e, Void p) {
+                classes.put(e.getSimpleName().toString(), e);
+                return super.visitType(e, p);
+            }
+
+        }, null);
+
+        unit.accept(new CodeElementScanner<Void, Void>() {
+            @Override
+            public Void visitExecutable(CodeExecutableElement e, Void p) {
+                if (e.getReturnType() instanceof GeneratedTypeMirror) {
+                    e.setReturnType(patchType(e.getReturnType()));
+                }
+                for (VariableElement element : e.getParameters()) {
+                    if (element instanceof CodeVariableElement) {
+                        CodeVariableElement var = ((CodeVariableElement) element);
+                        if (var.getType() instanceof GeneratedTypeMirror) {
+                            var.setType(patchType(var.getType()));
+                        }
+                    }
+                }
+                return super.visitExecutable(e, p);
+            }
+
+            @Override
+            public void visitTree(CodeTree e, Void p) {
+                if (e.getType() instanceof GeneratedTypeMirror) {
+                    e.setType(patchType(e.asType()));
+                }
+            }
+
+            private TypeMirror patchType(TypeMirror typeMirror) {
+                assert typeMirror instanceof GeneratedTypeMirror;
+                GeneratedTypeMirror type = (GeneratedTypeMirror) typeMirror;
+                CodeTypeElement generatedType = classes.get(Utils.fromTypeMirror(type).getSimpleName().toString());
+                if (generatedType == null) {
+                    return type;
+                }
+                return generatedType.asType();
+            }
+        }, null);
+
+    }
+
     private static class CodeWriter extends AbstractCodeWriter {
 
         private final Element originalElement;