From d72312c6023fa36997eecb2ef7b5028c28f605d5 Mon Sep 17 00:00:00 2001
From: Hong Chen <hchen99@users.noreply.github.com>
Date: Tue, 14 May 2024 10:53:11 -0500
Subject: [PATCH] Customize clang diagnostic handling. (#1707)

* Cutomize clang diagnostic handling.

* Removed code that was commented out.
---
 .../ICGDiagnosticConsumer.cpp                 | 39 +++++++++++++++++
 .../ICGDiagnosticConsumer.hh                  | 42 +++++++++++++++++++
 .../codegen/Interface_Code_Gen/main.cpp       |  8 ++++
 3 files changed, 89 insertions(+)
 create mode 100644 trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.cpp
 create mode 100644 trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.hh

diff --git a/trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.cpp b/trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.cpp
new file mode 100644
index 00000000..f753ab55
--- /dev/null
+++ b/trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.cpp
@@ -0,0 +1,39 @@
+#include <iostream>
+#include "clang/Basic/SourceManager.h"
+#include "clang/Basic/SourceLocation.h"
+
+#include "ICGDiagnosticConsumer.hh"
+#include "HeaderSearchDirs.hh"
+#include "Utilities.hh"
+
+
+ICGDiagnosticConsumer::ICGDiagnosticConsumer(llvm::raw_ostream &os, clang::DiagnosticOptions *diags, clang::CompilerInstance &in_ci, HeaderSearchDirs &in_hsd) 
+                                                : clang::TextDiagnosticPrinter(os, diags), ci(in_ci), hsd(in_hsd) {
+    error_in_user_code = false;
+};
+ICGDiagnosticConsumer::~ICGDiagnosticConsumer() {
+
+};
+
+
+/**
+ * @details
+ * -# Check the diagnostic level to see if an error is from user code. 
+ * -# Terminate the build if yes, continue otherwise.
+*/
+void ICGDiagnosticConsumer::HandleDiagnostic(clang::DiagnosticsEngine::Level DiagLevel, const clang::Diagnostic &Info) {
+    // Use TextDiagnosticPrinter to handle diagnostic if the code is user code.
+    // Otherwise use base DiagnosticConsumer to handle diagnostic for system code.
+    if (isInUserCode(ci , Info.getLocation(), hsd)) {
+        // Parent class implementation for handling diagnostic
+        clang::TextDiagnosticPrinter::HandleDiagnostic(DiagLevel, Info);
+
+        // Flag it if an error is from user code
+        if (DiagLevel == clang::DiagnosticsEngine::Level::Fatal || DiagLevel == clang::DiagnosticsEngine::Level::Error) {
+            error_in_user_code = true;
+        }
+    } else {
+        // Base class implementation for handling diagnostic
+        clang::DiagnosticConsumer::HandleDiagnostic(DiagLevel, Info);
+    }
+}
diff --git a/trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.hh b/trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.hh
new file mode 100644
index 00000000..92b7de70
--- /dev/null
+++ b/trick_source/codegen/Interface_Code_Gen/ICGDiagnosticConsumer.hh
@@ -0,0 +1,42 @@
+#ifndef ICG_DIAGNOSTICCONSUMER_HH
+#define ICG_DIAGNOSTICCONSUMER_HH
+
+#include "llvm/Support/raw_ostream.h"
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/TextDiagnosticPrinter.h"
+
+
+class HeaderSearchDirs;
+
+/**
+
+  This class is passed to the clang parser as a DiagnosticConsumer. 
+  It will terminate the trick build if an error found is from user code.
+
+  @date May 2024
+  
+ */
+
+
+class ICGDiagnosticConsumer : public clang::TextDiagnosticPrinter {
+public:
+    ICGDiagnosticConsumer(llvm::raw_ostream &os, clang::DiagnosticOptions *diags, clang::CompilerInstance &in_ci, HeaderSearchDirs &in_hsd);
+    ~ICGDiagnosticConsumer() override;
+
+    void HandleDiagnostic(clang::DiagnosticsEngine::Level DiagLevel,
+                          const clang::Diagnostic &Info) override;
+
+    /** Flag for if any error found in user code. */
+    bool error_in_user_code;
+
+protected:
+    /** The compiler instance. */
+    clang::CompilerInstance &ci ;
+
+    /** The header search directories. */
+    HeaderSearchDirs &hsd ;
+    
+};
+
+#endif
diff --git a/trick_source/codegen/Interface_Code_Gen/main.cpp b/trick_source/codegen/Interface_Code_Gen/main.cpp
index b2c4df79..84380329 100644
--- a/trick_source/codegen/Interface_Code_Gen/main.cpp
+++ b/trick_source/codegen/Interface_Code_Gen/main.cpp
@@ -25,6 +25,7 @@
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Parse/ParseAST.h"
 
+#include "ICGDiagnosticConsumer.hh"
 #include "ICGASTConsumer.hh"
 #include "HeaderSearchDirs.hh"
 #include "CommentSaver.hh"
@@ -323,6 +324,8 @@ int main(int argc, char * argv[]) {
 #else
     ci.getSourceManager().createMainFileID(fileEntry);
 #endif
+    ICGDiagnosticConsumer *icgDiagConsumer = new ICGDiagnosticConsumer(llvm::errs(), &ci.getDiagnosticOpts(), ci, hsd);
+    ci.getDiagnostics().setClient(icgDiagConsumer);
     ci.getDiagnosticClient().BeginSourceFile(ci.getLangOpts(), &ci.getPreprocessor());
     clang::ParseAST(ci.getSema());
     ci.getDiagnosticClient().EndSourceFile();
@@ -337,5 +340,10 @@ int main(int argc, char * argv[]) {
     // Print the list of headers that have the ICG:(No) comment
     printAttributes.printICGNoFiles();
 
+    if (icgDiagConsumer->error_in_user_code) {
+        std::cout << color(ERROR, "Trick build was terminated due to error in user code!") << std::endl;
+        exit(-1);
+    }
+
     return 0;
 }