[mlir][IR] Support multiple listeners in ForwardingRewriteListener

[springerm]

Summary

Depends On: D155608

mlir/include/mlir/IR/Listeners.h

struct RewriteListener {
   /// display it to users.
   virtual LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback) {
     return failure();
   }
 };
 
-/// A listener that forwards all notifications to another listener. This
+/// A listener that forwards all notifications to one or more listeners. This
 /// struct can be used as a base to create listener chains, so that multiple
 /// listeners can be notified of IR changes.
 struct ForwardingRewriteListener : public RewriteListener {
-  ForwardingRewriteListener(RewriteListener *listener) : listener(listener) {}
+  ForwardingRewriteListener(ArrayRef<RewriteListener *> listeners)
+      : listeners(listeners) {
+#ifndef NDEBUG
+    for (RewriteListener *listener : listeners) {
+      assert(listener && "expected non-null listener");
+    }
+#endif // NDEBUG
+  }
 
   void notifyOperationInserted(Operation *op) override {
+    for (RewriteListener *listener : listeners)
       listener->notifyOperationInserted(op);
   }
   void notifyBlockCreated(Block *block) override {
+    for (RewriteListener *listener : listeners)
       listener->notifyBlockCreated(block);
   }
   void notifyOperationModified(Operation *op) override {
+    for (RewriteListener *listener : listeners)
       listener->notifyOperationModified(op);
   }
   void notifyOperationReplaced(Operation *op, Operation *newOp) override {
+    for (RewriteListener *listener : listeners)
       listener->notifyOperationReplaced(op, newOp);
   }
   void notifyOperationReplaced(Operation *op, ValueRange replacement) override {
+    for (RewriteListener *listener : listeners)
       listener->notifyOperationReplaced(op, replacement);
   }
   void notifyOperationRemoved(Operation *op) override {
+    for (RewriteListener *listener : listeners)
       listener->notifyOperationRemoved(op);
   }
   LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback) override {
-    return listener->notifyMatchFailure(loc, reasonCallback);
+    for (RewriteListener *listener : listeners)
+      if (succeeded(listener->notifyMatchFailure(loc, reasonCallback)))
+        return success();
+    return failure();
   }
 
 private:
-  RewriteListener *listener;
+  SmallVector<RewriteListener *> listeners;
 };
 } // namespace mlir
 
 #endif // MLIR_IR_LISTENERS_H