2 files changed, 38 insertions, 2 deletions

diff --git a/include/llvm/ADT/SmallVector.h b/include/llvm/ADT/SmallVector.h
index f59a438d3e..dd3a6d0060 100644
--- a/include/llvm/ADT/SmallVector.h
+++ b/include/llvm/ADT/SmallVector.h
@@ -122,11 +122,11 @@ public:
 
 
   reference operator[](unsigned idx) {
-    assert (Begin + idx < End);
+    assert(Begin + idx < End);
     return Begin[idx];
   }
   const_reference operator[](unsigned idx) const {
-    assert (Begin + idx < End);
+    assert(Begin + idx < End);
     return Begin[idx];
   }
 
@@ -399,6 +399,24 @@ public:
                                         RHS.begin(), RHS.end());
   }
 
+  /// capacity - Return the total number of elements in the currently allocated
+  /// buffer.
+  size_t capacity() const { return Capacity - Begin; }
+
+  /// set_size - Set the array size to \arg N, which the current array must have
+  /// enough capacity for.
+  ///
+  /// This does not construct or destroy any elements in the vector.
+  ///
+  /// Clients can use this in conjunction with capacity() to write past the end
+  /// of the buffer when they know that more elements are available, and only
+  /// update the size later. This avoids the cost of value initializing elements
+  /// which will only be overwritten.
+  void set_size(unsigned N) {
+    assert(N <= capacity());
+    End = Begin + N;
+  }
+
 private:
   /// isSmall - Return true if this is a smallvector which has not had dynamic
   /// memory allocated for it.
diff --git a/unittests/ADT/SmallVectorTest.cpp b/unittests/ADT/SmallVectorTest.cpp
index addd904361..8a817966cb 100644
--- a/unittests/ADT/SmallVectorTest.cpp
+++ b/unittests/ADT/SmallVectorTest.cpp
@@ -381,4 +381,22 @@ TEST_F(SmallVectorTest, ConstVectorTest) {
   EXPECT_TRUE(constVector.begin() == constVector.end());
 }
 
+// Direct array access.
+TEST_F(SmallVectorTest, DirectVectorTest) {
+  EXPECT_EQ(0u, theVector.size());
+  EXPECT_EQ(4u, theVector.capacity());
+  EXPECT_EQ(0, Constructable::getNumConstructorCalls());
+  theVector.end()[0] = 1;
+  theVector.end()[1] = 2;
+  theVector.end()[2] = 3;
+  theVector.end()[3] = 4;
+  theVector.set_size(4);
+  EXPECT_EQ(4u, theVector.size());
+  EXPECT_EQ(4, Constructable::getNumConstructorCalls());
+  EXPECT_EQ(1, theVector[0].getValue());
+  EXPECT_EQ(2, theVector[1].getValue());
+  EXPECT_EQ(3, theVector[2].getValue());
+  EXPECT_EQ(4, theVector[3].getValue());
+}
+
 }