[StripDebugInfoWithPred] Use worklist instead of recursion (#8435)

As it caused stack overflow it is changed to use worklist. The behavior is slightly different since it flattens fused loc and deduplicates locations.

Author: uenoku

lib/Transforms/StripDebugInfoWithPred.cpp

@@ -34,23 +34,23 @@ struct StripDebugInfoWithPred
 
   // Return stripped location for the given `loc`.
   mlir::Location getStrippedLoc(Location loc) {
-    // If `pred` return true, strip the location.
-    if (pred(loc))
-      return UnknownLoc::get(loc.getContext());
-
-    if (auto fusedLoc = dyn_cast<FusedLoc>(loc)) {
-      SmallVector<mlir::Location> newLocations;
-      newLocations.reserve(fusedLoc.getLocations().size());
-      for (auto loc : fusedLoc.getLocations())
-        newLocations.push_back(getStrippedLoc(loc));
-      // NOTE: Don't use FusedLoc::get(&getContext(), newLocations,
-      //       fusedLoc.getMetadata()) to avoid a bytecode reader bug
-      //       llvm-project#99626.
-      return FusedLoc::get(newLocations, fusedLoc.getMetadata(), &getContext());
+    SetVector<Location> newLocations;
+    SmallVector<Location, 8> worklist({loc});
+    while (!worklist.empty()) {
+      auto currentLoc = worklist.pop_back_val();
+      if (auto fusedLoc = dyn_cast<FusedLoc>(currentLoc)) {
+        for (auto loc : fusedLoc.getLocations())
+          worklist.push_back(loc);
+      } else {
+        if (!pred(currentLoc))
+          newLocations.insert(currentLoc);
+      }
     }
 
-    // TODO: Handle other loc type.
-    return loc;
+    // NOTE: Don't use FusedLoc::get(&getContext(), locations, metadata) to
+    // avoid a bytecode reader bug llvm-project#99626.
+    return FusedLoc::get(newLocations.getArrayRef(), Attribute(),
+                         &getContext());
   }
 
   void updateLocArray(Operation *op, StringRef attributeName) {

test/Transforms/strip-debuginfo-pred.mlir

@@ -9,7 +9,7 @@ func.func @inline_notation() {
 // CHECK: affine.for
 // CHECK-NEXT: loc("foo")
   affine.for %i0 = 0 to 8 {
-  } loc(fused["foo", "foo.txt":10:8]) 
+  } loc(fused["foo", "foo.txt":10:8, "foo"])
   return
 }