Use Opaque location for handling deferred references to aliases internally to the parser (NFC)

[mehdi_amini]

Summary

This will allow to return to the client of parseLocationInstance a location that is resolved later.

[mehdi_amini]

Update to use OpaqueLoc instead of introduce a new location.

[mehdi_amini]

Replace the pointer in the OpaqueLoc with an index to the vector

[mehdi_amini]

Fix test

[rriddle]

LG, thanks.

mlir/lib/Parser/Parser.cpp

483

484

Why not StringRef here? All of the strings are coming from the file buffer right?

565–568

570–573

[mehdi_amini]

Address comments

mlir/include/mlir/IR/BuiltinLocationAttributes.td

AttrBuilderWithInferredContext<(ins "uintptr_t":$underlyingLocation,
       return $_get(fallbackLocation->getContext(), underlyingLocation,
                    underlyingTypeID, fallbackLocation);
     }]>
   ];
   let extraClassDeclaration = [{
     /// Returns an instance of opaque location which contains a given pointer to
     /// an object. The corresponding MLIR location is set to UnknownLoc.
     template <typename T>
-    static Location get(T underlyingLocation, MLIRContext *context);
+    static OpaqueLoc get(T underlyingLocation, MLIRContext *context);
 
     /// Returns an instance of opaque location which contains a given pointer to
     /// an object and an additional MLIR location.
     template <typename T>
-    static Location get(T underlyingLocation, Location fallbackLocation) {
+    static OpaqueLoc get(T underlyingLocation, Location fallbackLocation) {
       return get(reinterpret_cast<uintptr_t>(underlyingLocation),
                  TypeID::get<T>(), fallbackLocation);
     }
 
     /// Returns a pointer to some data structure that opaque location stores.
     template <typename T> static T getUnderlyingLocation(Location location) {
       assert(isa<T>(location));
       return reinterpret_cast<T>(

mlir/include/mlir/IR/Location.h

 
 //===----------------------------------------------------------------------===//
 // OpaqueLoc
 //===----------------------------------------------------------------------===//
 
 /// Returns an instance of opaque location which contains a given pointer to
 /// an object. The corresponding MLIR location is set to UnknownLoc.
 template <typename T>
-inline Location OpaqueLoc::get(T underlyingLocation, MLIRContext *context) {
+inline OpaqueLoc OpaqueLoc::get(T underlyingLocation, MLIRContext *context) {
   return get(reinterpret_cast<uintptr_t>(underlyingLocation), TypeID::get<T>(),
              UnknownLoc::get(context));
 }
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // LLVM Utilities
 //===----------------------------------------------------------------------===//

mlir/lib/Parser/Parser.cpp

private:
   /// block references and memorize them.
   SmallVector<DenseMap<StringRef, BlockDefinition>, 2> blocksByName;
   SmallVector<DenseMap<Block *, SMLoc>, 2> forwardRef;
 
   /// These are all of the placeholders we've made along with the location of
   /// their first reference, to allow checking for use of undefined values.
   DenseMap<Value, SMLoc> forwardRefPlaceholders;
 
-  /// A set of operations whose locations reference aliases that have yet to
-  /// be resolved.
-  SmallVector<std::pair<OpOrArgument, Token>, 8>
-      opsAndArgumentsWithDeferredLocs;
+  /// Deffered locations: when parsing `loc(#loc42)` we add an entry to this
+  /// map. After parsing the definition `#loc42 = ...` we'll patch back users
+  /// of this location.
+  struct DeferredLocInfo {
+    SMLoc loc;
+    StringRef identifier;
+  };
+  std::vector<DeferredLocInfo> deferredLocsReferences;
 
   /// The builder used when creating parsed operation instances.
   OpBuilder opBuilder;
 
   /// The top level operation that holds all of the parsed operations.
   Operation *topLevelOp;
 };
 } // namespace
for (const char *entry : errors) {
       auto loc = SMLoc::getFromPointer(entry);
       emitError(loc, "use of undeclared SSA value name");
     }
     return failure();
   }
 
   // Resolve the locations of any deferred operations.
   auto &attributeAliases = state.symbols.attributeAliasDefinitions;
-  for (std::pair<OpOrArgument, Token> &it : opsAndArgumentsWithDeferredLocs) {
-    llvm::SMLoc tokLoc = it.second.getLoc();
-    StringRef identifier = it.second.getSpelling().drop_front();
-    Attribute attr = attributeAliases.lookup(identifier);
+  auto locID = TypeID::get<DeferredLocInfo *>();
+  auto resolveLocation = [&](auto &opOrArgument) -> LogicalResult {
+    auto fwdLoc = opOrArgument.getLoc().template dyn_cast<OpaqueLoc>();
+    if (!fwdLoc || fwdLoc.getUnderlyingTypeID() != locID)
+      return success();
+    auto locInfo = deferredLocsReferences[fwdLoc.getUnderlyingLocation()];
+    Attribute attr = attributeAliases.lookup(locInfo.identifier);
     if (!attr)
-      return emitError(tokLoc) << "operation location alias was never defined";
-
-    LocationAttr locAttr = attr.dyn_cast<LocationAttr>();
+      return emitError(locInfo.loc)
+             << "operation location alias was never defined";
+    auto locAttr = attr.dyn_cast<LocationAttr>();
     if (!locAttr)
-      return emitError(tokLoc)
+      return emitError(locInfo.loc)
              << "expected location, but found '" << attr << "'";
-    auto opOrArgument = it.first;
-    if (auto *op = opOrArgument.dyn_cast<Operation *>())
-      op->setLoc(locAttr);
-    else
-      opOrArgument.get<BlockArgument>().setLoc(locAttr);
-  }
+    opOrArgument.setLoc(locAttr);
+    return success();
+  };
+
+  auto walkRes = topLevelOp->walk([&](Operation *op) {
+    if (failed(resolveLocation(*op)))
+      return WalkResult::interrupt();
+    for (Region &region : op->getRegions())
+      for (Block &block : region.getBlocks())
+        for (BlockArgument arg : block.getArguments())
+          if (failed(resolveLocation(arg)))
+            return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+  if (walkRes.wasInterrupted())
+    return failure();
 
   // Pop the top level name scope.
   if (failed(popSSANameScope()))
     return failure();
 
   // Verify that the parsed operations are valid.
   if (failed(verify(topLevelOp)))
     return failure();
if (tok.is(Token::hash_identifier)) {
     // If this alias can be resolved, do it now.
     Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier);
     if (attr) {
       if (!(directLoc = attr.dyn_cast<LocationAttr>()))
         return emitError(tok.getLoc())
                << "expected location, but found '" << attr << "'";
     } else {
       // Otherwise, remember this operation and resolve its location later.
-      opsAndArgumentsWithDeferredLocs.emplace_back(opOrArgument, tok);
+      // In the meantime, use a special OpaqueLoc as a marker.
+      directLoc = OpaqueLoc::get(deferredLocsReferences.size(),
+                                 TypeID::get<DeferredLocInfo *>(),
+                                 UnknownLoc::get(getContext()));
+      deferredLocsReferences.push_back(
+          DeferredLocInfo{tok.getLoc(), identifier});
     }
 
     // Otherwise, we parse the location directly.
   } else if (parseLocationInstance(directLoc)) {
     return failure();
   }
 
   if (parseToken(Token::r_paren, "expected ')' in location"))
     return failure();
 
-  if (directLoc) {
-    if (auto *op = opOrArgument.dyn_cast<Operation *>())
-      op->setLoc(directLoc);
-    else
-      opOrArgument.get<BlockArgument>().setLoc(directLoc);
-  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Region Parsing
 //===----------------------------------------------------------------------===//
 
 ParseResult OperationParser::parseRegion(