Index: include/clang/Lex/MacroArgs.h =================================================================== --- include/clang/Lex/MacroArgs.h +++ include/clang/Lex/MacroArgs.h @@ -53,9 +53,12 @@ /// Preprocessor owns which we use to avoid thrashing malloc/free. MacroArgs *ArgCache; - MacroArgs(unsigned NumToks, bool varargsElided) - : NumUnexpArgTokens(NumToks), VarargsElided(varargsElided), - ArgCache(nullptr) {} + /// MacroArgs - The number of arguments the invoked macro expects. + unsigned NumMacroArgs; + + MacroArgs(unsigned NumToks, bool varargsElided, unsigned MacroArgs) + : NumUnexpArgTokens(NumToks), VarargsElided(varargsElided), + ArgCache(nullptr), NumMacroArgs(MacroArgs) {} ~MacroArgs() = default; public: @@ -94,10 +97,9 @@ SourceLocation ExpansionLocStart, SourceLocation ExpansionLocEnd); - /// getNumArguments - Return the number of arguments passed into this macro - /// invocation. - unsigned getNumArguments() const { return NumUnexpArgTokens; } - + /// getNumMacroArguments - Return the number of arguments the invoked macro + /// expects. + unsigned getNumMacroArguments() const { return NumMacroArgs; } /// isVarargsElidedUse - Return true if this is a C99 style varargs macro /// invocation and there was no argument specified for the "..." argument. If Index: lib/Lex/MacroArgs.cpp =================================================================== --- lib/Lex/MacroArgs.cpp +++ lib/Lex/MacroArgs.cpp @@ -44,20 +44,22 @@ // Otherwise, use the best fit. ClosestMatch = (*Entry)->NumUnexpArgTokens; } - + MacroArgs *Result; if (!ResultEnt) { // Allocate memory for a MacroArgs object with the lexer tokens at the end. - Result = (MacroArgs*)malloc(sizeof(MacroArgs) + - UnexpArgTokens.size() * sizeof(Token)); + Result = (MacroArgs *)malloc(sizeof(MacroArgs) + + UnexpArgTokens.size() * sizeof(Token)); // Construct the MacroArgs object. - new (Result) MacroArgs(UnexpArgTokens.size(), VarargsElided); + new (Result) + MacroArgs(UnexpArgTokens.size(), VarargsElided, MI->getNumArgs()); } else { Result = *ResultEnt; // Unlink this node from the preprocessors singly linked list. *ResultEnt = Result->ArgCache; Result->NumUnexpArgTokens = UnexpArgTokens.size(); Result->VarargsElided = VarargsElided; + Result->NumMacroArgs = MI->getNumArgs(); } // Copy the actual unexpanded tokens to immediately after the result ptr. @@ -298,12 +300,10 @@ Preprocessor &PP, SourceLocation ExpansionLocStart, SourceLocation ExpansionLocEnd) { - assert(ArgNo < NumUnexpArgTokens && "Invalid argument number!"); - if (StringifiedArgs.empty()) { - StringifiedArgs.resize(getNumArguments()); - memset((void*)&StringifiedArgs[0], 0, - sizeof(StringifiedArgs[0])*getNumArguments()); - } + assert(ArgNo < getNumMacroArguments() && "Invalid argument number!"); + if (StringifiedArgs.empty()) + StringifiedArgs.resize(getNumMacroArguments(), {}); + if (StringifiedArgs[ArgNo].isNot(tok::string_literal)) StringifiedArgs[ArgNo] = StringifyArgument(getUnexpArgument(ArgNo), PP, /*Charify=*/false, Index: unittests/Lex/LexerTest.cpp =================================================================== --- unittests/Lex/LexerTest.cpp +++ unittests/Lex/LexerTest.cpp @@ -21,6 +21,8 @@ #include "clang/Lex/ModuleLoader.h" #include "clang/Lex/Preprocessor.h" #include "clang/Lex/PreprocessorOptions.h" +#include "clang/Lex/MacroArgs.h" +#include "clang/Lex/MacroInfo.h" #include "gtest/gtest.h" using namespace clang; @@ -41,26 +43,31 @@ Target = TargetInfo::CreateTargetInfo(Diags, TargetOpts); } - std::vector Lex(StringRef Source) { + std::unique_ptr CreatePP(StringRef Source, TrivialModuleLoader& ModLoader) { std::unique_ptr Buf = llvm::MemoryBuffer::getMemBuffer(Source); SourceMgr.setMainFileID(SourceMgr.createFileID(std::move(Buf))); - TrivialModuleLoader ModLoader; MemoryBufferCache PCMCache; HeaderSearch HeaderInfo(std::make_shared(), SourceMgr, Diags, LangOpts, Target.get()); - Preprocessor PP(std::make_shared(), Diags, LangOpts, - SourceMgr, PCMCache, HeaderInfo, ModLoader, - /*IILookup =*/nullptr, - /*OwnsHeaderSearch =*/false); - PP.Initialize(*Target); - PP.EnterMainSourceFile(); + std::unique_ptr PP = llvm::make_unique( + std::make_shared(), Diags, LangOpts, SourceMgr, + PCMCache, HeaderInfo, ModLoader, + /*IILookup =*/nullptr, + /*OwnsHeaderSearch =*/false); + PP->Initialize(*Target); + PP->EnterMainSourceFile(); + return PP; + } + std::vector Lex(StringRef Source) { + TrivialModuleLoader ModLoader; + auto PP = CreatePP(Source, ModLoader); std::vector toks; while (1) { Token tok; - PP.Lex(tok); + PP->Lex(tok); if (tok.is(tok::eof)) break; toks.push_back(tok); @@ -365,4 +372,48 @@ 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 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 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