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

[dneilson]

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.

[reames]

A few minor changes needed, but generally looks close to ready.

lib/IR/Attributes.cpp

564 ↗

(On Diff #129460)

Incomplete handling. I think the assert is simply stale and missing a test case.

1101 ↗

(On Diff #129460)

I believe this check is implied by the hasAttribute check above it.

unittests/IR/AttributesTest.cpp

81 ↗

(On Diff #129460)

Add a test for remove AttrBuilder case.

Also add a test case for a non-param alignment or at least one of them. Say, stackalign for the AttrBuilder case.

[rnk]

lib/IR/Attributes.cpp

1106 ↗

(On Diff #129460)

I think we can just assert(Index < AttrSets.size()) because otherwise hasAttribute would return false.

1109–1111 ↗

(On Diff #129460)

Can you simplify this to AttrSets[Index] = AttrSets[Index].removeAttribute(Kind);? It seems you fixed up AttributeSet::removeAttribute to handle this case.

1124 ↗

(On Diff #129460)

ditto, can strengthen to assert.

1129 ↗

(On Diff #129460)

ditto, can use AttributeSet::removeAttribute.

[dneilson]

Address comments by Philip & Reid.

[rnk]

lgtm, thanks!

llvm/trunk/lib/IR/Attributes.cpp

for (Attribute I : *this)
     B.addAttribute(I);
 
  return get(C, B);
 }
 
 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);
 }
 
 unsigned AttributeSet::getNumAttributes() const {
   return SetNode ? SetNode->getNumAttributes() : 0;
 }
AttributeList AttributeList::addParamAttribute(LLVMContext &C,
   }
 
   return getImpl(C, AttrSets);
 }
 
 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
 AttributeList::removeAttributes(LLVMContext &C, unsigned Index,
                                 const AttrBuilder &AttrsToRemove) const {
   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);
 }
 
 AttributeList AttributeList::removeAttributes(LLVMContext &C,
                                               unsigned WithoutIndex) const {
   if (!pImpl)
     return AttributeList();

llvm/trunk/unittests/IR/AttributesTest.cpp

TEST(Attributes, AddAttributes) {
   EXPECT_TRUE(AL.hasFnAttribute(Attribute::NoReturn));
   B.clear();
   B.addAttribute(Attribute::SExt);
   AL = AL.addAttributes(C, AttributeList::ReturnIndex, B);
   EXPECT_TRUE(AL.hasAttribute(AttributeList::ReturnIndex, Attribute::SExt));
   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;
   AL = AL.addAttribute(C, AttributeList::FirstArgIndex,
                        Attribute::getWithAlignment(C, 8));
   AL = AL.addAttribute(C, AttributeList::FirstArgIndex + 1,
                        Attribute::getWithAlignment(C, 32));
   EXPECT_EQ(8U, AL.getParamAlignment(0));