Add a breakpoint manager that matches based on File/Line/Col Locations

This will match the locations attached to the IRunits passed in as context
with an action.

This is a recommit of d09c80515d0e after fixing the test on Windows.

Differential Revision: https://reviews.llvm.org/D144815

2 files changed, 233 insertions, 0 deletions

diff --git a/mlir/unittests/Debug/CMakeLists.txt b/mlir/unittests/Debug/CMakeLists.txt
index 5ea18d2751de..59728bc819d9 100644
--- a/mlir/unittests/Debug/CMakeLists.txt
+++ b/mlir/unittests/Debug/CMakeLists.txt
@@ -1,6 +1,7 @@
 add_mlir_unittest(MLIRDebugTests
   DebugCounterTest.cpp
   ExecutionContextTest.cpp
+  FileLineColLocBreakpointManagerTest.cpp
 )
 
 target_link_libraries(MLIRDebugTests
diff --git a/mlir/unittests/Debug/FileLineColLocBreakpointManagerTest.cpp b/mlir/unittests/Debug/FileLineColLocBreakpointManagerTest.cpp
new file mode 100644
index 000000000000..e9cac3949ddb
--- /dev/null
+++ b/mlir/unittests/Debug/FileLineColLocBreakpointManagerTest.cpp
@@ -0,0 +1,232 @@
+//===- FileLineColLocBreakpointManagerTest.cpp - --------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Debug/BreakpointManagers/FileLineColLocBreakpointManager.h"
+#include "mlir/Debug/ExecutionContext.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/OperationSupport.h"
+#include "llvm/ADT/STLExtras.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::tracing;
+
+static Operation *createOp(MLIRContext *context, Location loc,
+                           StringRef operationName,
+                           unsigned int numRegions = 0) {
+  context->allowUnregisteredDialects();
+  return Operation::create(loc, OperationName(operationName, context),
+                           std::nullopt, std::nullopt, std::nullopt,
+                           std::nullopt, numRegions);
+}
+
+namespace {
+struct FileLineColLocTestingAction
+    : public ActionImpl<FileLineColLocTestingAction> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FileLineColLocTestingAction)
+  static constexpr StringLiteral tag = "file-line-col-loc-testing-action";
+  FileLineColLocTestingAction(ArrayRef<IRUnit> irUnits)
+      : ActionImpl<FileLineColLocTestingAction>(irUnits) {}
+};
+
+TEST(FileLineColLocBreakpointManager, OperationMatch) {
+  // This test will process a sequence of operation and check various situation
+  // with a breakpoint hitting or not based on the location attached to the
+  // operation. When a breakpoint hits, the action is skipped and the counter is
+  // not incremented.
+  ExecutionContext executionCtx(
+      [](const ActionActiveStack *) { return ExecutionContext::Skip; });
+  int counter = 0;
+  auto counterInc = [&]() { counter++; };
+
+  // Setup
+
+  MLIRContext context;
+  // Miscellaneous information to define operations
+  std::vector<StringRef> fileNames = {
+      StringRef("foo.bar"), StringRef("baz.qux"), StringRef("quux.corge")};
+  std::vector<std::pair<unsigned, unsigned>> lineColLoc = {{42, 7}, {24, 3}};
+  Location callee = UnknownLoc::get(&context),
+           caller = UnknownLoc::get(&context), loc = UnknownLoc::get(&context);
+
+  // Set of operations over where we are going to be testing the functionality
+  std::vector<Operation *> operations = {
+      createOp(&context, CallSiteLoc::get(callee, caller),
+               "callSiteLocOperation"),
+      createOp(&context,
+               FileLineColLoc::get(&context, fileNames[0], lineColLoc[0].first,
+                                   lineColLoc[0].second),
+               "fileLineColLocOperation"),
+      createOp(&context, FusedLoc::get(&context, {}, Attribute()),
+               "fusedLocOperation"),
+      createOp(&context, NameLoc::get(StringAttr::get(&context, fileNames[2])),
+               "nameLocOperation"),
+      createOp(&context, OpaqueLoc::get<void *>(nullptr, loc),
+               "opaqueLocOperation"),
+      createOp(&context,
+               FileLineColLoc::get(&context, fileNames[1], lineColLoc[1].first,
+                                   lineColLoc[1].second),
+               "anotherFileLineColLocOperation"),
+      createOp(&context, UnknownLoc::get(&context), "unknownLocOperation"),
+  };
+
+  FileLineColLocBreakpointManager breakpointManager;
+  executionCtx.addBreakpointManager(&breakpointManager);
+
+  // Test
+
+  // Basic case is that no breakpoint is set and the counter is incremented for
+  // every op.
+  auto checkNoMatch = [&]() {
+    counter = 0;
+    for (auto enumeratedOp : llvm::enumerate(operations)) {
+      executionCtx(counterInc,
+                   FileLineColLocTestingAction({enumeratedOp.value()}));
+      EXPECT_EQ(counter, static_cast<int>(enumeratedOp.index() + 1));
+    }
+  };
+  checkNoMatch();
+
+  // Set a breakpoint matching only the second operation in the list.
+  auto *breakpoint = breakpointManager.addBreakpoint(
+      fileNames[0], lineColLoc[0].first, lineColLoc[0].second);
+  auto checkMatchIdxs = [&](DenseSet<int> idxs) {
+    counter = 0;
+    int reference = 0;
+    for (int i = 0; i < (int)operations.size(); ++i) {
+      executionCtx(counterInc, FileLineColLocTestingAction({operations[i]}));
+      if (!idxs.contains(i))
+        reference++;
+      EXPECT_EQ(counter, reference);
+    }
+  };
+  checkMatchIdxs({1});
+
+  // Check that disabling the breakpoing brings us back to the original
+  // behavior.
+  breakpoint->disable();
+  checkNoMatch();
+
+  // Adding a breakpoint that won't match any location shouldn't affect the
+  // behavior.
+  breakpointManager.addBreakpoint(StringRef("random.file"), 3, 14);
+  checkNoMatch();
+
+  // Set a breakpoint matching only the fifth operation in the list.
+  breakpointManager.addBreakpoint(fileNames[1], lineColLoc[1].first,
+                                  lineColLoc[1].second);
+  counter = 0;
+  checkMatchIdxs({5});
+
+  // Re-enable the breakpoint matching only the second operation in the list.
+  // We now expect matching of operations 1 and 5.
+  breakpoint->enable();
+  checkMatchIdxs({1, 5});
+
+  for (auto *op : operations) {
+    op->destroy();
+  }
+}
+
+TEST(FileLineColLocBreakpointManager, BlockMatch) {
+  // This test will process a block and check various situation with
+  // a breakpoint hitting or not based on the location attached.
+  // When a breakpoint hits, the action is skipped and the counter is not
+  // incremented.
+  ExecutionContext executionCtx(
+      [](const ActionActiveStack *) { return ExecutionContext::Skip; });
+  int counter = 0;
+  auto counterInc = [&]() { counter++; };
+
+  // Setup
+
+  MLIRContext context;
+  std::vector<StringRef> fileNames = {StringRef("grault.garply"),
+                                      StringRef("waldo.fred")};
+  std::vector<std::pair<unsigned, unsigned>> lineColLoc = {{42, 7}, {24, 3}};
+  Operation *frontOp = createOp(&context,
+                                FileLineColLoc::get(&context, fileNames.front(),
+                                                    lineColLoc.front().first,
+                                                    lineColLoc.front().second),
+                                "firstOperation");
+  Operation *backOp = createOp(&context,
+                               FileLineColLoc::get(&context, fileNames.back(),
+                                                   lineColLoc.back().first,
+                                                   lineColLoc.back().second),
+                               "secondOperation");
+  Block block;
+  block.push_back(frontOp);
+  block.push_back(backOp);
+
+  FileLineColLocBreakpointManager breakpointManager;
+  executionCtx.addBreakpointManager(&breakpointManager);
+
+  // Test
+
+  executionCtx(counterInc, FileLineColLocTestingAction({&block}));
+  EXPECT_EQ(counter, 1);
+
+  auto *breakpoint = breakpointManager.addBreakpoint(
+      fileNames.front(), lineColLoc.front().first, lineColLoc.front().second);
+  counter = 0;
+  executionCtx(counterInc, FileLineColLocTestingAction({&block}));
+  EXPECT_EQ(counter, 0);
+  breakpoint->disable();
+  executionCtx(counterInc, FileLineColLocTestingAction({&block}));
+  EXPECT_EQ(counter, 1);
+
+  breakpoint = breakpointManager.addBreakpoint(
+      fileNames.back(), lineColLoc.back().first, lineColLoc.back().second);
+  counter = 0;
+  executionCtx(counterInc, FileLineColLocTestingAction({&block}));
+  EXPECT_EQ(counter, 0);
+  breakpoint->disable();
+  executionCtx(counterInc, FileLineColLocTestingAction({&block}));
+  EXPECT_EQ(counter, 1);
+}
+
+TEST(FileLineColLocBreakpointManager, RegionMatch) {
+  // This test will process a region and check various situation with
+  // a breakpoint hitting or not based on the location attached.
+  // When a breakpoint hits, the action is skipped and the counter is not
+  // incremented.
+  ExecutionContext executionCtx(
+      [](const ActionActiveStack *) { return ExecutionContext::Skip; });
+  int counter = 0;
+  auto counterInc = [&]() { counter++; };
+
+  // Setup
+
+  MLIRContext context;
+  StringRef fileName("plugh.xyzzy");
+  unsigned line = 42, col = 7;
+  Operation *containerOp =
+      createOp(&context, FileLineColLoc::get(&context, fileName, line, col),
+               "containerOperation", 1);
+  Region &region = containerOp->getRegion(0);
+
+  FileLineColLocBreakpointManager breakpointManager;
+  executionCtx.addBreakpointManager(&breakpointManager);
+
+  // Test
+  counter = 0;
+  executionCtx(counterInc, FileLineColLocTestingAction({&region}));
+  EXPECT_EQ(counter, 1);
+  auto *breakpoint = breakpointManager.addBreakpoint(fileName, line, col);
+  executionCtx(counterInc, FileLineColLocTestingAction({&region}));
+  EXPECT_EQ(counter, 1);
+  breakpoint->disable();
+  executionCtx(counterInc, FileLineColLocTestingAction({&region}));
+  EXPECT_EQ(counter, 2);
+
+  containerOp->destroy();
+}
+} // namespace