Refactor NameLoc so that it also holds a child location. This removes the awkward use of CallSiteLoc as a variable usage location.

--

PiperOrigin-RevId: 248014642

mlir/include/mlir/IR/Location.h

@@ -53,10 +53,10 @@ public:
     /// This represents a file/line/column location.
     FileLineCol,
 
-    /// This represents an identity name, such as variable and function name.
+    /// This represents an identity name attached to a child location.
     Name,
 
-    /// This represents a location as call site or variable usage site. .
+    /// This represents a location as a call site.
     CallSite,
 
     // Represents a location as a 'void*' pointer to a front-end's opaque
@@ -170,17 +170,25 @@ public:
   static bool kindof(Kind kind) { return kind == Kind::FileLineCol; }
 };
 
-/// Represents an identity name. It is usually the callee of a CallLocation.
+/// Represents an identity name attached to a child location.
 class NameLoc : public Location {
 public:
   using ImplType = detail::NameLocationStorage;
   using Location::Location;
 
-  /// Return a uniqued name location object.
+  /// Return a uniqued name location object. The child location must not be
+  /// another NameLoc.
+  static NameLoc get(Identifier name, Location child, MLIRContext *context);
+
+  /// Return a uniqued name location object with an unknown child.
   static NameLoc get(Identifier name, MLIRContext *context);
 
+  /// Return the name identifier.
   Identifier getName() const;
 
+  /// Return the child location.
+  Location getChildLoc() const;
+
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool kindof(Kind kind) { return kind == Kind::Name; }
 };

mlir/lib/IR/AsmPrinter.cpp

@@ -397,7 +397,16 @@ void ModulePrinter::printLocationInternal(Location loc, bool pretty) {
     break;
   }
   case Location::Kind::Name: {
-    os << '\"' << loc.cast<NameLoc>().getName() << '\"';
+    auto nameLoc = loc.cast<NameLoc>();
+    os << '\"' << nameLoc.getName() << '\"';
+
+    // Print the child if it isn't unknown.
+    auto childLoc = nameLoc.getChildLoc();
+    if (!childLoc.isa<UnknownLoc>()) {
+      os << '(';
+      printLocationInternal(childLoc, pretty);
+      os << ')';
+    }
     break;
   }
   case Location::Kind::CallSite: {

mlir/lib/IR/Location.cpp

@@ -45,10 +45,20 @@ unsigned FileLineColLoc::getColumn() const {
 // NameLoc
 //===----------------------------------------------------------------------===//
 
+NameLoc NameLoc::get(Identifier name, MLIRContext *context) {
+  return get(name, UnknownLoc::get(context), context);
+}
+
+/// Return the name identifier.
 Identifier NameLoc::getName() const {
   return static_cast<ImplType *>(loc)->name;
 }
 
+/// Return the child location.
+Location NameLoc::getChildLoc() const {
+  return static_cast<ImplType *>(loc)->child;
+}
+
 //===----------------------------------------------------------------------===//
 // CallSiteLoc
 //===----------------------------------------------------------------------===//

mlir/lib/IR/LocationDetail.h

@@ -56,10 +56,11 @@ struct FileLineColLocationStorage : public LocationStorage {
 };
 
 struct NameLocationStorage : public LocationStorage {
-  NameLocationStorage(Identifier name)
-      : LocationStorage(Location::Kind::Name), name(name) {}
+  NameLocationStorage(Identifier name, Location child)
+      : LocationStorage(Location::Kind::Name), name(name), child(child) {}
 
   Identifier name;
+  Location child;
 };
 
 struct CallSiteLocationStorage : public LocationStorage {

mlir/lib/IR/MLIRContext.cpp

@@ -638,13 +638,15 @@ FileLineColLoc FileLineColLoc::get(UniquedFilename filename, unsigned line,
   });
 }
 
-NameLoc NameLoc::get(Identifier name, MLIRContext *context) {
+NameLoc NameLoc::get(Identifier name, Location child, MLIRContext *context) {
   auto &impl = context->getImpl();
+  assert(!child.isa<NameLoc>() &&
+         "a NameLoc cannot be used as a child of another NameLoc");
 
   // Safely get or create a location instance.
   return safeGetOrCreate(impl.nameLocs, name.data(), impl.locationMutex, [&] {
     return new (impl.locationAllocator.Allocate<NameLocationStorage>())
-        NameLocationStorage(name);
+        NameLocationStorage(name, child);
   });
 }
 

mlir/lib/Parser/Parser.cpp

@@ -1502,7 +1502,30 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
     }
 
     // Otherwise, this is a NameLoc.
-    *loc = NameLoc::get(Identifier::get(str, ctx), ctx);
+
+    // Check for a child location.
+    if (consumeIf(Token::l_paren)) {
+      auto childSourceLoc = getToken().getLoc();
+
+      // Parse the child location.
+      llvm::Optional<Location> childLoc;
+      if (parseLocationInstance(&childLoc))
+        return failure();
+
+      // The child must not be another NameLoc.
+      if (childLoc->isa<NameLoc>())
+        return emitError(childSourceLoc,
+                         "child of NameLoc cannot be another NameLoc");
+      *loc = NameLoc::get(Identifier::get(str, ctx), *childLoc, ctx);
+
+      // Parse the closing ')'.
+      if (parseToken(Token::r_paren,
+                     "expected ')' after child location of NameLoc"))
+        return failure();
+    } else {
+      *loc = NameLoc::get(Identifier::get(str, ctx), ctx);
+    }
+
     return success();
   }
 

mlir/test/IR/invalid-locations.mlir

@@ -23,6 +23,20 @@ func @location_invalid_instance() {
 
 // -----
 
+func @location_name_missing_r_paren() {
+^bb:
+  return loc("foo"(unknown]) // expected-error {{expected ')' after child location of NameLoc}}
+}
+
+// -----
+
+func @location_name_child_is_name() {
+^bb:
+  return loc("foo"("foo")) // expected-error {{child of NameLoc cannot be another NameLoc}}
+}
+
+// -----
+
 func @location_callsite_missing_l_paren() {
 ^bb:
   return loc(callsite unknown  // expected-error {{expected '(' in callsite location}}