[Preprocessor]Correct Macro-Arg allocation of StringifiedArguments,

correct getNumArguments

StringifiedArguments is allocated (resized) based on the size the 
getNumArguments function. However, this function ACTUALLY currently 
returns the amount of total UnexpArgTokens which is minimum the same as 
the new implementation of getNumMacroArguments, since empty/omitted arguments 
result in 1 UnexpArgToken, and included ones at minimum include 2 
(1 for the arg itself, 1 for eof).

This patch renames the otherwise unused getNumArguments to be more clear 
that it is the number of arguments that the Macro expects, and thus the maximum 
number that can be stringified. This patch also replaces the explicit memset 
(which results in value instantiation of the new tokens, PLUS clearing the 
memory) with brace initialization.

Differential Revision: https://reviews.llvm.org/D32046

llvm-svn: 305425

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