[Attributes] Make type attribute handling more generic (NFCI)

Followup to D105658 to make AttrBuilder automatically work with
new type attributes. TableGen is tweaked to emit First/LastTypeAttr
markers, based on which we can handle type attributes
programmatically.

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

llvm/include/llvm/IR/Attributes.h

@@ -71,14 +71,19 @@ public:
   enum AttrKind {
     // IR-Level Attributes
     None,                  ///< No attributes have been set
-    #define GET_ATTR_NAMES
+    #define GET_ATTR_ENUM
-    #define ATTRIBUTE_ENUM(ENUM_NAME, OTHER) ENUM_NAME,
     #include "llvm/IR/Attributes.inc"
     EndAttrKinds,          ///< Sentinal value useful for loops
     EmptyKey,              ///< Use as Empty key for DenseMap of AttrKind
     TombstoneKey,          ///< Use as Tombstone key for DenseMap of AttrKind
   };
 
+  static const unsigned NumTypeAttrKinds = LastTypeAttr - FirstTypeAttr + 1;
+
+  static bool isTypeAttrKind(AttrKind Kind) {
+    return Kind >= FirstTypeAttr && Kind <= LastTypeAttr;
+  }
+
 private:
   AttributeImpl *pImpl = nullptr;
 
@@ -802,16 +807,6 @@ 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.
-  enum : unsigned {
-    ByValTypeIndex = 0,
-    StructRetTypeIndex = 1,
-    ByRefTypeIndex = 2,
-    PreallocatedTypeIndex = 3,
-    InAllocaTypeIndex = 4,
-    NumTypeIndices = 5,
-  };
-
   std::bitset<Attribute::EndAttrKinds> Attrs;
   std::map<SmallString<32>, SmallString<32>, std::less<>> TargetDepAttrs;
   MaybeAlign Alignment;
@@ -820,7 +815,7 @@ class AttrBuilder {
   uint64_t DerefOrNullBytes = 0;
   uint64_t AllocSizeArgs = 0;
   uint64_t VScaleRangeArgs = 0;
-  std::array<Type *, NumTypeIndices> TypeAttrs = {};
+  std::array<Type *, Attribute::NumTypeAttrKinds> TypeAttrs = {};
 
   Optional<unsigned> kindToTypeIndex(Attribute::AttrKind Kind) const;
 
@@ -905,20 +900,25 @@ public:
   /// dereferenceable_or_null attribute exists (zero is returned otherwise).
   uint64_t getDereferenceableOrNullBytes() const { return DerefOrNullBytes; }
 
+  /// Retrieve type for the given type attribute.
+  Type *getTypeAttr(Attribute::AttrKind Kind) const;
+
   /// Retrieve the byval type.
-  Type *getByValType() const { return TypeAttrs[ByValTypeIndex]; }
+  Type *getByValType() const { return getTypeAttr(Attribute::ByVal); }
 
   /// Retrieve the sret type.
-  Type *getStructRetType() const { return TypeAttrs[StructRetTypeIndex]; }
+  Type *getStructRetType() const { return getTypeAttr(Attribute::StructRet); }
 
   /// Retrieve the byref type.
-  Type *getByRefType() const { return TypeAttrs[ByRefTypeIndex]; }
+  Type *getByRefType() const { return getTypeAttr(Attribute::ByRef); }
 
   /// Retrieve the preallocated type.
-  Type *getPreallocatedType() const { return TypeAttrs[PreallocatedTypeIndex]; }
+  Type *getPreallocatedType() const {
+    return getTypeAttr(Attribute::Preallocated);
+  }
 
   /// Retrieve the inalloca type.
-  Type *getInAllocaType() const { return TypeAttrs[InAllocaTypeIndex]; }
+  Type *getInAllocaType() const { return getTypeAttr(Attribute::InAlloca); }
 
   /// Retrieve the allocsize args, if the allocsize attribute exists.  If it
   /// doesn't exist, pair(0, 0) is returned.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -1642,17 +1642,10 @@ Error BitcodeReader::parseAttributeGroupBlock() {
           Attribute::AttrKind Kind;
           if (Error Err = parseAttrKind(Record[++i], &Kind))
             return Err;
-          if (Kind == Attribute::ByVal) {
+          if (!Attribute::isTypeAttrKind(Kind))
-            B.addByValAttr(HasType ? getTypeByID(Record[++i]) : nullptr);
+            return error("Not a type attribute");
-          } else if (Kind == Attribute::StructRet) {
+
-            B.addStructRetAttr(HasType ? getTypeByID(Record[++i]) : nullptr);
+          B.addTypeAttr(Kind, HasType ? getTypeByID(Record[++i]) : nullptr);
-          } else if (Kind == Attribute::ByRef) {
-            B.addByRefAttr(getTypeByID(Record[++i]));
-          } else if (Kind == Attribute::Preallocated) {
-            B.addPreallocatedAttr(getTypeByID(Record[++i]));
-          } else if (Kind == Attribute::InAlloca) {
-            B.addInAllocaAttr(HasType ? getTypeByID(Record[++i]) : nullptr);
-          }
         }
       }
 

llvm/lib/IR/Attributes.cpp

