[mlir] Implement the SubElement interfaces for the builtin locations

[rriddle]

Summary

This enables find/replace of nested components for location attributes.

Depends on D136407

[Mogball]

mlir/include/mlir/IR/BuiltinLocationAttributes.td

32

Could this be added to the base class if they all implement it?

mlir/lib/IR/Location.cpp

152

Can you spell out this auto

[rriddle]

mlir/include/mlir/IR/BuiltinLocationAttributes.td

32

Not all of them implement it (e.g. FileLineColLoc, UnknownLoc).

mlir/include/mlir/IR/BuiltinLocationAttributes.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BUILTIN_LOCATION_ATTRIBUTES_TD
 #define BUILTIN_LOCATION_ATTRIBUTES_TD
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinDialect.td"
+include "mlir/IR/SubElementInterfaces.td"
 
 // Base class for Builtin dialect location attributes.
-class Builtin_LocationAttr<string name>
-    : AttrDef<Builtin_Dialect, name, [], "::mlir::LocationAttr"> {
+class Builtin_LocationAttr<string name, list<Trait> traits = []>
+    : AttrDef<Builtin_Dialect, name, traits, "::mlir::LocationAttr"> {
   let cppClassName = name;
   let mnemonic = ?;
 }
 
 //===----------------------------------------------------------------------===//
 // CallSiteLoc
 //===----------------------------------------------------------------------===//
 
-def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc"> {
+def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "A callsite source location";
   let description = [{
     Syntax:
 
     ```
     callsite-location ::= `callsite` `(` location `at` location `)`
     ```
 
def FileLineColLoc : Builtin_LocationAttr<"FileLineColLoc"> {
   ];
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FusedLoc
 //===----------------------------------------------------------------------===//
 
-def FusedLoc : Builtin_LocationAttr<"FusedLoc"> {
+def FusedLoc : Builtin_LocationAttr<"FusedLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "A tuple of other source locations";
   let description = [{
     Syntax:
 
     ```
     fused-location ::= `fused` fusion-metadata? `[` location (location `,`)* `]`
     fusion-metadata ::= `<` attribute-value `>`
     ```
let extraClassDeclaration = [{
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // NameLoc
 //===----------------------------------------------------------------------===//
 
-def NameLoc : Builtin_LocationAttr<"NameLoc"> {
+def NameLoc : Builtin_LocationAttr<"NameLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "A named source location";
   let description = [{
     Syntax:
 
     ```
     name-location ::= string-literal (`(` location `)`)?
     ```
 
def NameLoc : Builtin_LocationAttr<"NameLoc", [
   ];
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // OpaqueLoc
 //===----------------------------------------------------------------------===//
 
-def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc"> {
+def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "An opaque source location";
   let description = [{
     An instance of this location essentially contains a pointer to some data
     structure that is external to MLIR and an optional location that can be used
     if the first one is not suitable. Since it contains an external structure,
     only the optional location is used during serialization.
   }];
   let parameters = (ins "uintptr_t":$underlyingLocation,

mlir/include/mlir/IR/Location.h

 // location position information.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_LOCATION_H
 #define MLIR_IR_LOCATION_H
 
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/SubElementInterfaces.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 
 class Location;
 class WalkResult;
 
 //===----------------------------------------------------------------------===//

mlir/lib/IR/Location.cpp

 CallSiteLoc CallSiteLoc::get(Location name, ArrayRef<Location> frames) {
   assert(!frames.empty() && "required at least 1 call frame");
   Location caller = frames.back();
   for (auto frame : llvm::reverse(frames.drop_back()))
     caller = CallSiteLoc::get(frame, caller);
   return CallSiteLoc::get(name, caller);
 }
 
+void CallSiteLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkAttrsFn(getCallee());
+  walkAttrsFn(getCaller());
+}
+
+Attribute
+CallSiteLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                         ArrayRef<Type> replTypes) const {
+  return get(replAttrs[0].cast<LocationAttr>(),
+             replAttrs[1].cast<LocationAttr>());
+}
+
 //===----------------------------------------------------------------------===//
 // FusedLoc
 //===----------------------------------------------------------------------===//
 
 Location FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
                        MLIRContext *context) {
   // Unique the set of locations to be fused.
   llvm::SmallSetVector<Location, 4> decomposedLocs;
if (locs.empty()) {
     return Base::get(context, ArrayRef<Location>{UnknownLoc::get(context)},
                      metadata);
   }
   if (locs.size() == 1 && !metadata)
     return locs.front();
 
   return Base::get(context, locs, metadata);
 }
+
+void FusedLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  for (Attribute attr : getLocations())
+    walkAttrsFn(attr);
+  walkAttrsFn(getMetadata());
+}
+
+Attribute
+FusedLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                      ArrayRef<Type> replTypes) const {
+  SmallVector<Location> newLocs;
+  newLocs.reserve(replAttrs.size() - 1);
+  for (Attribute attr : replAttrs.drop_back())
+    newLocs.push_back(attr.cast<LocationAttr>());
+  return get(getContext(), newLocs, replAttrs.back());
+}
+
+//===----------------------------------------------------------------------===//
+// NameLoc
+//===----------------------------------------------------------------------===//
+
+void NameLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkAttrsFn(getName());
+  walkAttrsFn(getChildLoc());
+}
+
+Attribute NameLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                               ArrayRef<Type> replTypes) const {
+  return get(replAttrs[0].cast<StringAttr>(),
+             replAttrs[1].cast<LocationAttr>());
+}
+
+//===----------------------------------------------------------------------===//
+// OpaqueLoc
+//===----------------------------------------------------------------------===//
+
+void OpaqueLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkAttrsFn(getFallbackLocation());
+}
+
+Attribute
+OpaqueLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                       ArrayRef<Type> replTypes) const {
+  return get(getUnderlyingLocation(), getUnderlyingTypeID(),
+             replAttrs[0].cast<LocationAttr>());
+}