diff --git a/llvm/include/llvm/ADT/SmallString.h b/llvm/include/llvm/ADT/SmallString.h index cd6f2173d04f4..2fa47c7f4e237 100644 --- a/llvm/include/llvm/ADT/SmallString.h +++ b/llvm/include/llvm/ADT/SmallString.h @@ -30,6 +30,12 @@ class SmallString : public SmallVector { /// Initialize from a StringRef. SmallString(StringRef S) : SmallVector(S.begin(), S.end()) {} + /// Initialize by concatenating a list of StringRefs. + SmallString(std::initializer_list Refs) + : SmallVector() { + this->append(Refs); + } + /// Initialize with a range. template SmallString(ItTy S, ItTy E) : SmallVector(S, E) {} @@ -65,6 +71,12 @@ class SmallString : public SmallVector { SmallVectorImpl::append(RHS.begin(), RHS.end()); } + /// Assign from a list of StringRefs. + void assign(std::initializer_list Refs) { + this->clear(); + append(Refs); + } + /// @} /// @name String Concatenation /// @{ @@ -89,6 +101,20 @@ class SmallString : public SmallVector { SmallVectorImpl::append(RHS.begin(), RHS.end()); } + /// Append from a list of StringRefs. + void append(std::initializer_list Refs) { + size_t SizeNeeded = this->size(); + for (const StringRef &Ref : Refs) + SizeNeeded += Ref.size(); + this->reserve(SizeNeeded); + auto CurEnd = this->end(); + for (const StringRef &Ref : Refs) { + this->uninitialized_copy(Ref.begin(), Ref.end(), CurEnd); + CurEnd += Ref.size(); + } + this->set_size(SizeNeeded); + } + /// @} /// @name String Comparison /// @{ diff --git a/llvm/unittests/ADT/SmallStringTest.cpp b/llvm/unittests/ADT/SmallStringTest.cpp index e78da9fe52898..3401a8266a6cc 100644 --- a/llvm/unittests/ADT/SmallStringTest.cpp +++ b/llvm/unittests/ADT/SmallStringTest.cpp @@ -71,6 +71,12 @@ TEST_F(SmallStringTest, AssignSmallVector) { EXPECT_STREQ("abc", theString.c_str()); } +TEST_F(SmallStringTest, AssignStringRefs) { + theString.assign({"abc", "def", "ghi"}); + EXPECT_EQ(9u, theString.size()); + EXPECT_STREQ("abcdefghi", theString.c_str()); +} + TEST_F(SmallStringTest, AppendIterPair) { StringRef abc = "abc"; theString.append(abc.begin(), abc.end()); @@ -96,6 +102,19 @@ TEST_F(SmallStringTest, AppendSmallVector) { EXPECT_STREQ("abcabc", theString.c_str()); } +TEST_F(SmallStringTest, AppendStringRefs) { + theString.append({"abc", "def", "ghi"}); + EXPECT_EQ(9u, theString.size()); + EXPECT_STREQ("abcdefghi", theString.c_str()); + StringRef Jkl = "jkl"; + std::string Mno = "mno"; + SmallString<4> Pqr("pqr"); + const char *Stu = "stu"; + theString.append({Jkl, Mno, Pqr, Stu}); + EXPECT_EQ(21u, theString.size()); + EXPECT_STREQ("abcdefghijklmnopqrstu", theString.c_str()); +} + TEST_F(SmallStringTest, StringRefConversion) { StringRef abc = "abc"; theString.assign(abc.begin(), abc.end());