diff --git a/include/llvm/ADT/SmallVector.h b/include/llvm/ADT/SmallVector.h index f59a438d3eb..dd3a6d00603 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 addd904361c..8a817966cb8 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()); +} + }