[mlir:SubElementsInterface] Add support for "skipping" when replacing attributes/types

[rriddle]

Summary

This is used to fix a bug in SymbolTable::replaceAllSymbolUses where we replace symbols that
we shouldn't.

[csigg]

Amazing! Thanks a lot River!

mlir/include/mlir/IR/SubElementInterfaces.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_SUBELEMENTINTERFACES_H
 #define MLIR_IR_SUBELEMENTINTERFACES_H
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Types.h"
+#include "mlir/IR/Visitors.h"
+
+namespace mlir {
+template <typename T>
+using SubElementReplFn = function_ref<T(T)>;
+template <typename T>
+using SubElementResultReplFn = function_ref<std::pair<T, WalkResult>(T)>;
+} // namespace mlir
 
 /// Include the definitions of the sub elemnt interfaces.
 #include "mlir/IR/SubElementAttrInterfaces.h.inc"
 #include "mlir/IR/SubElementTypeInterfaces.h.inc"
 
 #endif // MLIR_IR_SUBELEMENTINTERFACES_H

mlir/include/mlir/IR/SubElementInterfaces.td

class SubElementInterfaceBase<string interfaceName, string attrOrType,
   code extraClassDeclaration = [{
     /// Walk all of the held sub-attributes and sub-types.
     void walkSubElements(llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
                          llvm::function_ref<void(mlir::Type)> walkTypesFn);
 
     /// Recursively replace all of the nested sub-attributes and sub-types using the
     /// provided map functions. Returns nullptr in the case of failure.
     }] # attrOrType # [{ replaceSubElements(
-      llvm::function_ref<mlir::Attribute(mlir::Attribute)> replaceAttrFn,
-      llvm::function_ref<mlir::Type(mlir::Type)> replaceTypeFn
+      mlir::SubElementReplFn<mlir::Attribute> replaceAttrFn,
+      mlir::SubElementReplFn<mlir::Type> replaceTypeFn
+    ) {
+      return replaceSubElements(
+        [&](Attribute attr) { return std::make_pair(replaceAttrFn(attr), WalkResult::advance()); },
+        [&](Type type) { return std::make_pair(replaceTypeFn(type), WalkResult::advance()); }
+      );
+    }
+    /// Recursively replace all of the nested sub-attributes and sub-types using the
+    /// provided map functions. This variant allows for the map function to return an
+    /// additional walk result. Returns nullptr in the case of failure.
+    }] # attrOrType # [{ replaceSubElements(
+      llvm::function_ref<
+        std::pair<mlir::Attribute, mlir::WalkResult>(mlir::Attribute)> replaceAttrFn,
+      llvm::function_ref<
+        std::pair<mlir::Type, mlir::WalkResult>(mlir::Type)> replaceTypeFn
     );
   }];
   code extraTraitClassDeclaration = [{
     /// Walk all of the held sub-attributes and sub-types.
     void walkSubElements(llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
                          llvm::function_ref<void(mlir::Type)> walkTypesFn) {
       }] # interfaceName # " interface(" # derivedValue # [{);
       interface.walkSubElements(walkAttrsFn, walkTypesFn);
     }
 
     /// Recursively replace all of the nested sub-attributes and sub-types using the
     /// provided map functions. Returns nullptr in the case of failure.
     }] # attrOrType # [{ replaceSubElements(
-      llvm::function_ref<mlir::Attribute(mlir::Attribute)> replaceAttrFn,
-      llvm::function_ref<mlir::Type(mlir::Type)> replaceTypeFn) {
+      mlir::SubElementReplFn<mlir::Attribute> replaceAttrFn,
+      mlir::SubElementReplFn<mlir::Type> replaceTypeFn) {
       }] # interfaceName # " interface(" # derivedValue # [{);
       return interface.replaceSubElements(replaceAttrFn, replaceTypeFn);
     }
-
-    /// Recursively replace all of the nested sub-attributes and sub-types using the
-    /// provided map functions. Returns nullptr in the case of failure.
-    }] # attrOrType # [{ replaceImmediateSubElements(
-      llvm::ArrayRef<mlir::Attribute> replAttrs,
-      llvm::function_ref<mlir::Type(mlir::Type)> replTypes) {
-      return nullptr;
+    }] # attrOrType # [{ replaceSubElements(
+      mlir::SubElementResultReplFn<mlir::Attribute> replaceAttrFn,
+      mlir::SubElementResultReplFn<mlir::Type> replaceTypeFn) {
+      }] # interfaceName # " interface(" # derivedValue # [{);
+      return interface.replaceSubElements(replaceAttrFn, replaceTypeFn);
     }
   }];
   code extraSharedClassDeclaration = [{
     /// Walk all of the held sub-attributes.
     void walkSubAttrs(llvm::function_ref<void(mlir::Attribute)> walkFn) {
       walkSubElements(walkFn, /*walkTypesFn=*/[](mlir::Type) {});
     }
     /// Walk all of the held sub-types.
     void walkSubTypes(llvm::function_ref<void(mlir::Type)> walkFn) {
       walkSubElements(/*walkAttrsFn=*/[](mlir::Attribute) {}, walkFn);
     }
     
     /// Recursively replace all of the nested sub-attributes using the provided
     /// map function. Returns nullptr in the case of failure.
     }] # attrOrType # [{ replaceSubElements(
