[mlir][IR] Add `OpBuilder::attachScopedListener` and make `setListener` protected

[springerm]

Summary

OpBuilder::setListener is dangerous: users can accidentally replace a listener from an OpBuilder that was passed to them.

This revision makes setListener protected: it should only be used from the constructor of derived classes (typically rewriters that set up themselves as listeners).

This revision also adds a new function attachScopedListener that can be used to attach additional listeners in a safe way. Detaching existing listener (whether accidental or on purpose) is no longer possible with the public OpBuilder API.

Depends On: D155609

mlir/include/mlir/IR/Builders.h

 //===- Builders.h - Helpers for constructing MLIR Classes -------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILDERS_H
 #define MLIR_IR_BUILDERS_H
 
 #include "mlir/IR/Listeners.h"
 #include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Compiler.h"
 #include <optional>
 
 namespace mlir {
 
 class AffineExpr;
 class IRMapping;
 class UnknownLoc;
public:
   /// terminator.
   static OpBuilder atBlockTerminator(Block *block,
                                      RewriteListener *listener = nullptr) {
     auto *terminator = block->getTerminator();
     assert(terminator != nullptr && "the block has no terminator");
     return OpBuilder(block, Block::iterator(terminator), listener);
   }
 
-  /// Sets the listener of this builder to the one provided.
-  void setListener(RewriteListener *newListener) { listener = newListener; }
+  /// Attach the given listener to this builder and return an object that resets
+  /// the listener to the original state when the object is destructed at scope
+  /// exit.
+  auto attachScopedListener(RewriteListener *newListener) {
+    assert(newListener && "expected non-null listener");
+    RewriteListener *previousListener = getListener();
+    SmallVector<RewriteListener *> delegatedListeners({newListener});
+    if (previousListener)
+      delegatedListeners.push_back(previousListener);
+    // Forwarding listener is deleted when the returned object is destructed.
+    auto forwardingListener =
+        std::make_unique<ForwardingRewriteListener>(delegatedListeners);
+    setListener(forwardingListener.get());
+    return llvm::make_scope_exit(
+        [this, previousListener, u = move(forwardingListener)]() {
+          // Reset the listener to the original value.
+          this->setListener(previousListener);
+        });
+  }
 
   /// Returns the current listener of this builder, or nullptr if this builder
   /// doesn't have a listener.
   RewriteListener *getListener() const { return listener; }
 
   //===--------------------------------------------------------------------===//
   // Insertion Point Management
   //===--------------------------------------------------------------------===//
Operation *cloneWithoutRegions(Operation &op) {
     return insert(op.cloneWithoutRegions());
   }
   template <typename OpT>
   OpT cloneWithoutRegions(OpT op) {
     return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
   }
 
 protected:
+  /// Sets the listener of this builder to the one provided.
+  void setListener(RewriteListener *newListener) { listener = newListener; }
+
   /// The optional listener for events of this builder.
   RewriteListener *listener = nullptr;
 
 private:
   /// The current block this builder is inserting into.
   Block *block = nullptr;
   /// The insertion point within the block that this builder is inserting
   /// before.
   Block::iterator insertPoint;
 };
 
 } // namespace mlir
 
 #endif

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

void transform::BufferizeToAllocationOp::build(OpBuilder &b,
   resultTypes.push_back(b.getType<transform::AnyOpType>());
   return build(b, result,
                /*resultTypes=*/resultTypes,
                /*target=*/target,
                /*memorySpace=*/b.getI64IntegerAttr(memorySpace));
 }
 
 namespace {
-class NewOpsListener : public ForwardingRewriteListener {
+class NewOpsListener : public RewriteListener {
 public:
-  using ForwardingRewriteListener::ForwardingRewriteListener;
+  using RewriteListener::RewriteListener;
 
   SmallVector<Operation *> getNewOps() const {
     return SmallVector<Operation *>(newOps.begin(), newOps.end());
   }
 
 private:
   void notifyOperationInserted(Operation *op) override {
-    ForwardingRewriteListener::notifyOperationInserted(op);
+    RewriteListener::notifyOperationInserted(op);
     auto inserted = newOps.insert(op);
     (void)inserted;
     assert(inserted.second && "expected newly created op");
   }
 
   void notifyOperationRemoved(Operation *op) override {
-    ForwardingRewriteListener::notifyOperationRemoved(op);
+    RewriteListener::notifyOperationRemoved(op);
     op->walk([&](Operation *op) { newOps.erase(op); });
   }
 
   DenseSet<Operation *> newOps;
 };
 } // namespace
 
 DiagnosedSilenceableFailure transform::BufferizeToAllocationOp::apply(
     transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   // Attach listener to keep track of newly created ops.
-  RewriteListener *previousListener = rewriter.getListener();
-  auto resetListener =
-      llvm::make_scope_exit([&]() { rewriter.setListener(previousListener); });
-  NewOpsListener newOpsListener(previousListener);
-  rewriter.setListener(&newOpsListener);
+  NewOpsListener newOpsListener;
+  auto resetListener = rewriter.attachScopedListener(&newOpsListener);
 
   linalg::BufferizeToAllocationOptions options;
   if (getMemcpyOp() == "memref.tensor_store") {
     options.memcpyOp =
         linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefTensorStore;
   } else if (getMemcpyOp() == "memref.copy") {
     options.memcpyOp =
         linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefCopy;