2 files changed, 47 insertions, 22 deletions

diff --git a/llvm/include/llvm/ADT/SmallVector.h b/llvm/include/llvm/ADT/SmallVector.h
index d9a7c6f9538..acb4426b4f4 100644
--- a/llvm/include/llvm/ADT/SmallVector.h
+++ b/llvm/include/llvm/ADT/SmallVector.h
@@ -70,37 +70,42 @@ public:
   }
 };
 
+/// Figure out the offset of the first element.
+template <class T, typename = void> struct SmallVectorAlignmentAndSize {
+  AlignedCharArrayUnion<SmallVectorBase> Base;
+  AlignedCharArrayUnion<T> FirstEl;
+};
+
 /// This is the part of SmallVectorTemplateBase which does not depend on whether
 /// the type T is a POD. The extra dummy template argument is used by ArrayRef
 /// to avoid unnecessarily requiring T to be complete.
 template <typename T, typename = void>
 class SmallVectorTemplateCommon : public SmallVectorBase {
-private:
-  template <typename, unsigned> friend struct SmallVectorStorage;
-
-  // Allocate raw space for N elements of type T.  If T has a ctor or dtor, we
-  // don't want it to be automatically run, so we need to represent the space as
-  // something else.  Use an array of char of sufficient alignment.
-  using U = AlignedCharArrayUnion<T>;
-  U FirstEl;
+  /// Find the address of the first element.  For this pointer math to be valid
+  /// with small-size of 0 for T with lots of alignment, it's important that
+  /// SmallVectorStorage is properly-aligned even for small-size of 0.
+  void *getFirstEl() const {
+    return const_cast<void *>(reinterpret_cast<const void *>(
+        reinterpret_cast<const char *>(this) +
+        offsetof(SmallVectorAlignmentAndSize<T>, FirstEl)));
+  }
   // Space after 'FirstEl' is clobbered, do not add any instance vars after it.
 
 protected:
-  SmallVectorTemplateCommon(size_t Size) : SmallVectorBase(&FirstEl, Size) {}
+  SmallVectorTemplateCommon(size_t Size)
+      : SmallVectorBase(getFirstEl(), Size) {}
 
   void grow_pod(size_t MinCapacity, size_t TSize) {
-    SmallVectorBase::grow_pod(&FirstEl, MinCapacity, TSize);
+    SmallVectorBase::grow_pod(getFirstEl(), MinCapacity, TSize);
   }
 
   /// Return true if this is a smallvector which has not had dynamic
   /// memory allocated for it.
-  bool isSmall() const {
-    return BeginX == static_cast<const void*>(&FirstEl);
-  }
+  bool isSmall() const { return BeginX == getFirstEl(); }
 
   /// Put this vector in a state of being small.
   void resetToSmall() {
-    BeginX = &FirstEl;
+    BeginX = getFirstEl();
     Size = Capacity = 0; // FIXME: Setting Capacity to 0 is suspect.
   }
 
@@ -818,16 +823,17 @@ SmallVectorImpl<T> &SmallVectorImpl<T>::operator=(SmallVectorImpl<T> &&RHS) {
   return *this;
 }
 
-/// Storage for the SmallVector elements which aren't contained in
-/// SmallVectorTemplateCommon. There are 'N-1' elements here. The remaining '1'
-/// element is in the base class. This is specialized for the N=1 and N=0 cases
+/// Storage for the SmallVector elements.  This is specialized for the N=0 case
 /// to avoid allocating unnecessary storage.
 template <typename T, unsigned N>
 struct SmallVectorStorage {
-  typename SmallVectorTemplateCommon<T>::U InlineElts[N - 1];
+  AlignedCharArrayUnion<T> InlineElts[N];
 };
-template <typename T> struct SmallVectorStorage<T, 1> {};
-template <typename T> struct SmallVectorStorage<T, 0> {};
+
+/// We need the storage to be properly aligned even for small-size of 0 so that
+/// the pointer math in \a SmallVectorTemplateCommon::getFirstEl() is
+/// well-defined.
+template <typename T> struct alignas(alignof(T)) SmallVectorStorage<T, 0> {};
 
 /// This is a 'vector' (really, a variable-sized array), optimized
 /// for the case when the array is small.  It contains some number of elements
diff --git a/llvm/lib/Support/SmallVector.cpp b/llvm/lib/Support/SmallVector.cpp
index 7208572d08d..1070c6672ed 100644
--- a/llvm/lib/Support/SmallVector.cpp
+++ b/llvm/lib/Support/SmallVector.cpp
@@ -14,10 +14,29 @@
 #include "llvm/ADT/SmallVector.h"
 using namespace llvm;
 
-// Check that no bytes are wasted.
+// Check that no bytes are wasted and everything is well-aligned.
+namespace {
+struct Struct16B {
+  alignas(16) void *X;
+};
+struct Struct32B {
+  alignas(32) void *X;
+};
+}
+static_assert(sizeof(SmallVector<void *, 0>) ==
+                  sizeof(unsigned) * 2 + sizeof(void *),
+              "wasted space in SmallVector size 0");
+static_assert(alignof(SmallVector<Struct16B, 0>) >= alignof(Struct16B),
+              "wrong alignment for 16-byte aligned T");
+static_assert(alignof(SmallVector<Struct32B, 0>) >= alignof(Struct32B),
+              "wrong alignment for 32-byte aligned T");
+static_assert(sizeof(SmallVector<Struct16B, 0>) >= alignof(Struct16B),
+              "missing padding for 16-byte aligned T");
+static_assert(sizeof(SmallVector<Struct32B, 0>) >= alignof(Struct32B),
+              "missing padding for 32-byte aligned T");
 static_assert(sizeof(SmallVector<void *, 1>) ==
                   sizeof(unsigned) * 2 + sizeof(void *) * 2,
-              "wasted space in SmallVector size 1; missing EBO?");
+              "wasted space in SmallVector size 1");
 
 /// grow_pod - This is an implementation of the grow() method which only works
 /// on POD-like datatypes and is out of line to reduce code duplication.