-      llvm::function_ref<mlir::Attribute(mlir::Attribute)> replaceAttrFn) {
+      mlir::SubElementReplFn<mlir::Attribute> replaceAttrFn) {
       return replaceSubElements(
         replaceAttrFn, [](mlir::Type type) { return type; });
     }
+    }] # attrOrType # [{ replaceSubElements(
+      mlir::SubElementResultReplFn<mlir::Attribute> replaceAttrFn) {
+      return replaceSubElements(
+        replaceAttrFn,
+        [](mlir::Type type) { return std::make_pair(type, WalkResult::advance()); }
+      );
+    }
     /// Recursively replace all of the nested sub-types using the provided map
     /// function. Returns nullptr in the case of failure.
     }] # attrOrType # [{ replaceSubElements(
-      llvm::function_ref<mlir::Type(mlir::Type)> replaceTypeFn) {
+      mlir::SubElementReplFn<mlir::Type> replaceTypeFn) {
       return replaceSubElements(
         [](mlir::Attribute attr) { return attr; }, replaceTypeFn);
     }
+    }] # attrOrType # [{ replaceSubElements(
+      mlir::SubElementResultReplFn<mlir::Type> replaceTypeFn) {
+      return replaceSubElements(
+        [](mlir::Attribute attr) { return std::make_pair(attr, WalkResult::advance()); },
+        replaceTypeFn
+      );
+    }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SubElementAttrInterface
 //===----------------------------------------------------------------------===//
 
 def SubElementAttrInterface

mlir/lib/IR/SubElementInterfaces.cpp

 static bool isMutable(Attribute attr) {
   return attr.hasTrait<AttributeTrait::IsMutable>();
 }
 static bool isMutable(Type type) {
   return type.hasTrait<TypeTrait::IsMutable>();
 }
 
 template <typename InterfaceT, typename T, typename ReplaceSubElementFnT>
