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

482

483

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

554–557

559–562

[mehdi_amini]

Address comments

mlir/include/mlir/IR/BuiltinLocationAttributes.td

let description = [{
     loc(?)
     ```
   }];
   let extraClassDeclaration = [{
     static UnknownLoc get(MLIRContext *context);
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// AliasLoc
+//===----------------------------------------------------------------------===//
+
+def FwdAliasLoc : Builtin_LocationAttr<"FwdAliasLoc"> {
+  let summary = "A location private to the parser, used to resolve aliases.";
+  let description = [{
+  }];
+  let parameters = (ins
+    "StringAttr":$aliasIdentifier,
+    "const char *":$smLoc
+  );
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "StringAttr":$name, "const char *":$smLoc), [{
+      return $_get(name.getContext(), name, smLoc);
+    }]>
+  ];
+}
+
+
 #endif // BUILTIN_LOCATION_ATTRIBUTES_TD

mlir/lib/IR/Location.cpp

return TypeSwitch<LocationAttr, WalkResult>(*this)
       .Case([&](OpaqueLoc opaqueLoc) -> WalkResult {
         return opaqueLoc.getFallbackLocation()->walk(walkFn);
       })
       .Default(WalkResult::advance());
 }
 
 /// Methods for support type inquiry through isa, cast, and dyn_cast.
 bool LocationAttr::classof(Attribute attr) {
-  return attr.isa<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, OpaqueLoc,
-                  UnknownLoc>();
+  return attr.isa<CallSiteLoc, FileLineColLoc, FusedLoc, FwdAliasLoc, NameLoc,
+                  OpaqueLoc, UnknownLoc>();
 }
 
 //===----------------------------------------------------------------------===//
 // CallSiteLoc
 //===----------------------------------------------------------------------===//
 
 CallSiteLoc CallSiteLoc::get(Location name, ArrayRef<Location> frames) {
   assert(!frames.empty() && "required at least 1 call frame");

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;
-
   /// The builder used when creating parsed operation instances.
   OpBuilder opBuilder;
 
   /// The top level operation that holds all of the parsed operations.
   Operation *topLevelOp;
 };
 } // namespace
 
 OperationParser::OperationParser(ParserState &state, ModuleOp topLevelOp)
     : Parser(state), opBuilder(topLevelOp.getRegion()), topLevelOp(topLevelOp) {
   // The top level operation starts a new name scope.
   pushSSANameScope(/*isIsolated=*/true);
 
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);
-    if (!attr)
-      return emitError(tokLoc) << "operation location alias was never defined";
+  auto getLocation = [&](OpOrArgument opOrArgument,
+                         FwdAliasLoc fwdLoc) -> LocationAttr {
+    LocationAttr locAttr;
+    SMLoc tokLoc = SMLoc::getFromPointer(fwdLoc.getSmLoc());
+    Attribute attr = attributeAliases.lookup(fwdLoc.getAliasIdentifier());
+    if (!attr) {
+      emitError(tokLoc) << "operation location alias was never defined";
+      return locAttr;
+    }
 
-    LocationAttr locAttr = attr.dyn_cast<LocationAttr>();
+    locAttr = attr.dyn_cast<LocationAttr>();
     if (!locAttr)
-      return emitError(tokLoc)
-             << "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);
-  }
+      emitError(tokLoc) << "expected location, but found '" << attr << "'";
+    return locAttr;
+  };
+
+  auto walkRes = topLevelOp->walk([&](Operation *op) {
+    auto fwdLoc = op->getLoc().dyn_cast<FwdAliasLoc>();
+    if (!fwdLoc)
+      return WalkResult::advance();
+    auto locAttr = getLocation(op, fwdLoc);
+    if (!locAttr)
+      return WalkResult::interrupt();
+    // FIXME: block arguments
+    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 FwdAliasLoc as a marker.
+      directLoc = FwdAliasLoc::get(StringAttr::get(getContext(), identifier),
+                                   tok.getLoc().getPointer());
     }
 
     // Otherwise, we parse the location directly.
   } else if (parseLocationInstance(directLoc)) {
     return failure();
   }
 
   if (parseToken(Token::r_paren, "expected ')' in location"))