[LLParser] Handle mixed blockaddress forward references with names and IDs

[MichielDerhaeg]

Summary

ValIDs of the LocalID and LocalName kind were mixed in the same std::map.
std::less is not defined in that case. Operations on the std::map were
undefined and the incorrect basic block is parsed.

[MichielDerhaeg]

This was a bug reported a while ago on the old bugzilla: https://bugs.llvm.org/show_bug.cgi?id=46183

[dexonsmith]

Would it be simpler to define < for that case? Or, we could define a custom one to send to the std::map:

struct CompareValID {
  static bool operator<(const ValID &LHS, const ValID &RHS) {
    if (LHS.Kind != RHS.Kind)
      return LHS.Kind < RHS.Kind;
    return LHS < RHS;
  }
};

[MichielDerhaeg]

In D137164#3915823, @dexonsmith wrote:

Would it be simpler to define < for that case? Or, we could define a custom one to send to the std::map:

struct CompareValID {
  static bool operator<(const ValID &LHS, const ValID &RHS) {
    if (LHS.Kind != RHS.Kind)
      return LHS.Kind < RHS.Kind;
    return LHS < RHS;
  }
};

That was my initial idea, but if you search for other ForwardRef-related datastructures in LLParser.cpp, you'll see that they solved it this way for all other cases as well.
I did it like this to remain consistent.

[dexonsmith]

In D137164#3917926, @MichielDerhaeg wrote:

In D137164#3915823, @dexonsmith wrote:

Would it be simpler to define < for that case? Or, we could define a custom one to send to the std::map:

struct CompareValID {
  static bool operator<(const ValID &LHS, const ValID &RHS) {
    if (LHS.Kind != RHS.Kind)
      return LHS.Kind < RHS.Kind;
    return LHS < RHS;
  }
};

That was my initial idea, but if you search for other ForwardRef-related datastructures in LLParser.cpp, you'll see that they solved it this way for all other cases as well.
I did it like this to remain consistent.

Hmm. Okay, the pattern followed elsewhere is reasonable. LGTM then.

That said, the pattern of splitting the map seems to add boilerplate and I'm not convinced it's less error-prone than a shared map. @dblaikie, any thoughts? (I.e., would it be better to tidy it up, or is there some value in using the split maps?)

[dblaikie]

In D137164#3918030, @dexonsmith wrote:

In D137164#3917926, @MichielDerhaeg wrote:

In D137164#3915823, @dexonsmith wrote:

Would it be simpler to define < for that case? Or, we could define a custom one to send to the std::map:

struct CompareValID {
  static bool operator<(const ValID &LHS, const ValID &RHS) {
    if (LHS.Kind != RHS.Kind)
      return LHS.Kind < RHS.Kind;
    return LHS < RHS;
  }
};

That was my initial idea, but if you search for other ForwardRef-related datastructures in LLParser.cpp, you'll see that they solved it this way for all other cases as well.
I did it like this to remain consistent.

Hmm. Okay, the pattern followed elsewhere is reasonable. LGTM then.

That said, the pattern of splitting the map seems to add boilerplate and I'm not convinced it's less error-prone than a shared map. @dblaikie, any thoughts? (I.e., would it be better to tidy it up, or is there some value in using the split maps?)

Yeah, at least at a glance, and with your framing here - I tend to agree that all this duplication seems unfortunate & I guess, by the sounds of it, there's more of this duplication pre-existing in other instances that could be cleaned up if we generalized the op<?

llvm/include/llvm/AsmParser/LLParser.h

private:
     std::map<unsigned, std::pair<GlobalValue*, LocTy> > ForwardRefValIDs;
     std::vector<GlobalValue*> NumberedVals;
 
     // Comdat forward reference information.
     std::map<std::string, LocTy> ForwardRefComdats;
 
     // References to blockaddress.  The key is the function ValID, the value is
     // a list of references to blocks in that function.
-    std::map<ValID, std::map<ValID, GlobalValue *>> ForwardRefBlockAddresses;
+    std::map<ValID, std::map<ValID, GlobalValue *>> ForwardRefBlockAddressNames;
+    std::map<ValID, std::map<ValID, GlobalValue *>> ForwardRefBlockAddressIDs;
+
     class PerFunctionState;
     /// Reference to per-function state to allow basic blocks to be
     /// forward-referenced by blockaddress instructions within the same
     /// function.
     PerFunctionState *BlockAddressPFS;
 
     // References to dso_local_equivalent. The key is the global's ValID, the
     // value is a placeholder value that will be replaced. Note there are two

llvm/lib/AsmParser/LLParser.cpp

if (Function *Fn = dyn_cast<Function>(V)) {
       AttrBuilder Attrs(M->getContext(), GV->getAttributes());
       Attrs.merge(B);
       GV->setAttributes(AttributeSet::get(Context,Attrs));
     } else {
       llvm_unreachable("invalid object with forward attribute group reference");
     }
   }
 
