diff --git a/lib/CodeGen/ModuleBuilder.cpp b/lib/CodeGen/ModuleBuilder.cpp
index 7329ed1..6ce4e68 100644
--- a/lib/CodeGen/ModuleBuilder.cpp
+++ b/lib/CodeGen/ModuleBuilder.cpp
@@ -49,13 +49,7 @@
     
     virtual ~CodeGeneratorImpl() {}
     
-    virtual llvm::Module* ReleaseModule() {      
-      if (Diags.hasErrorOccurred())
-        return 0;
-      
-      if (Builder)
-        Builder->Release();
-      
+    virtual llvm::Module* ReleaseModule() {
       return M.take();
     }
     
@@ -134,7 +128,16 @@
     virtual void HandleTagDeclDefinition(TagDecl *D) {
       Builder->UpdateCompletedType(D);
     }
-    
+
+    virtual void HandleTranslationUnit(TranslationUnit& TU) {
+      if (Diags.hasErrorOccurred()) {
+        M.reset();
+        return;
+      }
+
+      if (Builder)
+        Builder->Release();
+    };
   };
 }