From da49cf8c877115c8139ba7b93f5dbd6e74a20d4a Mon Sep 17 00:00:00 2001
From: Jan Rehders <jan.rehders@amd.com>
Date: Thu, 21 Dec 2023 16:52:57 +0100
Subject: [PATCH] Allow replacing instrs while visiting

This reduces the need to keep around lists of instructions to be replaced later
until the visitor finished and will avoid some memory allocations.
---
 lib/Dialect/Visitor.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/lib/Dialect/Visitor.cpp b/lib/Dialect/Visitor.cpp
index 2096be0..ec5daac 100644
--- a/lib/Dialect/Visitor.cpp
+++ b/lib/Dialect/Visitor.cpp
@@ -244,7 +244,7 @@ void VisitorBase::visitByDeclarations(void *payload, llvm::Module &module,
       continue;
     }
 
-    for (Use &use : decl.uses()) {
+    for (Use &use : make_early_inc_range(decl.uses())) {
       if (auto *inst = dyn_cast<Instruction>(use.getUser())) {
         if (!filter(*inst))
           continue;
@@ -260,7 +260,7 @@ void VisitorBase::visitByDeclarations(void *payload, llvm::Module &module,
 void VisitorBase::visit(void *payload, Function &fn) const {
   if (m_strategy == VisitorStrategy::ByInstruction) {
     for (BasicBlock &bb : fn) {
-      for (Instruction &inst : bb)
+      for (Instruction &inst : make_early_inc_range(bb))
         visit(payload, inst);
     }
     return;
@@ -269,7 +269,7 @@ void VisitorBase::visit(void *payload, Function &fn) const {
   if (m_strategy == VisitorStrategy::ReversePostOrder) {
     ReversePostOrderTraversal<Function *> rpot(&fn);
     for (BasicBlock *bb : rpot) {
-      for (Instruction &inst : *bb)
+      for (Instruction &inst : make_early_inc_range(*bb))
         visit(payload, inst);
     }
     return;