2 files changed, 83 insertions, 12 deletions

diff --git a/clang/include/clang/AST/ASTVector.h b/clang/include/clang/AST/ASTVector.h
index 3b856a79193..6ec054582e2 100644
--- a/clang/include/clang/AST/ASTVector.h
+++ b/clang/include/clang/AST/ASTVector.h
@@ -236,14 +236,14 @@ public:
 
   iterator insert(const ASTContext &C, iterator I, size_type NumToInsert,
                   const T &Elt) {
-    if (I == this->end()) {  // Important special case for empty vector.
-      append(C, NumToInsert, Elt);
-      return this->end()-1;
-    }
-
     // Convert iterator to elt# to avoid invalidating iterator when we reserve()
     size_t InsertElt = I - this->begin();
 
+    if (I == this->end()) { // Important special case for empty vector.
+      append(C, NumToInsert, Elt);
+      return this->begin() + InsertElt;
+    }
+
     // Ensure there is enough space.
     reserve(C, static_cast<unsigned>(this->size() + NumToInsert));
 
@@ -284,14 +284,15 @@ public:
 
   template<typename ItTy>
   iterator insert(const ASTContext &C, iterator I, ItTy From, ItTy To) {
-    if (I == this->end()) {  // Important special case for empty vector.
+    // Convert iterator to elt# to avoid invalidating iterator when we reserve()
+    size_t InsertElt = I - this->begin();
+
+    if (I == this->end()) { // Important special case for empty vector.
       append(C, From, To);
-      return this->end()-1;
+      return this->begin() + InsertElt;
     }
 
     size_t NumToInsert = std::distance(From, To);
-    // Convert iterator to elt# to avoid invalidating iterator when we reserve()
-    size_t InsertElt = I - this->begin();
 
     // Ensure there is enough space.
     reserve(C, static_cast<unsigned>(this->size() + NumToInsert));
diff --git a/clang/unittests/AST/ASTVectorTest.cpp b/clang/unittests/AST/ASTVectorTest.cpp
index ce6d0a07ce8..3a8e8e89b6e 100644
--- a/clang/unittests/AST/ASTVectorTest.cpp
+++ b/clang/unittests/AST/ASTVectorTest.cpp
@@ -14,11 +14,81 @@
 #include "llvm/Support/Compiler.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTVector.h"
+#include "clang/Basic/Builtins.h"
+
+#include "gtest/gtest.h"
+
+#include <vector>
 
 using namespace clang;
 
-LLVM_ATTRIBUTE_UNUSED void CompileTest() {
-  ASTContext *C = nullptr;
+namespace clang {
+namespace ast {
+
+namespace {
+class ASTVectorTest : public ::testing::Test {
+protected:
+  ASTVectorTest()
+      : FileMgr(FileMgrOpts), DiagID(new DiagnosticIDs()),
+        Diags(DiagID, new DiagnosticOptions, new IgnoringDiagConsumer()),
+        SourceMgr(Diags, FileMgr), Idents(LangOpts, nullptr),
+        Ctxt(LangOpts, SourceMgr, Idents, Sels, Builtins) {}
+
+  FileSystemOptions FileMgrOpts;
+  FileManager FileMgr;
+  IntrusiveRefCntPtr<DiagnosticIDs> DiagID;
+  DiagnosticsEngine Diags;
+  SourceManager SourceMgr;
+  LangOptions LangOpts;
+  IdentifierTable Idents;
+  SelectorTable Sels;
+  Builtin::Context Builtins;
+  ASTContext Ctxt;
+};
+} // unnamed namespace
+
+TEST_F(ASTVectorTest, Compile) {
   ASTVector<int> V;
-  V.insert(*C, V.begin(), 0);
+  V.insert(Ctxt, V.begin(), 0);
 }
+
+TEST_F(ASTVectorTest, InsertFill) {
+  ASTVector<double> V;
+
+  // Ensure returned iterator points to first of inserted elements
+  auto I = V.insert(Ctxt, V.begin(), 5, 1.0);
+  ASSERT_EQ(V.begin(), I);
+
+  // Check non-empty case as well
+  I = V.insert(Ctxt, V.begin() + 1, 5, 1.0);
+  ASSERT_EQ(V.begin() + 1, I);
+
+  // And insert-at-end
+  I = V.insert(Ctxt, V.end(), 5, 1.0);
+  ASSERT_EQ(V.end() - 5, I);
+}
+
+TEST_F(ASTVectorTest, InsertEmpty) {
+  ASTVector<double> V;
+
+  // Ensure no pointer overflow when inserting empty range
+  std::vector<int> IntVec{0, 1, 2, 3};
+  auto I = V.insert(Ctxt, V.begin(), IntVec.begin(), IntVec.begin());
+  ASSERT_EQ(V.begin(), I);
+  ASSERT_TRUE(V.empty());
+
+  // Non-empty range
+  I = V.insert(Ctxt, V.begin(), IntVec.begin(), IntVec.end());
+  ASSERT_EQ(V.begin(), I);
+
+  // Non-Empty Vector, empty range
+  I = V.insert(Ctxt, V.end(), IntVec.begin(), IntVec.begin());
+  ASSERT_EQ(V.begin() + IntVec.size(), I);
+
+  // Non-Empty Vector, non-empty range
+  I = V.insert(Ctxt, V.end(), IntVec.begin(), IntVec.end());
+  ASSERT_EQ(V.begin() + IntVec.size(), I);
+}
+
+} // end namespace ast
+} // end namespace clang