-  // If there are entries in ForwardRefBlockAddresses at this point, the
-  // function was never defined.
-  if (!ForwardRefBlockAddresses.empty())
-    return error(ForwardRefBlockAddresses.begin()->first.Loc,
+  // If there are entries in ForwardRefBlockAddress{Names,IDs} at this
+  // point, the function was never defined.
+  if (!ForwardRefBlockAddressNames.empty())
+    return error(ForwardRefBlockAddressNames.begin()->first.Loc,
+                 "expected function name in blockaddress");
+  if (!ForwardRefBlockAddressIDs.empty())
+    return error(ForwardRefBlockAddressIDs.begin()->first.Loc,
                  "expected function name in blockaddress");
 
   auto ResolveForwardRefDSOLocalEquivalents = [&](const ValID &GVRef,
                                                   GlobalValue *FwdRef) {
     GlobalValue *GV = nullptr;
     if (GVRef.Kind == ValID::t_GlobalName) {
       GV = M->getNamedValue(GVRef.StrVal);
     } else if (GVRef.UIntVal < NumberedVals.size()) {
if (GV) {
       if (!isa<Function>(GV))
         return error(Fn.Loc, "expected function name in blockaddress");
       F = cast<Function>(GV);
       if (F->isDeclaration())
         return error(Fn.Loc, "cannot take blockaddress inside a declaration");
     }
 
     if (!F) {
+      auto &forwardRefBlockAddresses = Label.Kind == ValID::t_LocalName
+                                           ? ForwardRefBlockAddressNames
+                                           : ForwardRefBlockAddressIDs;
+
       // Make a global variable as a placeholder for this reference.
       GlobalValue *&FwdRef =
-          ForwardRefBlockAddresses.insert(std::make_pair(
-                                              std::move(Fn),
-                                              std::map<ValID, GlobalValue *>()))
+          forwardRefBlockAddresses
+              .insert(std::make_pair(std::move(Fn),
               .first->second.insert(std::make_pair(std::move(Label), nullptr))
               .first->second;
       if (!FwdRef) {
         unsigned FwdDeclAS;
         if (ExpectedTy) {
           // If we know the type that the blockaddress is being assigned to,
           // we can use the address space of that type.
           if (!ExpectedTy->isPointerTy())
bool LLParser::parseFunctionHeader(Function *&Fn, bool IsDefine) {
   ValID ID;
   if (FunctionName.empty()) {
     ID.Kind = ValID::t_GlobalID;
     ID.UIntVal = NumberedVals.size() - 1;
   } else {
     ID.Kind = ValID::t_GlobalName;
     ID.StrVal = FunctionName;
   }
-  auto Blocks = ForwardRefBlockAddresses.find(ID);
-  if (Blocks != ForwardRefBlockAddresses.end())
+
+  auto Blocks = ForwardRefBlockAddressNames.find(ID);
+  if (Blocks != ForwardRefBlockAddressNames.end())
+    return error(Blocks->first.Loc,
+                 "cannot take blockaddress inside a declaration");
+  Blocks = ForwardRefBlockAddressIDs.find(ID);
+  if (Blocks != ForwardRefBlockAddressIDs.end())
     return error(Blocks->first.Loc,
                  "cannot take blockaddress inside a declaration");
+
   return false;
 }
 
 bool LLParser::PerFunctionState::resolveForwardRefBlockAddresses() {
   ValID ID;
   if (FunctionNumber == -1) {
     ID.Kind = ValID::t_GlobalName;
     ID.StrVal = std::string(F.getName());
   } else {
     ID.Kind = ValID::t_GlobalID;
     ID.UIntVal = FunctionNumber;
   }
 
-  auto Blocks = P.ForwardRefBlockAddresses.find(ID);
-  if (Blocks == P.ForwardRefBlockAddresses.end())
+  auto ResolveForwardRefBlockAddresses = [&](auto &forwardRefBlockAddresses) {
+    auto Blocks = forwardRefBlockAddresses.find(ID);
+    if (Blocks == forwardRefBlockAddresses.end())
       return false;
 
     for (const auto &I : Blocks->second) {
       const ValID &BBID = I.first;
       GlobalValue *GV = I.second;
 
-    assert((BBID.Kind == ValID::t_LocalID || BBID.Kind == ValID::t_LocalName) &&
-           "Expected local id or name");
+      assert(
+          (BBID.Kind == ValID::t_LocalID || BBID.Kind == ValID::t_LocalName) &&
       BasicBlock *BB;
       if (BBID.Kind == ValID::t_LocalName)
         BB = getBB(BBID.StrVal, BBID.Loc);
       else
         BB = getBB(BBID.UIntVal, BBID.Loc);
       if (!BB)
         return P.error(BBID.Loc, "referenced value is not a basic block");
 
       Value *ResolvedVal = BlockAddress::get(&F, BB);
-    ResolvedVal = P.checkValidVariableType(BBID.Loc, BBID.StrVal, GV->getType(),
-                                           ResolvedVal);
+      ResolvedVal = P.checkValidVariableType(BBID.Loc, BBID.StrVal,
+                                             GV->getType(), ResolvedVal);
       if (!ResolvedVal)
         return true;
       GV->replaceAllUsesWith(ResolvedVal);
       GV->eraseFromParent();
     }
 
-  P.ForwardRefBlockAddresses.erase(Blocks);
+    forwardRefBlockAddresses.erase(Blocks);
     return false;
+  };
+
+  return ResolveForwardRefBlockAddresses(P.ForwardRefBlockAddressNames) ||
+         ResolveForwardRefBlockAddresses(P.ForwardRefBlockAddressIDs);
 }
 
 /// parseFunctionBody
 ///   ::= '{' BasicBlock+ UseListOrderDirective* '}'
 bool LLParser::parseFunctionBody(Function &Fn) {
   if (Lex.getKind() != lltok::lbrace)
     return tokError("expected '{' in function body");
   Lex.Lex();  // eat the {.

llvm/test/Bitcode/blockaddress-forwardref.ll

+; RUN: llvm-as < %s | llvm-dis | FileCheck %s
+; RUN: verify-uselistorder < %s
+
+; CHECK: i8* blockaddress(@f, %0), i8* blockaddress(@f, %bb)
+@g = global [2 x i8*] [i8* blockaddress(@f, %0), i8* blockaddress(@f, %bb)]
+
+define void @f() {
+entry:
+  br label %bb
+bb:
+  br label %0
+; %0
+  ret void
+}