[Attributes] Fix crash when attempting to remove alignment from an attribute list/set

Daniel Neilson · Daniel Neilson · commit 88dddb89483e · 2018-01-17T19:15:21.000Z
Summary: Discovered while working on a patch to move alignment in @llvm.memcpy/move/set from an arg into parameter attributes. The current implementations of AttributeSet::removeAttribute() and AttributeList::removeAttribute crash when attempting to remove the alignment attribute. Currently, these implementations add the to-be-removed attributes to an AttrBuilder and then remove the builder from the list/set. Alignment is special in that it must be added to a builder with an integer value for the alignment; attempts to add alignment to a builder without a value is an error. This change fixes the removeAttribute implementations for AttributeSet and AttributeList to make them able to remove the alignment, and other similar, attributes. Reviewers: rnk, chandlerc, pete, javed.absar, reames Reviewed By: rnk Subscribers: llvm-commits Differential Revision: https://reviews.llvm.org/D41951 llvm-svn: 322735
diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
@@ -543,26 +543,21 @@ AttributeSet AttributeSet::addAttributes(LLVMContext &C,
 AttributeSet AttributeSet::removeAttribute(LLVMContext &C,
                                              Attribute::AttrKind Kind) const {
   if (!hasAttribute(Kind)) return *this;
-  AttrBuilder B;
-  B.addAttribute(Kind);
-  return removeAttributes(C, B);
+  AttrBuilder B(*this);
+  B.removeAttribute(Kind);
+  return get(C, B);
 }
 
 AttributeSet AttributeSet::removeAttribute(LLVMContext &C,
                                              StringRef Kind) const {
   if (!hasAttribute(Kind)) return *this;
-  AttrBuilder B;
-  B.addAttribute(Kind);
-  return removeAttributes(C, B);
+  AttrBuilder B(*this);
+  B.removeAttribute(Kind);
+  return get(C, B);
 }
 
 AttributeSet AttributeSet::removeAttributes(LLVMContext &C,
                                               const AttrBuilder &Attrs) const {
-
-  // FIXME it is not obvious how this should work for alignment.
-  // For now, say we can't pass in alignment, which no current use does.
-  assert(!Attrs.hasAlignmentAttr() && "Attempt to change alignment!");
-
   AttrBuilder B(*this);
   B.remove(Attrs);
   return get(C, B);
@@ -1098,17 +1093,27 @@ AttributeList AttributeList::addParamAttribute(LLVMContext &C,
 AttributeList AttributeList::removeAttribute(LLVMContext &C, unsigned Index,
                                              Attribute::AttrKind Kind) const {
   if (!hasAttribute(Index, Kind)) return *this;
-  AttrBuilder B;
-  B.addAttribute(Kind);
-  return removeAttributes(C, Index, B);
+
+  Index = attrIdxToArrayIdx(Index);
+  SmallVector<AttributeSet, 4> AttrSets(this->begin(), this->end());
+  assert(Index < AttrSets.size());
+
+  AttrSets[Index] = AttrSets[Index].removeAttribute(C, Kind);
+
+  return getImpl(C, AttrSets);
 }
 
 AttributeList AttributeList::removeAttribute(LLVMContext &C, unsigned Index,
                                              StringRef Kind) const {
   if (!hasAttribute(Index, Kind)) return *this;
-  AttrBuilder B;
-  B.addAttribute(Kind);
-  return removeAttributes(C, Index, B);
+
+  Index = attrIdxToArrayIdx(Index);
+  SmallVector<AttributeSet, 4> AttrSets(this->begin(), this->end());
+  assert(Index < AttrSets.size());
+
+  AttrSets[Index] = AttrSets[Index].removeAttribute(C, Kind);
+
+  return getImpl(C, AttrSets);
 }
 
 AttributeList
@@ -1117,18 +1122,12 @@ AttributeList::removeAttributes(LLVMContext &C, unsigned Index,
   if (!pImpl)
     return AttributeList();
 
-  // FIXME it is not obvious how this should work for alignment.
-  // For now, say we can't pass in alignment, which no current use does.
-  assert(!AttrsToRemove.hasAlignmentAttr() && "Attempt to change alignment!");
-
   Index = attrIdxToArrayIdx(Index);
   SmallVector<AttributeSet, 4> AttrSets(this->begin(), this->end());
   if (Index >= AttrSets.size())
     AttrSets.resize(Index + 1);
 
-  AttrBuilder B(AttrSets[Index]);
-  B.remove(AttrsToRemove);
-  AttrSets[Index] = AttributeSet::get(C, B);
+  AttrSets[Index] = AttrSets[Index].removeAttributes(C, AttrsToRemove);
 
   return getImpl(C, AttrSets);
 }
diff --git a/llvm/unittests/IR/AttributesTest.cpp b/llvm/unittests/IR/AttributesTest.cpp
@@ -63,6 +63,76 @@ TEST(Attributes, AddAttributes) {
   EXPECT_TRUE(AL.hasFnAttribute(Attribute::NoReturn));
 }
 
+TEST(Attributes, RemoveAlign) {
+  LLVMContext C;
+
+  Attribute AlignAttr = Attribute::getWithAlignment(C, 8);
+  Attribute StackAlignAttr = Attribute::getWithStackAlignment(C, 32);
+  AttrBuilder B_align_readonly;
+  B_align_readonly.addAttribute(AlignAttr);
+  B_align_readonly.addAttribute(Attribute::ReadOnly);
+  AttrBuilder B_align;
+  B_align.addAttribute(AlignAttr);
+  AttrBuilder B_stackalign_optnone;
+  B_stackalign_optnone.addAttribute(StackAlignAttr);
+  B_stackalign_optnone.addAttribute(Attribute::OptimizeNone);
+  AttrBuilder B_stackalign;
+  B_stackalign.addAttribute(StackAlignAttr);
+
+  AttributeSet AS = AttributeSet::get(C, B_align_readonly);
+  EXPECT_TRUE(AS.getAlignment() == 8);
+  EXPECT_TRUE(AS.hasAttribute(Attribute::ReadOnly));
+  AS = AS.removeAttribute(C, Attribute::Alignment);
+  EXPECT_FALSE(AS.hasAttribute(Attribute::Alignment));
+  EXPECT_TRUE(AS.hasAttribute(Attribute::ReadOnly));
+  AS = AttributeSet::get(C, B_align_readonly);
+  AS = AS.removeAttributes(C, B_align);
+  EXPECT_TRUE(AS.getAlignment() == 0);
+  EXPECT_TRUE(AS.hasAttribute(Attribute::ReadOnly));
+
+  AttributeList AL;
+  AL = AL.addParamAttributes(C, 0, B_align_readonly);
+  AL = AL.addAttributes(C, 0, B_stackalign_optnone);
+  EXPECT_TRUE(AL.hasAttributes(0));
+  EXPECT_TRUE(AL.hasAttribute(0, Attribute::StackAlignment));
+  EXPECT_TRUE(AL.hasAttribute(0, Attribute::OptimizeNone));
+  EXPECT_TRUE(AL.getStackAlignment(0) == 32);
+  EXPECT_TRUE(AL.hasParamAttrs(0));
+  EXPECT_TRUE(AL.hasParamAttr(0, Attribute::Alignment));
+  EXPECT_TRUE(AL.hasParamAttr(0, Attribute::ReadOnly));
+  EXPECT_TRUE(AL.getParamAlignment(0) == 8);
+
+  AL = AL.removeParamAttribute(C, 0, Attribute::Alignment);
+  EXPECT_FALSE(AL.hasParamAttr(0, Attribute::Alignment));
+  EXPECT_TRUE(AL.hasParamAttr(0, Attribute::ReadOnly));
+  EXPECT_TRUE(AL.hasAttribute(0, Attribute::StackAlignment));
+  EXPECT_TRUE(AL.hasAttribute(0, Attribute::OptimizeNone));
+  EXPECT_TRUE(AL.getStackAlignment(0) == 32);
+
+  AL = AL.removeAttribute(C, 0, Attribute::StackAlignment);
+  EXPECT_FALSE(AL.hasParamAttr(0, Attribute::Alignment));
+  EXPECT_TRUE(AL.hasParamAttr(0, Attribute::ReadOnly));
+  EXPECT_FALSE(AL.hasAttribute(0, Attribute::StackAlignment));
+  EXPECT_TRUE(AL.hasAttribute(0, Attribute::OptimizeNone));
+
+  AttributeList AL2;
+  AL2 = AL2.addParamAttributes(C, 0, B_align_readonly);
+  AL2 = AL2.addAttributes(C, 0, B_stackalign_optnone);
+
+  AL2 = AL2.removeParamAttributes(C, 0, B_align);
+  EXPECT_FALSE(AL2.hasParamAttr(0, Attribute::Alignment));
+  EXPECT_TRUE(AL2.hasParamAttr(0, Attribute::ReadOnly));
+  EXPECT_TRUE(AL2.hasAttribute(0, Attribute::StackAlignment));
+  EXPECT_TRUE(AL2.hasAttribute(0, Attribute::OptimizeNone));
+  EXPECT_TRUE(AL2.getStackAlignment(0) == 32);
+
+  AL2 = AL2.removeAttributes(C, 0, B_stackalign);
+  EXPECT_FALSE(AL2.hasParamAttr(0, Attribute::Alignment));
+  EXPECT_TRUE(AL2.hasParamAttr(0, Attribute::ReadOnly));
+  EXPECT_FALSE(AL2.hasAttribute(0, Attribute::StackAlignment));
+  EXPECT_TRUE(AL2.hasAttribute(0, Attribute::OptimizeNone));
+}
+
 TEST(Attributes, AddMatchingAlignAttr) {
   LLVMContext C;
   AttributeList AL;
