1 files changed, 62 insertions, 9 deletions

diff --git a/clang/unittests/Lex/LexerTest.cpp b/clang/unittests/Lex/LexerTest.cpp
index e2507d3..a708de8 100644
--- a/clang/unittests/Lex/LexerTest.cpp
+++ b/clang/unittests/Lex/LexerTest.cpp
@@ -18,6 +18,8 @@
 #include "clang/Basic/TargetOptions.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/HeaderSearchOptions.h"
+#include "clang/Lex/MacroArgs.h"
+#include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/ModuleLoader.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
@@ -41,26 +43,33 @@ protected:
     Target = TargetInfo::CreateTargetInfo(Diags, TargetOpts);
   }
 
-  std::vector<Token> Lex(StringRef Source) {
+  std::unique_ptr<Preprocessor> CreatePP(StringRef Source,
+                                         TrivialModuleLoader &ModLoader) {
     std::unique_ptr<llvm::MemoryBuffer> Buf =
         llvm::MemoryBuffer::getMemBuffer(Source);
     SourceMgr.setMainFileID(SourceMgr.createFileID(std::move(Buf)));
 
-    TrivialModuleLoader ModLoader;
     MemoryBufferCache PCMCache;
     HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
                             Diags, LangOpts, Target.get());
-    Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
-                    SourceMgr, PCMCache, HeaderInfo, ModLoader,
-                    /*IILookup =*/nullptr,
-                    /*OwnsHeaderSearch =*/false);
-    PP.Initialize(*Target);
-    PP.EnterMainSourceFile();
+    std::unique_ptr<Preprocessor> PP = llvm::make_unique<Preprocessor>(
+        std::make_shared<PreprocessorOptions>(), Diags, LangOpts, SourceMgr,
+        PCMCache, HeaderInfo, ModLoader,
+        /*IILookup =*/nullptr,
+        /*OwnsHeaderSearch =*/false);
+    PP->Initialize(*Target);
+    PP->EnterMainSourceFile();
+    return PP;
+  }
+
+  std::vector<Token> Lex(StringRef Source) {
+    TrivialModuleLoader ModLoader;
+    auto PP = CreatePP(Source, ModLoader);
 
     std::vector<Token> toks;
     while (1) {
       Token tok;
-      PP.Lex(tok);
+      PP->Lex(tok);
       if (tok.is(tok::eof))
         break;
       toks.push_back(tok);
@@ -365,4 +374,48 @@ TEST_F(LexerTest, DontMergeMacroArgsFromDifferentMacroFiles) {
   EXPECT_EQ(SourceMgr.getFileIDSize(SourceMgr.getFileID(helper1ArgLoc)), 8U);
 }
 
+TEST_F(LexerTest, DontOverallocateStringifyArgs) {
+  TrivialModuleLoader ModLoader;
+  auto PP = CreatePP("\"StrArg\", 5, 'C'", ModLoader);
+
+  llvm::BumpPtrAllocator Allocator;
+  std::array<IdentifierInfo *, 3> ArgList;
+  MacroInfo *MI = PP->AllocateMacroInfo({});
+  MI->setIsFunctionLike();
+  MI->setArgumentList(ArgList, Allocator);
+  EXPECT_EQ(3, MI->getNumArgs());
+  EXPECT_TRUE(MI->isFunctionLike());
+
+  Token Eof;
+  Eof.setKind(tok::eof);
+  std::vector<Token> ArgTokens;
+  while (1) {
+    Token tok;
+    PP->Lex(tok);
+    if (tok.is(tok::eof)) {
+      ArgTokens.push_back(Eof);
+      break;
+    }
+    if (tok.is(tok::comma))
+      ArgTokens.push_back(Eof);
+    else
+      ArgTokens.push_back(tok);
+  }
+
+  MacroArgs *MA = MacroArgs::create(MI, ArgTokens, false, *PP);
+  Token Result = MA->getStringifiedArgument(0, *PP, {}, {});
+  EXPECT_EQ(tok::string_literal, Result.getKind());
+  EXPECT_STREQ("\"\\\"StrArg\\\"\"", Result.getLiteralData());
+  Result = MA->getStringifiedArgument(1, *PP, {}, {});
+  EXPECT_EQ(tok::string_literal, Result.getKind());
+  EXPECT_STREQ("\"5\"", Result.getLiteralData());
+  Result = MA->getStringifiedArgument(2, *PP, {}, {});
+  EXPECT_EQ(tok::string_literal, Result.getKind());
+  EXPECT_STREQ("\"'C'\"", Result.getLiteralData());
+#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
+  EXPECT_DEATH(MA->getStringifiedArgument(3, *PP, {}, {}),
+               "Invalid argument number!");
+#endif
+}
+
 } // anonymous namespace