1 files changed, 66 insertions, 14 deletions

diff --git a/lib/IR/AttributeImpl.h b/lib/IR/AttributeImpl.h
index 659f9568b7c6..267a0dab2f25 100644
--- a/lib/IR/AttributeImpl.h
+++ b/lib/IR/AttributeImpl.h
@@ -17,8 +17,11 @@
 #define LLVM_LIB_IR_ATTRIBUTEIMPL_H
 
 #include "llvm/ADT/FoldingSet.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/IR/Attributes.h"
+#include "llvm/Support/DataTypes.h"
 #include "llvm/Support/TrailingObjects.h"
+#include <climits>
 #include <string>
 
 namespace llvm {
@@ -118,7 +121,8 @@ public:
       : EnumAttributeImpl(IntAttrEntry, Kind), Val(Val) {
     assert((Kind == Attribute::Alignment || Kind == Attribute::StackAlignment ||
             Kind == Attribute::Dereferenceable ||
-            Kind == Attribute::DereferenceableOrNull) &&
+            Kind == Attribute::DereferenceableOrNull ||
+            Kind == Attribute::AllocSize) &&
            "Wrong kind for int attribute!");
   }
 
@@ -148,19 +152,37 @@ class AttributeSetNode final
   friend TrailingObjects;
 
   unsigned NumAttrs; ///< Number of attributes in this node.
+  /// Bitset with a bit for each available attribute Attribute::AttrKind.
+  uint64_t AvailableAttrs;
 
