[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

1 files changed, 47 insertions, 75 deletions

diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
index 9c06e19..86b296e 100644
--- a/llvm/lib/IR/Attributes.cpp
+++ b/llvm/lib/IR/Attributes.cpp
@@ -1597,10 +1597,25 @@ void AttrBuilder::clear() {
   DerefBytes = DerefOrNullBytes = 0;
   AllocSizeArgs = 0;
   VScaleRangeArgs = 0;
-  ByValType = nullptr;
-  StructRetType = nullptr;
-  ByRefType = nullptr;
-  PreallocatedType = nullptr;
+  TypeAttrs = {};
+}
+
+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) {
-  Attrs[Attribute::ByVal] = true;
-  ByValType = Ty;
+AttrBuilder &AttrBuilder::addTypeAttr(Attribute::AttrKind Kind, Type *Ty) {
+  Optional<unsigned> TypeIndex = kindToTypeIndex(Kind);
+  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;
-  StructRetType = Ty;
-  return *this;
+  return addTypeAttr(Attribute::StructRet, Ty);
 }
 
 AttrBuilder &AttrBuilder::addByRefAttr(Type *Ty) {
-  Attrs[Attribute::ByRef] = true;
-  ByRefType = Ty;
-  return *this;
+  return addTypeAttr(Attribute::ByRef, Ty);
 }
 
 AttrBuilder &AttrBuilder::addPreallocatedAttr(Type *Ty) {
-  Attrs[Attribute::Preallocated] = true;
-  PreallocatedType = Ty;
-  return *this;
+  return addTypeAttr(Attribute::Preallocated, Ty);
 }
 
 AttrBuilder &AttrBuilder::addInAllocaAttr(Type *Ty) {
-  Attrs[Attribute::InAlloca] = true;
-  InAllocaType = Ty;
-  return *this;
+  return addTypeAttr(Attribute::InAlloca, Ty);
 }
 
 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 &&
-         StructRetType == B.StructRetType && ByRefType == B.ByRefType &&
-         PreallocatedType == B.PreallocatedType &&
-         InAllocaType == B.InAllocaType &&
+         DerefBytes == B.DerefBytes && TypeAttrs == B.TypeAttrs &&
          VScaleRangeArgs == B.VScaleRangeArgs;
 }