@@ -138,6 +138,7 @@ Attribute Attribute::get(LLVMContext &Context, StringRef Kind, StringRef Val) {
 
 Attribute Attribute::get(LLVMContext &Context, Attribute::AttrKind Kind,
                          Type *Ty) {
+  assert(Attribute::isTypeAttrKind(Kind) && "Not a type attribute");
   LLVMContextImpl *pImpl = Context.pImpl;
   FoldingSetNodeID ID;
   ID.AddInteger(Kind);
@@ -800,23 +801,13 @@ AttributeSetNode *AttributeSetNode::get(LLVMContext &C, const AttrBuilder &B) {
     if (!B.contains(Kind))
       continue;
 
+    if (Attribute::isTypeAttrKind(Kind)) {
+      Attrs.push_back(Attribute::get(C, Kind, B.getTypeAttr(Kind)));
+      continue;
+    }
+
     Attribute Attr;
     switch (Kind) {
-    case Attribute::ByVal:
-      Attr = Attribute::getWithByValType(C, B.getByValType());
-      break;
-    case Attribute::StructRet:
-      Attr = Attribute::getWithStructRetType(C, B.getStructRetType());
-      break;
-    case Attribute::ByRef:
-      Attr = Attribute::getWithByRefType(C, B.getByRefType());
-      break;
-    case Attribute::Preallocated:
-      Attr = Attribute::getWithPreallocatedType(C, B.getPreallocatedType());
-      break;
-    case Attribute::InAlloca:
-      Attr = Attribute::getWithInAllocaType(C, B.getInAllocaType());
-      break;
     case Attribute::Alignment:
       assert(B.getAlignment() && "Alignment must be set");
       Attr = Attribute::getWithAlignment(C, *B.getAlignment());
@@ -1602,20 +1593,9 @@ void AttrBuilder::clear() {
 
 Optional<unsigned>
 AttrBuilder::kindToTypeIndex(Attribute::AttrKind Kind) const {
-  switch (Kind) {
+  if (Attribute::isTypeAttrKind(Kind))
-  case Attribute::ByVal:
+    return Kind - Attribute::FirstTypeAttr;
-    return ByValTypeIndex;
+  return None;
-  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) {
@@ -1765,6 +1745,12 @@ AttrBuilder &AttrBuilder::addVScaleRangeAttrFromRawRepr(uint64_t RawArgs) {
   return *this;
 }
 
+Type *AttrBuilder::getTypeAttr(Attribute::AttrKind Kind) const {
+  Optional<unsigned> TypeIndex = kindToTypeIndex(Kind);
+  assert(TypeIndex && "Not a type attribute");
+  return TypeAttrs[*TypeIndex];
+}
+
 AttrBuilder &AttrBuilder::addTypeAttr(Attribute::AttrKind Kind, Type *Ty) {
   Optional<unsigned> TypeIndex = kindToTypeIndex(Kind);
   assert(TypeIndex && "Not a type attribute");
@@ -1813,7 +1799,7 @@ AttrBuilder &AttrBuilder::merge(const AttrBuilder &B) {
   if (!VScaleRangeArgs)
     VScaleRangeArgs = B.VScaleRangeArgs;
 
-  for (unsigned Index = 0; Index < NumTypeIndices; ++Index)
+  for (unsigned Index = 0; Index < Attribute::NumTypeAttrKinds; ++Index)
     if (!TypeAttrs[Index])
       TypeAttrs[Index] = B.TypeAttrs[Index];
 
@@ -1845,7 +1831,7 @@ AttrBuilder &AttrBuilder::remove(const AttrBuilder &B) {
   if (B.VScaleRangeArgs)
     VScaleRangeArgs = 0;
 
-  for (unsigned Index = 0; Index < NumTypeIndices; ++Index)
+  for (unsigned Index = 0; Index < Attribute::NumTypeAttrKinds; ++Index)
     if (B.TypeAttrs[Index])
       TypeAttrs[Index] = nullptr;
 

llvm/utils/TableGen/Attributes.cpp

@@ -58,7 +58,20 @@ void Attributes::emitTargetIndependentNames(raw_ostream &OS) {
   Emit({"StrBoolAttr"}, "ATTRIBUTE_STRBOOL");
 
   OS << "#undef ATTRIBUTE_ALL\n";
-  OS << "#endif\n";
+  OS << "#endif\n\n";
+
+  OS << "#ifdef GET_ATTR_ENUM\n";
+  OS << "#undef GET_ATTR_ENUM\n";
+  unsigned Value = 1; // Leave zero for AttrKind::None.
+  for (StringRef KindName : {"EnumAttr", "TypeAttr", "IntAttr"}) {
+    OS << "First" << KindName << " = " << Value << ",\n";
+    for (auto A : Records.getAllDerivedDefinitions(KindName)) {
+      OS << A->getName() << " = " << Value << ",\n";
+      Value++;
+    }
+    OS << "Last" << KindName << " = " << (Value - 1) << ",\n";
+  }
+  OS << "#endif\n\n";
 }
 
 void Attributes::emitFnAttrCompatCheck(raw_ostream &OS, bool IsStringAttr) {