[AttrBuilder] Make handling of type attributes more generic (NFCI)

While working on the elementtype attribute, I felt that the type
attribute handling in AttrBuilder is overly repetitive. This patch
converts the separate Type* members into an std::array<Type*>, so
that all type attribute kinds can be handled generically.

There's more room for improvement here (especially when it comes to
converting the AttrBuilder to an Attribute), but this seems like a
good starting point.

Differential Revision: https://reviews.llvm.org/D105658

llvm/include/llvm/AsmParser/LLParser.h

@@ -328,10 +328,8 @@ namespace llvm {
     bool parseFnAttributeValuePairs(AttrBuilder &B,
                                     std::vector<unsigned> &FwdRefAttrGrps,
                                     bool inAttrGrp, LocTy &BuiltinLoc);
-    bool parseRequiredTypeAttr(Type *&Result, lltok::Kind AttrName);
+    bool parseRequiredTypeAttr(AttrBuilder &B, lltok::Kind AttrToken,
-    bool parsePreallocated(Type *&Result);
+                               Attribute::AttrKind AttrKind);
-    bool parseInalloca(Type *&Result);
-    bool parseByRef(Type *&Result);
 
     // Module Summary Index Parsing.
     bool skipModuleSummaryEntry();

llvm/include/llvm/IR/Attributes.h

@@ -802,6 +802,14 @@ template <> struct DenseMapInfo<AttributeList> {
 /// value, however, is not. So this can be used as a quick way to test for
 /// equality, presence of attributes, etc.
 class AttrBuilder {
+  // Indices into the TypeAttrs array.
+  static const unsigned ByValTypeIndex = 0;
+  static const unsigned StructRetTypeIndex = 1;
+  static const unsigned ByRefTypeIndex = 2;
+  static const unsigned PreallocatedTypeIndex = 3;
+  static const unsigned InAllocaTypeIndex = 4;
+  static const unsigned NumTypeIndices = 5;
+
   std::bitset<Attribute::EndAttrKinds> Attrs;
   std::map<SmallString<32>, SmallString<32>, std::less<>> TargetDepAttrs;
   MaybeAlign Alignment;
@@ -810,11 +818,9 @@ class AttrBuilder {
   uint64_t DerefOrNullBytes = 0;
   uint64_t AllocSizeArgs = 0;
   uint64_t VScaleRangeArgs = 0;
-  Type *ByValType = nullptr;
+  std::array<Type *, NumTypeIndices> TypeAttrs = {};
-  Type *StructRetType = nullptr;
+
-  Type *ByRefType = nullptr;
+  Optional<unsigned> kindToTypeIndex(Attribute::AttrKind Kind) const;
-  Type *PreallocatedType = nullptr;
-  Type *InAllocaType = nullptr;
 
 public:
   AttrBuilder() = default;
@@ -898,19 +904,19 @@ public:
   uint64_t getDereferenceableOrNullBytes() const { return DerefOrNullBytes; }
 
   /// Retrieve the byval type.
-  Type *getByValType() const { return ByValType; }
+  Type *getByValType() const { return TypeAttrs[ByValTypeIndex]; }
 
   /// Retrieve the sret type.
-  Type *getStructRetType() const { return StructRetType; }
+  Type *getStructRetType() const { return TypeAttrs[StructRetTypeIndex]; }
 
   /// Retrieve the byref type.
-  Type *getByRefType() const { return ByRefType; }
+  Type *getByRefType() const { return TypeAttrs[ByRefTypeIndex]; }
 
   /// Retrieve the preallocated type.
-  Type *getPreallocatedType() const { return PreallocatedType; }
+  Type *getPreallocatedType() const { return TypeAttrs[PreallocatedTypeIndex]; }
 
   /// Retrieve the inalloca type.
-  Type *getInAllocaType() const { return InAllocaType; }
+  Type *getInAllocaType() const { return TypeAttrs[InAllocaTypeIndex]; }
 
   /// Retrieve the allocsize args, if the allocsize attribute exists.  If it
   /// doesn't exist, pair(0, 0) is returned.
@@ -959,6 +965,9 @@ public:
   /// This turns two ints into the form used internally in Attribute.
   AttrBuilder &addVScaleRangeAttr(unsigned MinValue, unsigned MaxValue);
 
+  /// Add a type attribute with the given type.
+  AttrBuilder &addTypeAttr(Attribute::AttrKind Kind, Type *Ty);
+
   /// This turns a byval type into the form used internally in Attribute.
   AttrBuilder &addByValAttr(Type *Ty);
 

llvm/lib/AsmParser/LLParser.cpp

@@ -1411,10 +1411,9 @@ bool LLParser::parseFnAttributeValuePairs(AttrBuilder &B,
     case lltok::kw_willreturn: B.addAttribute(Attribute::WillReturn); break;
     case lltok::kw_writeonly: B.addAttribute(Attribute::WriteOnly); break;
     case lltok::kw_preallocated: {
-      Type *Ty;
+      if (parseRequiredTypeAttr(B, lltok::kw_preallocated,
-      if (parsePreallocated(Ty))
+                                Attribute::Preallocated))
         return true;
-      B.addPreallocatedAttr(Ty);
       break;
     }
 
@@ -1717,34 +1716,27 @@ bool LLParser::parseOptionalParamAttrs(AttrBuilder &B) {
       B.addStackAlignmentAttr(Alignment);
       continue;
     }
-    case lltok::kw_byval: {
+    case lltok::kw_byval:
-      Type *Ty;
+      if (parseRequiredTypeAttr(B, lltok::kw_byval, Attribute::ByVal))
-      if (parseRequiredTypeAttr(Ty, lltok::kw_byval))
         return true;
-      B.addByValAttr(Ty);
       continue;
-    }
+    case lltok::kw_byref:
-    case lltok::kw_sret: {
+      if (parseRequiredTypeAttr(B, lltok::kw_byref, Attribute::ByRef))
-      Type *Ty;
-      if (parseRequiredTypeAttr(Ty, lltok::kw_sret))
         return true;
-      B.addStructRetAttr(Ty);
       continue;
-    }
+    case lltok::kw_inalloca:
-    case lltok::kw_preallocated: {
+      if (parseRequiredTypeAttr(B, lltok::kw_inalloca, Attribute::InAlloca))
-      Type *Ty;
-      if (parsePreallocated(Ty))
         return true;
-      B.addPreallocatedAttr(Ty);
       continue;
-    }
+    case lltok::kw_preallocated:
-    case lltok::kw_inalloca: {
+      if (parseRequiredTypeAttr(B, lltok::kw_preallocated,
-      Type *Ty;
+                                Attribute::Preallocated))
-      if (parseInalloca(Ty))
+        return true;
+      continue;
+    case lltok::kw_sret:
+      if (parseRequiredTypeAttr(B, lltok::kw_sret, Attribute::StructRet))
         return true;
-      B.addInAllocaAttr(Ty);
       continue;
-    }
     case lltok::kw_dereferenceable: {
       uint64_t Bytes;
       if (parseOptionalDerefAttrBytes(lltok::kw_dereferenceable, Bytes))
@@ -1759,13 +1751,6 @@ bool LLParser::parseOptionalParamAttrs(AttrBuilder &B) {
       B.addDereferenceableOrNullAttr(Bytes);
       continue;
     }
-    case lltok::kw_byref: {
-      Type *Ty;
-      if (parseByRef(Ty))
-        return true;
-      B.addByRefAttr(Ty);
-      continue;
-    }
     case lltok::kw_inreg:           B.addAttribute(Attribute::InReg); break;
     case lltok::kw_nest:            B.addAttribute(Attribute::Nest); break;
     case lltok::kw_noundef:
@@ -2708,37 +2693,22 @@ bool LLParser::parseParameterList(SmallVectorImpl<ParamInfo> &ArgList,
 
 /// parseRequiredTypeAttr
 ///   ::= attrname(<ty>)
-bool LLParser::parseRequiredTypeAttr(Type *&Result, lltok::Kind AttrName) {
+bool LLParser::parseRequiredTypeAttr(AttrBuilder &B, lltok::Kind AttrToken,
-  Result = nullptr;
+                                     Attribute::AttrKind AttrKind) {
-  if (!EatIfPresent(AttrName))
+  Type *Ty = nullptr;
+  if (!EatIfPresent(AttrToken))
     return true;
   if (!EatIfPresent(lltok::lparen))
     return error(Lex.getLoc(), "expected '('");
-  if (parseType(Result))
+  if (parseType(Ty))
     return true;
   if (!EatIfPresent(lltok::rparen))
     return error(Lex.getLoc(), "expected ')'");
+
+  B.addTypeAttr(AttrKind, Ty);
   return false;
 }
 
-/// parsePreallocated
-///   ::= preallocated(<ty>)
-bool LLParser::parsePreallocated(Type *&Result) {
-  return parseRequiredTypeAttr(Result, lltok::kw_preallocated);
-}
-
-/// parseInalloca
-///   ::= inalloca(<ty>)
-bool LLParser::parseInalloca(Type *&Result) {
-  return parseRequiredTypeAttr(Result, lltok::kw_inalloca);
-}
-
-/// parseByRef
-///   ::= byref(<type>)
-bool LLParser::parseByRef(Type *&Result) {
-  return parseRequiredTypeAttr(Result, lltok::kw_byref);
-}
-
 /// parseOptionalOperandBundles
 ///    ::= /*empty*/
 ///    ::= '[' OperandBundle [, OperandBundle ]* ']'

llvm/lib/IR/Attributes.cpp

@@ -1597,10 +1597,25 @@ void AttrBuilder::clear() {
   DerefBytes = DerefOrNullBytes = 0;
   AllocSizeArgs = 0;
   VScaleRangeArgs = 0;
-  ByValType = nullptr;
+  TypeAttrs = {};
-  StructRetType = nullptr;
+}
-  ByRefType = nullptr;
+
-  PreallocatedType = nullptr;
+Optional<unsigned>
+AttrBuilder::kindToTypeIndex(Attribute::AttrKind Kind) const {
+  switch (Kind) {
+  case Attribute::ByVal:
+    return ByValTypeIndex;
+  case Attribute::ByRef:
+    return ByRefTypeIndex;
+  case Attribute::InAlloca:
+    return InAllocaTypeIndex;
+  case Attribute::Preallocated:
+    return PreallocatedTypeIndex;
+  case Attribute::StructRet:
+    return StructRetTypeIndex;
+  default:
+    return None;
+  }
 }
 
 AttrBuilder &AttrBuilder::addAttribute(Attribute Attr) {
@@ -1612,18 +1627,12 @@ AttrBuilder &AttrBuilder::addAttribute(Attribute Attr) {
   Attribute::AttrKind Kind = Attr.getKindAsEnum();
   Attrs[Kind] = true;
 
-  if (Kind == Attribute::Alignment)
+  if (Optional<unsigned> TypeIndex = kindToTypeIndex(Kind))
+    TypeAttrs[*TypeIndex] = Attr.getValueAsType();
+  else if (Kind == Attribute::Alignment)
     Alignment = Attr.getAlignment();
   else if (Kind == Attribute::StackAlignment)
     StackAlignment = Attr.getStackAlignment();
-  else if (Kind == Attribute::ByVal)
-    ByValType = Attr.getValueAsType();
-  else if (Kind == Attribute::StructRet)
-    StructRetType = Attr.getValueAsType();
-  else if (Kind == Attribute::ByRef)
-    ByRefType = Attr.getValueAsType();
-  else if (Kind == Attribute::Preallocated)
-    PreallocatedType = Attr.getValueAsType();
   else if (Kind == Attribute::Dereferenceable)
     DerefBytes = Attr.getDereferenceableBytes();
   else if (Kind == Attribute::DereferenceableOrNull)
@@ -1632,8 +1641,6 @@ AttrBuilder &AttrBuilder::addAttribute(Attribute Attr) {
     AllocSizeArgs = Attr.getValueAsInt();
   else if (Kind == Attribute::VScaleRange)
     VScaleRangeArgs = Attr.getValueAsInt();
-  else if (Kind == Attribute::InAlloca)
-    InAllocaType = Attr.getValueAsType();
 
   return *this;
 }
@@ -1647,20 +1654,12 @@ AttrBuilder &AttrBuilder::removeAttribute(Attribute::AttrKind Val) {
   assert((unsigned)Val < Attribute::EndAttrKinds && "Attribute out of range!");
   Attrs[Val] = false;
 
-  if (Val == Attribute::Alignment)
+  if (Optional<unsigned> TypeIndex = kindToTypeIndex(Val))
+    TypeAttrs[*TypeIndex] = nullptr;
+  else if (Val == Attribute::Alignment)
     Alignment.reset();
   else if (Val == Attribute::StackAlignment)
     StackAlignment.reset();
-  else if (Val == Attribute::ByVal)
-    ByValType = nullptr;
-  else if (Val == Attribute::StructRet)
-    StructRetType = nullptr;
-  else if (Val == Attribute::ByRef)
-    ByRefType = nullptr;
-  else if (Val == Attribute::Preallocated)
-    PreallocatedType = nullptr;
-  else if (Val == Attribute::InAlloca)
-    InAllocaType = nullptr;
   else if (Val == Attribute::Dereferenceable)
     DerefBytes = 0;
   else if (Val == Attribute::DereferenceableOrNull)
@@ -1766,34 +1765,32 @@ AttrBuilder &AttrBuilder::addVScaleRangeAttrFromRawRepr(uint64_t RawArgs) {
   return *this;
 }
 
-AttrBuilder &AttrBuilder::addByValAttr(Type *Ty) {
+AttrBuilder &AttrBuilder::addTypeAttr(Attribute::AttrKind Kind, Type *Ty) {
-  Attrs[Attribute::ByVal] = true;
+  Optional<unsigned> TypeIndex = kindToTypeIndex(Kind);
-  ByValType = Ty;
+  assert(TypeIndex && "Not a type attribute");
+  Attrs[Kind] = true;
+  TypeAttrs[*TypeIndex] = Ty;
   return *this;
 }
 
+AttrBuilder &AttrBuilder::addByValAttr(Type *Ty) {
+  return addTypeAttr(Attribute::ByVal, Ty);
+}
+
 AttrBuilder &AttrBuilder::addStructRetAttr(Type *Ty) {
-  Attrs[Attribute::StructRet] = true;
+  return addTypeAttr(Attribute::StructRet, Ty);
-  StructRetType = Ty;
-  return *this;
 }
 
 AttrBuilder &AttrBuilder::addByRefAttr(Type *Ty) {
-  Attrs[Attribute::ByRef] = true;
+  return addTypeAttr(Attribute::ByRef, Ty);
-  ByRefType = Ty;
-  return *this;
 }
 
 AttrBuilder &AttrBuilder::addPreallocatedAttr(Type *Ty) {
-  Attrs[Attribute::Preallocated] = true;
+  return addTypeAttr(Attribute::Preallocated, Ty);
-  PreallocatedType = Ty;
-  return *this;
 }
 
 AttrBuilder &AttrBuilder::addInAllocaAttr(Type *Ty) {
-  Attrs[Attribute::InAlloca] = true;
+  return addTypeAttr(Attribute::InAlloca, Ty);
-  InAllocaType = Ty;
-  return *this;
 }
 
 AttrBuilder &AttrBuilder::merge(const AttrBuilder &B) {
@@ -1813,24 +1810,13 @@ AttrBuilder &AttrBuilder::merge(const AttrBuilder &B) {
   if (!AllocSizeArgs)
     AllocSizeArgs = B.AllocSizeArgs;
 
-  if (!ByValType)
-    ByValType = B.ByValType;
-
-  if (!StructRetType)
-    StructRetType = B.StructRetType;
-
-  if (!ByRefType)
-    ByRefType = B.ByRefType;
-
-  if (!PreallocatedType)
-    PreallocatedType = B.PreallocatedType;
-
-  if (!InAllocaType)
-    InAllocaType = B.InAllocaType;
-
   if (!VScaleRangeArgs)
     VScaleRangeArgs = B.VScaleRangeArgs;
 
+  for (unsigned Index = 0; Index < NumTypeIndices; ++Index)
+    if (!TypeAttrs[Index])
+      TypeAttrs[Index] = B.TypeAttrs[Index];
+
   Attrs |= B.Attrs;
 
   for (const auto &I : B.td_attrs())
@@ -1856,24 +1842,13 @@ AttrBuilder &AttrBuilder::remove(const AttrBuilder &B) {
   if (B.AllocSizeArgs)
     AllocSizeArgs = 0;
 
-  if (B.ByValType)
-    ByValType = nullptr;
-
-  if (B.StructRetType)
-    StructRetType = nullptr;
-
-  if (B.ByRefType)
-    ByRefType = nullptr;
-
-  if (B.PreallocatedType)
-    PreallocatedType = nullptr;
-
-  if (B.InAllocaType)
-    InAllocaType = nullptr;
-
   if (B.VScaleRangeArgs)
     VScaleRangeArgs = 0;
 
+  for (unsigned Index = 0; Index < NumTypeIndices; ++Index)
+    if (B.TypeAttrs[Index])
+      TypeAttrs[Index] = nullptr;
+
   Attrs &= ~B.Attrs;
 
   for (const auto &I : B.td_attrs())
@@ -1932,10 +1907,7 @@ bool AttrBuilder::operator==(const AttrBuilder &B) const {
       return false;
 
   return Alignment == B.Alignment && StackAlignment == B.StackAlignment &&
-         DerefBytes == B.DerefBytes && ByValType == B.ByValType &&
+         DerefBytes == B.DerefBytes && TypeAttrs == B.TypeAttrs &&
-         StructRetType == B.StructRetType && ByRefType == B.ByRefType &&
-         PreallocatedType == B.PreallocatedType &&
-         InAllocaType == B.InAllocaType &&
          VScaleRangeArgs == B.VScaleRangeArgs;
 }