[mlir][python] Add debug helpersusers/nico/python-2

This PR lets users provide --debug-only flags with python decorators.
This greatly simplifies the debugging experience in python.

Co-authored-by: Tres <tpopp@users.noreply.github.com>

1 files changed, 23 insertions, 2 deletions

diff --git a/llvm/unittests/Support/DebugTest.cpp b/llvm/unittests/Support/DebugTest.cpp
index e8b7548..3874142 100644
--- a/llvm/unittests/Support/DebugTest.cpp
+++ b/llvm/unittests/Support/DebugTest.cpp
@@ -19,8 +19,8 @@ using namespace llvm;
 TEST(DebugTest, Basic) {
   std::string s1, s2;
   raw_string_ostream os1(s1), os2(s2);
-  static const char *DT[] = {"A", "B"};  
-  
+  static const char *DT[] = {"A", "B"};
+
   llvm::DebugFlag = true;
   setCurrentDebugTypes(DT, 2);
   DEBUG_WITH_TYPE("A", os1 << "A");
@@ -50,4 +50,25 @@ TEST(DebugTest, CommaInDebugBlock) {
   });
   EXPECT_EQ("ZYX", os1.str());
 }
+
+TEST(DebugTest, AppendAndPop) {
+  std::string s1, s2, s3;
+  raw_string_ostream os1(s1), os2(s2), os3(s3);
+
+  llvm::DebugFlag = true;
+  appendDebugType("A");
+  DEBUG_WITH_TYPE("A", os1 << "A");
+  DEBUG_WITH_TYPE("B", os1 << "B");
+  EXPECT_EQ("A", os1.str());
+
+  appendDebugType("B");
+  DEBUG_WITH_TYPE("A", os2 << "A");
+  DEBUG_WITH_TYPE("B", os2 << "B");
+  EXPECT_EQ("AB", os2.str());
+
+  popAppendedDebugTypes();
+  DEBUG_WITH_TYPE("A", os3 << "A");
+  DEBUG_WITH_TYPE("B", os3 << "B");
+  EXPECT_EQ("A", os3.str());
+}
 #endif