-static void updateSubElementImpl(T element, function_ref<T(T)> walkFn,
-                                 DenseMap<T, T> &visited,
-                                 SmallVectorImpl<T> &newElements,
-                                 FailureOr<bool> &changed,
-                                 ReplaceSubElementFnT &&replaceSubElementFn) {
+static void updateSubElementImpl(
+    T element, function_ref<std::pair<T, WalkResult>(T)> walkFn,
+    DenseMap<T, T> &visited, SmallVectorImpl<T> &newElements,
+    FailureOr<bool> &changed, ReplaceSubElementFnT &&replaceSubElementFn) {
   // Bail early if we failed at any point.
   if (failed(changed))
     return;
   newElements.push_back(element);
 
   // Guard against potentially null inputs. We always map null to null.
   if (!element)
     return;
 
   // Check for an existing mapping for this element, and walk it if we haven't
   // yet.
   T &mappedElement = visited[element];
   if (!mappedElement) {
+    WalkResult result = WalkResult::advance();
+    std::tie(mappedElement, result) = walkFn(element);
+
     // Try walking this element.
-    if (!(mappedElement = walkFn(element))) {
+    if (result.wasInterrupted() || !mappedElement) {
       changed = failure();
       return;
     }
 
     // Handle replacing sub-elements if this element is also a container.
+    if (!result.wasSkipped()) {
       if (auto interface = mappedElement.template dyn_cast<InterfaceT>()) {
         if (!(mappedElement = replaceSubElementFn(interface))) {
           changed = failure();
           return;
         }
       }
     }
+  }
 
   // Update to the mapped element.
   if (mappedElement != element) {
     newElements.back() = mappedElement;
     changed = true;
   }
 }
 
 template <typename InterfaceT>
 static typename InterfaceT::ValueType
 replaceSubElementsImpl(InterfaceT interface,
-                       function_ref<Attribute(Attribute)> walkAttrsFn,
-                       function_ref<Type(Type)> walkTypesFn,
+                       SubElementResultReplFn<Attribute> walkAttrsFn,
+                       SubElementResultReplFn<Type> walkTypesFn,
                        DenseMap<Attribute, Attribute> &visitedAttrs,
                        DenseMap<Type, Type> &visitedTypes) {
   // Walk the current sub-elements, replacing them as necessary.
   SmallVector<Attribute, 16> newAttrs;
   SmallVector<Type, 16> newTypes;
   FailureOr<bool> changed = false;
   auto replaceSubElementFn = [&](auto subInterface) {
     return replaceSubElementsImpl(subInterface, walkAttrsFn, walkTypesFn,
replaceSubElementsImpl(InterfaceT interface,
   if (isMutable(interface))
     return {};
 
   // Use the new elements during the replacement.
   return interface.replaceImmediateSubElements(newAttrs, newTypes);
 }
 
 Attribute SubElementAttrInterface::replaceSubElements(
-    function_ref<Attribute(Attribute)> replaceAttrFn,
-    function_ref<Type(Type)> replaceTypeFn) {
+    SubElementResultReplFn<Attribute> replaceAttrFn,
+    SubElementResultReplFn<Type> replaceTypeFn) {
   assert(replaceAttrFn && replaceTypeFn && "expected valid replace functions");
   DenseMap<Attribute, Attribute> visitedAttrs;
   DenseMap<Type, Type> visitedTypes;
   return replaceSubElementsImpl(*this, replaceAttrFn, replaceTypeFn,
                                 visitedAttrs, visitedTypes);
 }
 
 Type SubElementTypeInterface::replaceSubElements(
-    function_ref<Attribute(Attribute)> replaceAttrFn,
-    function_ref<Type(Type)> replaceTypeFn) {
+    SubElementResultReplFn<Attribute> replaceAttrFn,
+    SubElementResultReplFn<Type> replaceTypeFn) {
   assert(replaceAttrFn && replaceTypeFn && "expected valid replace functions");
   DenseMap<Attribute, Attribute> visitedAttrs;
   DenseMap<Type, Type> visitedTypes;
   return replaceSubElementsImpl(*this, replaceAttrFn, replaceTypeFn,
                                 visitedAttrs, visitedTypes);
 }
 
 //===----------------------------------------------------------------------===//
 // SubElementInterface Tablegen definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SubElementAttrInterfaces.cpp.inc"
 #include "mlir/IR/SubElementTypeInterfaces.cpp.inc"

mlir/lib/IR/SymbolTable.cpp

 replaceAllSymbolUsesImpl(SymbolT symbol, StringAttr newSymbol, IRUnitT *limit) {
   // Generate a new attribute to replace the given attribute.
   FlatSymbolRefAttr newLeafAttr = FlatSymbolRefAttr::get(newSymbol);
   for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
     SymbolRefAttr oldAttr = scope.symbol;
     SymbolRefAttr newAttr = generateNewRefAttr(scope.symbol, newLeafAttr);
 
     auto walkFn = [&](Operation *op) -> Optional<WalkResult> {
-      auto remapAttrFn = [&](Attribute attr) -> Attribute {
+      auto remapAttrFn =
+          [&](Attribute attr) -> std::pair<Attribute, WalkResult> {
+        // Regardless of the match, don't walk nested SymbolRefAttrs, we don't
+        // want to accidentally replace an inner reference.
         if (attr == oldAttr)
-          return newAttr;
+          return {newAttr, WalkResult::skip()};
         // Handle prefix matches.
         if (SymbolRefAttr symRef = attr.dyn_cast<SymbolRefAttr>()) {
           if (isReferencePrefixOf(oldAttr, symRef)) {
             auto oldNestedRefs = oldAttr.getNestedReferences();
             auto nestedRefs = symRef.getNestedReferences();
             if (oldNestedRefs.empty())
-              return SymbolRefAttr::get(newSymbol, nestedRefs);
+              return {SymbolRefAttr::get(newSymbol, nestedRefs),
+                      WalkResult::skip()};
 
             auto newNestedRefs = llvm::to_vector<4>(nestedRefs);
             newNestedRefs[oldNestedRefs.size() - 1] = newLeafAttr;
-            return SymbolRefAttr::get(symRef.getRootReference(), newNestedRefs);
+            return {
+                SymbolRefAttr::get(symRef.getRootReference(), newNestedRefs),
+                WalkResult::skip()};
           }
+          return {attr, WalkResult::skip()};
         }
-        return attr;
+        return {attr, WalkResult::advance()};
       };
       // Generate a new attribute dictionary by replacing references to the old
       // symbol.
       auto newDict = op->getAttrDictionary().replaceSubElements(remapAttrFn);
       if (!newDict)
         return WalkResult::interrupt();
 
       op->setAttrs(newDict.template cast<DictionaryAttr>());

mlir/test/IR/test-symbol-rauw.mlir

func.func @symbol_bar() {
       // CHECK-SAME: z_non_symbol_attr_3
       "foo.op"() {
         non_symbol_attr,
         use = [#test.sub_elements_access<[@symbol_foo],@symbol_bar,@symbol_foo>],
         z_non_symbol_attr_3
       } : () -> ()
     }
 }
+
+// -----
+
+module {
+  // CHECK: module @replaced_foo
+  module @foo attributes {sym.new_name = "replaced_foo" } {
+    // CHECK: func.func private @foo
+    func.func private @foo()
+  }
+
+  // CHECK: foo.op
+  // CHECK-SAME: use = @replaced_foo::@foo
+  "foo.op"() {
+    use = @foo::@foo
+  } : () -> ()
+}