-  AttributeSetNode(ArrayRef<Attribute> Attrs) : NumAttrs(Attrs.size()) {
+  AttributeSetNode(ArrayRef<Attribute> Attrs)
+    : NumAttrs(Attrs.size()), AvailableAttrs(0) {
+    static_assert(Attribute::EndAttrKinds <= sizeof(AvailableAttrs) * CHAR_BIT,
+                  "Too many attributes for AvailableAttrs");
     // There's memory after the node where we can store the entries in.
     std::copy(Attrs.begin(), Attrs.end(), getTrailingObjects<Attribute>());
+
+    for (Attribute I : *this) {
+      if (!I.isStringAttribute()) {
+        AvailableAttrs |= ((uint64_t)1) << I.getKindAsEnum();
+      }
+    }
   }
 
   // AttributesSetNode is uniqued, these should not be publicly available.
   void operator=(const AttributeSetNode &) = delete;
   AttributeSetNode(const AttributeSetNode &) = delete;
 public:
+  void operator delete(void *p) { ::operator delete(p); }
+
   static AttributeSetNode *get(LLVMContext &C, ArrayRef<Attribute> Attrs);
 
-  bool hasAttribute(Attribute::AttrKind Kind) const;
+  /// \brief Return the number of attributes this AttributeSet contains.
+  unsigned getNumAttributes() const { return NumAttrs; }
+
+  bool hasAttribute(Attribute::AttrKind Kind) const {
+    return AvailableAttrs & ((uint64_t)1) << Kind;
+  }
   bool hasAttribute(StringRef Kind) const;
   bool hasAttributes() const { return NumAttrs != 0; }
 
@@ -171,6 +193,7 @@ public:
   unsigned getStackAlignment() const;
   uint64_t getDereferenceableBytes() const;
   uint64_t getDereferenceableOrNullBytes() const;
+  std::pair<unsigned, Optional<unsigned>> getAllocSizeArgs() const;
   std::string getAsString(bool InAttrGrp) const;
 
   typedef const Attribute *iterator;
@@ -200,10 +223,12 @@ class AttributeSetImpl final
 
 private:
   LLVMContext &Context;
-  unsigned NumAttrs; ///< Number of entries in this set.
+  unsigned NumSlots; ///< Number of entries in this set.
+  /// Bitset with a bit for each available attribute Attribute::AttrKind.
+  uint64_t AvailableFunctionAttrs;
 
   // Helper fn for TrailingObjects class.
-  size_t numTrailingObjects(OverloadToken<IndexAttrPair>) { return NumAttrs; }
+  size_t numTrailingObjects(OverloadToken<IndexAttrPair>) { return NumSlots; }
 
   /// \brief Return a pointer to the IndexAttrPair for the specified slot.
   const IndexAttrPair *getNode(unsigned Slot) const {
@@ -215,27 +240,48 @@ private:
   AttributeSetImpl(const AttributeSetImpl &) = delete;
 public:
   AttributeSetImpl(LLVMContext &C,
-                   ArrayRef<std::pair<unsigned, AttributeSetNode *> > Attrs)
-      : Context(C), NumAttrs(Attrs.size()) {
+                   ArrayRef<std::pair<unsigned, AttributeSetNode *> > Slots)
+      : Context(C), NumSlots(Slots.size()), AvailableFunctionAttrs(0) {
+    static_assert(Attribute::EndAttrKinds <=
+                      sizeof(AvailableFunctionAttrs) * CHAR_BIT,
+                  "Too many attributes");
 
 #ifndef NDEBUG
-    if (Attrs.size() >= 2) {
-      for (const std::pair<unsigned, AttributeSetNode *> *i = Attrs.begin() + 1,
-                                                         *e = Attrs.end();
+    if (Slots.size() >= 2) {
+      for (const std::pair<unsigned, AttributeSetNode *> *i = Slots.begin() + 1,
+                                                         *e = Slots.end();
            i != e; ++i) {
         assert((i-1)->first <= i->first && "Attribute set not ordered!");
       }
     }
 #endif
     // There's memory after the node where we can store the entries in.
-    std::copy(Attrs.begin(), Attrs.end(), getTrailingObjects<IndexAttrPair>());
+    std::copy(Slots.begin(), Slots.end(), getTrailingObjects<IndexAttrPair>());
+
+    // Initialize AvailableFunctionAttrs summary bitset.
+    if (NumSlots > 0) {
+      static_assert(AttributeSet::FunctionIndex == ~0u,
+                    "FunctionIndex should be biggest possible index");
+      const std::pair<unsigned, AttributeSetNode *> &Last = Slots.back();
+      if (Last.first == AttributeSet::FunctionIndex) {
+        const AttributeSetNode *Node = Last.second;
+        for (Attribute I : *Node) {
+          if (!I.isStringAttribute())
+            AvailableFunctionAttrs |= ((uint64_t)1) << I.getKindAsEnum();
+        }
+      }
+    }
   }
 
+  void operator delete(void *p) { ::operator delete(p); }
+
   /// \brief Get the context that created this AttributeSetImpl.
   LLVMContext &getContext() { return Context; }
 
-  /// \brief Return the number of attributes this AttributeSet contains.
-  unsigned getNumAttributes() const { return NumAttrs; }
+  /// \brief Return the number of slots used in this attribute list. This is
+  /// the number of arguments that have an attribute set on them (including the
+  /// function itself).
+  unsigned getNumSlots() const { return NumSlots; }
 
   /// \brief Get the index of the given "slot" in the AttrNodes list. This index
   /// is the index of the return, parameter, or function object that the
@@ -258,12 +304,18 @@ public:
     return getNode(Slot)->second;
   }
 
+  /// \brief Return true if the AttributeSetNode for the FunctionIndex has an
+  /// enum attribute of the given kind.
+  bool hasFnAttribute(Attribute::AttrKind Kind) const {
+    return AvailableFunctionAttrs & ((uint64_t)1) << Kind;
+  }
+
   typedef AttributeSetNode::iterator iterator;
   iterator begin(unsigned Slot) const { return getSlotNode(Slot)->begin(); }
   iterator end(unsigned Slot) const { return getSlotNode(Slot)->end(); }
 
   void Profile(FoldingSetNodeID &ID) const {
-    Profile(ID, makeArrayRef(getNode(0), getNumAttributes()));
+    Profile(ID, makeArrayRef(getNode(0), getNumSlots()));
   }
   static void Profile(FoldingSetNodeID &ID,
                       ArrayRef<std::pair<unsigned, AttributeSetNode*> > Nodes) {