1 files changed, 48 insertions, 5 deletions

diff --git a/clang-tools-extra/test/clang-tidy/checkers/llvm/use-new-mlir-op-builder.cpp b/clang-tools-extra/test/clang-tidy/checkers/llvm/use-new-mlir-op-builder.cpp
index b57eab0..c4a1d8d 100644
--- a/clang-tools-extra/test/clang-tidy/checkers/llvm/use-new-mlir-op-builder.cpp
+++ b/clang-tools-extra/test/clang-tidy/checkers/llvm/use-new-mlir-op-builder.cpp
@@ -2,6 +2,7 @@
 
 namespace mlir {
 class Location {};
+class Value {};
 class OpBuilder {
 public:
   template <typename OpTy, typename... Args>
@@ -28,6 +29,13 @@ struct NamedOp {
   static NamedOp create(OpBuilder &builder, Location location, const char* name) {
     return NamedOp(name);
   }
+  Value getResult() { return Value(); }
+};
+struct OperandOp {
+  OperandOp(Value val) {}
+  static OperandOp create(OpBuilder &builder, Location location, Value val) {
+    return OperandOp(val);
+  }
 };
 } // namespace mlir
 
@@ -40,6 +48,22 @@ void g(mlir::OpBuilder &b) {
   b.create<T>(b.getUnknownLoc(), "gaz");
 }
 
+class CustomBuilder : public mlir::ImplicitLocOpBuilder {
+public:
+  mlir::NamedOp f(const char *name) {
+    // CHECK-MESSAGES: :[[@LINE+2]]:12: warning: use 'OpType::create(builder, ...)'
+    // CHECK-FIXES: return mlir::NamedOp::create(*this, name);
+    return create<mlir::NamedOp>(name);
+  }
+
+  mlir::NamedOp g(const char *name) {
+    using mlir::NamedOp;
+    // CHECK-MESSAGES: :[[@LINE+2]]:12: warning: use 'OpType::create(builder, ...)'
+    // CHECK-FIXES: return NamedOp::create(*this, name);
+    return create<NamedOp>(name);
+  }
+};
+
 void f() {
   mlir::OpBuilder builder;
   // CHECK-MESSAGES: :[[@LINE+2]]:3: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
@@ -47,15 +71,18 @@ void f() {
   builder.create<mlir::  ModuleOp>(builder.getUnknownLoc());
 
   using mlir::NamedOp;
+  using mlir::OperandOp;
+
   // CHECK-MESSAGES: :[[@LINE+2]]:3: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
   // CHECK-FIXES: NamedOp::create(builder, builder.getUnknownLoc(), "baz");
   builder.create<NamedOp>(builder.getUnknownLoc(), "baz");
 
-  // CHECK-MESSAGES: :[[@LINE+3]]:3: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
-  // CHECK-FIXES: NamedOp::create(builder, builder.getUnknownLoc(),
-  // CHECK-FIXES:   "caz");
+  // CHECK-MESSAGES: :[[@LINE+4]]:3: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
+  // CHECK-FIXES: NamedOp::create(builder,
+  // CHECK-FIXES:      builder.getUnknownLoc(),
+  // CHECK-FIXES:      "caz");
   builder.
-   create<NamedOp>(
+   create<NamedOp>  (
      builder.getUnknownLoc(),
      "caz");
 
@@ -66,10 +93,26 @@ void f() {
 
   mlir::ImplicitLocOpBuilder ib;
   // CHECK-MESSAGES: :[[@LINE+2]]:3: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
-  // CHECK-FIXES: mlir::ModuleOp::create(ib);
+  // CHECK-FIXES: mlir::ModuleOp::create(ib );
   ib.create<mlir::ModuleOp>(   );
 
   // CHECK-MESSAGES: :[[@LINE+2]]:3: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
   // CHECK-FIXES: mlir::OpBuilder().create<mlir::ModuleOp>(builder.getUnknownLoc());
   mlir::OpBuilder().create<mlir::ModuleOp>(builder.getUnknownLoc());
+
+  auto *p = &builder;
+  // CHECK-MESSAGES: :[[@LINE+2]]:3: warning: use 'OpType::create(builder, ...)'
+  // CHECK-FIXES: NamedOp::create(*p, builder.getUnknownLoc(), "eaz");
+  p->create<NamedOp>(builder.getUnknownLoc(), "eaz");
+
+  CustomBuilder cb;
+  cb.f("faz");
+  cb.g("gaz");
+
+  // CHECK-FIXES:      OperandOp::create(builder, builder.getUnknownLoc(),
+  // CHECK-FIXES-NEXT:   NamedOp::create(builder, builder.getUnknownLoc(), "haz").getResult());
+  // CHECK-MESSAGES: :[[@LINE+2]]:3: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
+  // CHECK-MESSAGES: :[[@LINE+2]]:5: warning: use 'OpType::create(builder, ...)' instead of 'builder.create<OpType>(...)' [llvm-use-new-mlir-op-builder]
+  builder.create<OperandOp>(builder.getUnknownLoc(),
+    builder.create<NamedOp>(builder.getUnknownLoc(), "haz").getResult());
 }