xref: /llvm-project/clang/unittests/Basic/SourceManagerTest.cpp (revision 41b51007e6376cba72b00fb655a63b06c554d4e1)
1 //===- unittests/Basic/SourceManagerTest.cpp ------ SourceManager tests ---===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "clang/Basic/SourceManager.h"
10 #include "clang/Basic/Diagnostic.h"
11 #include "clang/Basic/DiagnosticOptions.h"
12 #include "clang/Basic/FileManager.h"
13 #include "clang/Basic/LangOptions.h"
14 #include "clang/Basic/TargetInfo.h"
15 #include "clang/Basic/TargetOptions.h"
16 #include "clang/Lex/HeaderSearch.h"
17 #include "clang/Lex/HeaderSearchOptions.h"
18 #include "clang/Lex/ModuleLoader.h"
19 #include "clang/Lex/Preprocessor.h"
20 #include "clang/Lex/PreprocessorOptions.h"
21 #include "llvm/ADT/SmallString.h"
22 #include "llvm/Config/llvm-config.h"
23 #include "llvm/Support/Process.h"
24 #include "gtest/gtest.h"
25 #include <cstddef>
26 
27 using namespace clang;
28 
29 namespace {
30 
31 // The test fixture.
32 class SourceManagerTest : public ::testing::Test {
33 protected:
34   SourceManagerTest()
35     : FileMgr(FileMgrOpts),
36       DiagID(new DiagnosticIDs()),
37       Diags(DiagID, new DiagnosticOptions, new IgnoringDiagConsumer()),
38       SourceMgr(Diags, FileMgr),
39       TargetOpts(new TargetOptions) {
40     TargetOpts->Triple = "x86_64-apple-darwin11.1.0";
41     Target = TargetInfo::CreateTargetInfo(Diags, TargetOpts);
42   }
43 
44   FileSystemOptions FileMgrOpts;
45   FileManager FileMgr;
46   IntrusiveRefCntPtr<DiagnosticIDs> DiagID;
47   DiagnosticsEngine Diags;
48   SourceManager SourceMgr;
49   LangOptions LangOpts;
50   std::shared_ptr<TargetOptions> TargetOpts;
51   IntrusiveRefCntPtr<TargetInfo> Target;
52 };
53 
54 TEST_F(SourceManagerTest, isInMemoryBuffersNoSourceLocationInfo) {
55   // Check for invalid source location for each method
56   SourceLocation LocEmpty;
57   bool isWrittenInBuiltInFileFalse = SourceMgr.isWrittenInBuiltinFile(LocEmpty);
58   bool isWrittenInCommandLineFileFalse =
59       SourceMgr.isWrittenInCommandLineFile(LocEmpty);
60   bool isWrittenInScratchSpaceFalse =
61       SourceMgr.isWrittenInScratchSpace(LocEmpty);
62 
63   EXPECT_FALSE(isWrittenInBuiltInFileFalse);
64   EXPECT_FALSE(isWrittenInCommandLineFileFalse);
65   EXPECT_FALSE(isWrittenInScratchSpaceFalse);
66 
67   // Check for valid source location per filename for each method
68   const char *Source = "int x";
69 
70   std::unique_ptr<llvm::MemoryBuffer> BuiltInBuf =
71       llvm::MemoryBuffer::getMemBuffer(Source);
72   const FileEntry *BuiltInFile =
73       FileMgr.getVirtualFile("<built-in>", BuiltInBuf->getBufferSize(), 0);
74   SourceMgr.overrideFileContents(BuiltInFile, std::move(BuiltInBuf));
75   FileID BuiltInFileID =
76       SourceMgr.getOrCreateFileID(BuiltInFile, SrcMgr::C_User);
77   SourceMgr.setMainFileID(BuiltInFileID);
78   SourceLocation LocBuiltIn =
79       SourceMgr.getLocForStartOfFile(SourceMgr.getMainFileID());
80   bool isWrittenInBuiltInFileTrue =
81       SourceMgr.isWrittenInBuiltinFile(LocBuiltIn);
82 
83   std::unique_ptr<llvm::MemoryBuffer> CommandLineBuf =
84       llvm::MemoryBuffer::getMemBuffer(Source);
85   const FileEntry *CommandLineFile = FileMgr.getVirtualFile(
86       "<command line>", CommandLineBuf->getBufferSize(), 0);
87   SourceMgr.overrideFileContents(CommandLineFile, std::move(CommandLineBuf));
88   FileID CommandLineFileID =
89       SourceMgr.getOrCreateFileID(CommandLineFile, SrcMgr::C_User);
90   SourceMgr.setMainFileID(CommandLineFileID);
91   SourceLocation LocCommandLine =
92       SourceMgr.getLocForStartOfFile(SourceMgr.getMainFileID());
93   bool isWrittenInCommandLineFileTrue =
94       SourceMgr.isWrittenInCommandLineFile(LocCommandLine);
95 
96   std::unique_ptr<llvm::MemoryBuffer> ScratchSpaceBuf =
97       llvm::MemoryBuffer::getMemBuffer(Source);
98   const FileEntry *ScratchSpaceFile = FileMgr.getVirtualFile(
99       "<scratch space>", ScratchSpaceBuf->getBufferSize(), 0);
100   SourceMgr.overrideFileContents(ScratchSpaceFile, std::move(ScratchSpaceBuf));
101   FileID ScratchSpaceFileID =
102       SourceMgr.getOrCreateFileID(ScratchSpaceFile, SrcMgr::C_User);
103   SourceMgr.setMainFileID(ScratchSpaceFileID);
104   SourceLocation LocScratchSpace =
105       SourceMgr.getLocForStartOfFile(SourceMgr.getMainFileID());
106   bool isWrittenInScratchSpaceTrue =
107       SourceMgr.isWrittenInScratchSpace(LocScratchSpace);
108 
109   EXPECT_TRUE(isWrittenInBuiltInFileTrue);
110   EXPECT_TRUE(isWrittenInCommandLineFileTrue);
111   EXPECT_TRUE(isWrittenInScratchSpaceTrue);
112 }
113 
114 TEST_F(SourceManagerTest, isInSystemHeader) {
115   // Check for invalid source location
116   SourceLocation LocEmpty;
117   bool isInSystemHeaderFalse = SourceMgr.isInSystemHeader(LocEmpty);
118   ASSERT_FALSE(isInSystemHeaderFalse);
119 }
120 
121 TEST_F(SourceManagerTest, isBeforeInTranslationUnit) {
122   const char *source =
123     "#define M(x) [x]\n"
124     "M(foo)";
125   std::unique_ptr<llvm::MemoryBuffer> Buf =
126       llvm::MemoryBuffer::getMemBuffer(source);
127   FileID mainFileID = SourceMgr.createFileID(std::move(Buf));
128   SourceMgr.setMainFileID(mainFileID);
129 
130   TrivialModuleLoader ModLoader;
131   HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
132                           Diags, LangOpts, &*Target);
133   Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
134                   SourceMgr, HeaderInfo, ModLoader,
135                   /*IILookup =*/nullptr,
136                   /*OwnsHeaderSearch =*/false);
137   PP.Initialize(*Target);
138   PP.EnterMainSourceFile();
139 
140   std::vector<Token> toks;
141   while (1) {
142     Token tok;
143     PP.Lex(tok);
144     if (tok.is(tok::eof))
145       break;
146     toks.push_back(tok);
147   }
148 
149   // Make sure we got the tokens that we expected.
150   ASSERT_EQ(3U, toks.size());
151   ASSERT_EQ(tok::l_square, toks[0].getKind());
152   ASSERT_EQ(tok::identifier, toks[1].getKind());
153   ASSERT_EQ(tok::r_square, toks[2].getKind());
154 
155   SourceLocation lsqrLoc = toks[0].getLocation();
156   SourceLocation idLoc = toks[1].getLocation();
157   SourceLocation rsqrLoc = toks[2].getLocation();
158 
159   SourceLocation macroExpStartLoc = SourceMgr.translateLineCol(mainFileID, 2, 1);
160   SourceLocation macroExpEndLoc = SourceMgr.translateLineCol(mainFileID, 2, 6);
161   ASSERT_TRUE(macroExpStartLoc.isFileID());
162   ASSERT_TRUE(macroExpEndLoc.isFileID());
163 
164   SmallString<32> str;
165   ASSERT_EQ("M", PP.getSpelling(macroExpStartLoc, str));
166   ASSERT_EQ(")", PP.getSpelling(macroExpEndLoc, str));
167 
168   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(lsqrLoc, idLoc));
169   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(idLoc, rsqrLoc));
170   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(macroExpStartLoc, idLoc));
171   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(idLoc, macroExpEndLoc));
172 }
173 
174 TEST_F(SourceManagerTest, isBeforeInTranslationUnitWithTokenSplit) {
175   const char *main = R"cpp(
176     #define ID(X) X
177     ID(
178       ID(a >> b)
179       c
180     )
181   )cpp";
182 
183   SourceMgr.setMainFileID(
184       SourceMgr.createFileID(llvm::MemoryBuffer::getMemBuffer(main)));
185 
186   TrivialModuleLoader ModLoader;
187   HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
188                           Diags, LangOpts, &*Target);
189   Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
190                   SourceMgr, HeaderInfo, ModLoader,
191                   /*IILookup =*/nullptr,
192                   /*OwnsHeaderSearch =*/false);
193   PP.Initialize(*Target);
194   PP.EnterMainSourceFile();
195   llvm::SmallString<8> Scratch;
196 
197   std::vector<Token> toks;
198   while (1) {
199     Token tok;
200     PP.Lex(tok);
201     if (tok.is(tok::eof))
202       break;
203     toks.push_back(tok);
204   }
205 
206   // Make sure we got the tokens that we expected.
207   ASSERT_EQ(4U, toks.size()) << "a >> b c";
208   // Sanity check their order.
209   for (unsigned I = 0; I < toks.size() - 1; ++I) {
210     EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(toks[I].getLocation(),
211                                                     toks[I + 1].getLocation()));
212     EXPECT_FALSE(SourceMgr.isBeforeInTranslationUnit(toks[I + 1].getLocation(),
213                                                      toks[I].getLocation()));
214   }
215 
216   // Split the >> into two > tokens, as happens when parsing nested templates.
217   unsigned RightShiftIndex = 1;
218   SourceLocation RightShift = toks[RightShiftIndex].getLocation();
219   EXPECT_EQ(">>", Lexer::getSpelling(SourceMgr.getSpellingLoc(RightShift),
220                                      Scratch, SourceMgr, LangOpts));
221   SourceLocation Greater1 = PP.SplitToken(RightShift, /*Length=*/1);
222   SourceLocation Greater2 = RightShift.getLocWithOffset(1);
223   EXPECT_TRUE(Greater1.isMacroID());
224   EXPECT_EQ(">", Lexer::getSpelling(SourceMgr.getSpellingLoc(Greater1), Scratch,
225                                     SourceMgr, LangOpts));
226   EXPECT_EQ(">", Lexer::getSpelling(SourceMgr.getSpellingLoc(Greater2), Scratch,
227                                     SourceMgr, LangOpts));
228   EXPECT_EQ(SourceMgr.getImmediateExpansionRange(Greater1).getBegin(),
229             RightShift);
230 
231   for (unsigned I = 0; I < toks.size(); ++I) {
232     SCOPED_TRACE("Token " + std::to_string(I));
233     // Right-shift is the parent of Greater1, so it compares less.
234     EXPECT_EQ(
235         SourceMgr.isBeforeInTranslationUnit(toks[I].getLocation(), Greater1),
236         I <= RightShiftIndex);
237     EXPECT_EQ(
238         SourceMgr.isBeforeInTranslationUnit(toks[I].getLocation(), Greater2),
239         I <= RightShiftIndex);
240     EXPECT_EQ(
241         SourceMgr.isBeforeInTranslationUnit(Greater1, toks[I].getLocation()),
242         RightShiftIndex < I);
243     EXPECT_EQ(
244         SourceMgr.isBeforeInTranslationUnit(Greater2, toks[I].getLocation()),
245         RightShiftIndex < I);
246   }
247   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(Greater1, Greater2));
248   EXPECT_FALSE(SourceMgr.isBeforeInTranslationUnit(Greater2, Greater1));
249 }
250 
251 TEST_F(SourceManagerTest, getColumnNumber) {
252   const char *Source =
253     "int x;\n"
254     "int y;";
255 
256   std::unique_ptr<llvm::MemoryBuffer> Buf =
257       llvm::MemoryBuffer::getMemBuffer(Source);
258   FileID MainFileID = SourceMgr.createFileID(std::move(Buf));
259   SourceMgr.setMainFileID(MainFileID);
260 
261   bool Invalid;
262 
263   Invalid = false;
264   EXPECT_EQ(1U, SourceMgr.getColumnNumber(MainFileID, 0, &Invalid));
265   EXPECT_TRUE(!Invalid);
266 
267   Invalid = false;
268   EXPECT_EQ(5U, SourceMgr.getColumnNumber(MainFileID, 4, &Invalid));
269   EXPECT_TRUE(!Invalid);
270 
271   Invalid = false;
272   EXPECT_EQ(1U, SourceMgr.getColumnNumber(MainFileID, 7, &Invalid));
273   EXPECT_TRUE(!Invalid);
274 
275   Invalid = false;
276   EXPECT_EQ(5U, SourceMgr.getColumnNumber(MainFileID, 11, &Invalid));
277   EXPECT_TRUE(!Invalid);
278 
279   Invalid = false;
280   EXPECT_EQ(7U, SourceMgr.getColumnNumber(MainFileID, strlen(Source),
281                                          &Invalid));
282   EXPECT_TRUE(!Invalid);
283 
284   Invalid = false;
285   SourceMgr.getColumnNumber(MainFileID, strlen(Source)+1, &Invalid);
286   EXPECT_TRUE(Invalid);
287 
288   // Test invalid files
289   Invalid = false;
290   SourceMgr.getColumnNumber(FileID(), 0, &Invalid);
291   EXPECT_TRUE(Invalid);
292 
293   Invalid = false;
294   SourceMgr.getColumnNumber(FileID(), 1, &Invalid);
295   EXPECT_TRUE(Invalid);
296 
297   // Test with no invalid flag.
298   EXPECT_EQ(1U, SourceMgr.getColumnNumber(MainFileID, 0, nullptr));
299 }
300 
301 TEST_F(SourceManagerTest, locationPrintTest) {
302   const char *header = "#define IDENTITY(x) x\n";
303 
304   const char *Source = "int x;\n"
305                        "include \"test-header.h\"\n"
306                        "IDENTITY(int y);\n"
307                        "int z;";
308 
309   std::unique_ptr<llvm::MemoryBuffer> HeaderBuf =
310       llvm::MemoryBuffer::getMemBuffer(header);
311   std::unique_ptr<llvm::MemoryBuffer> Buf =
312       llvm::MemoryBuffer::getMemBuffer(Source);
313 
314   const FileEntry *SourceFile =
315       FileMgr.getVirtualFile("/mainFile.cpp", Buf->getBufferSize(), 0);
316   SourceMgr.overrideFileContents(SourceFile, std::move(Buf));
317 
318   const FileEntry *HeaderFile =
319       FileMgr.getVirtualFile("/test-header.h", HeaderBuf->getBufferSize(), 0);
320   SourceMgr.overrideFileContents(HeaderFile, std::move(HeaderBuf));
321 
322   FileID MainFileID = SourceMgr.getOrCreateFileID(SourceFile, SrcMgr::C_User);
323   FileID HeaderFileID = SourceMgr.getOrCreateFileID(HeaderFile, SrcMgr::C_User);
324   SourceMgr.setMainFileID(MainFileID);
325 
326   auto BeginLoc = SourceMgr.getLocForStartOfFile(MainFileID);
327   auto EndLoc = SourceMgr.getLocForEndOfFile(MainFileID);
328 
329   auto BeginEOLLoc = SourceMgr.translateLineCol(MainFileID, 1, 7);
330 
331   auto HeaderLoc = SourceMgr.getLocForStartOfFile(HeaderFileID);
332 
333   EXPECT_EQ(BeginLoc.printToString(SourceMgr), "/mainFile.cpp:1:1");
334   EXPECT_EQ(EndLoc.printToString(SourceMgr), "/mainFile.cpp:4:7");
335 
336   EXPECT_EQ(BeginEOLLoc.printToString(SourceMgr), "/mainFile.cpp:1:7");
337   EXPECT_EQ(HeaderLoc.printToString(SourceMgr), "/test-header.h:1:1");
338 
339   EXPECT_EQ(SourceRange(BeginLoc, BeginLoc).printToString(SourceMgr),
340             "</mainFile.cpp:1:1>");
341   EXPECT_EQ(SourceRange(BeginLoc, BeginEOLLoc).printToString(SourceMgr),
342             "</mainFile.cpp:1:1, col:7>");
343   EXPECT_EQ(SourceRange(BeginLoc, EndLoc).printToString(SourceMgr),
344             "</mainFile.cpp:1:1, line:4:7>");
345   EXPECT_EQ(SourceRange(BeginLoc, HeaderLoc).printToString(SourceMgr),
346             "</mainFile.cpp:1:1, /test-header.h:1:1>");
347 }
348 
349 TEST_F(SourceManagerTest, getInvalidBOM) {
350   ASSERT_EQ(SrcMgr::ContentCache::getInvalidBOM(""), nullptr);
351   ASSERT_EQ(SrcMgr::ContentCache::getInvalidBOM("\x00\x00\x00"), nullptr);
352   ASSERT_EQ(SrcMgr::ContentCache::getInvalidBOM("\xFF\xFF\xFF"), nullptr);
353   ASSERT_EQ(SrcMgr::ContentCache::getInvalidBOM("#include <iostream>"),
354             nullptr);
355 
356   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
357                 "\xFE\xFF#include <iostream>")),
358             "UTF-16 (BE)");
359   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
360                 "\xFF\xFE#include <iostream>")),
361             "UTF-16 (LE)");
362   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
363                 "\x2B\x2F\x76#include <iostream>")),
364             "UTF-7");
365   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
366                 "\xF7\x64\x4C#include <iostream>")),
367             "UTF-1");
368   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
369                 "\xDD\x73\x66\x73#include <iostream>")),
370             "UTF-EBCDIC");
371   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
372                 "\x0E\xFE\xFF#include <iostream>")),
373             "SCSU");
374   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
375                 "\xFB\xEE\x28#include <iostream>")),
376             "BOCU-1");
377   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
378                 "\x84\x31\x95\x33#include <iostream>")),
379             "GB-18030");
380   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
381                 llvm::StringLiteral::withInnerNUL(
382                     "\x00\x00\xFE\xFF#include <iostream>"))),
383             "UTF-32 (BE)");
384   ASSERT_EQ(StringRef(SrcMgr::ContentCache::getInvalidBOM(
385                 llvm::StringLiteral::withInnerNUL(
386                     "\xFF\xFE\x00\x00#include <iostream>"))),
387             "UTF-32 (LE)");
388 }
389 
390 // Regression test - there was an out of bound access for buffers not terminated by zero.
391 TEST_F(SourceManagerTest, getLineNumber) {
392   const unsigned pageSize = llvm::sys::Process::getPageSizeEstimate();
393   std::unique_ptr<char[]> source(new char[pageSize]);
394   for(unsigned i = 0; i < pageSize; ++i) {
395     source[i] = 'a';
396   }
397 
398   std::unique_ptr<llvm::MemoryBuffer> Buf =
399       llvm::MemoryBuffer::getMemBuffer(
400         llvm::MemoryBufferRef(
401           llvm::StringRef(source.get(), 3), "whatever"
402         ),
403         false
404       );
405 
406   FileID mainFileID = SourceMgr.createFileID(std::move(Buf));
407   SourceMgr.setMainFileID(mainFileID);
408 
409   ASSERT_NO_FATAL_FAILURE(SourceMgr.getLineNumber(mainFileID, 1, nullptr));
410 }
411 
412 #if defined(LLVM_ON_UNIX)
413 
414 TEST_F(SourceManagerTest, getMacroArgExpandedLocation) {
415   const char *header =
416     "#define FM(x,y) x\n";
417 
418   const char *main =
419     "#include \"/test-header.h\"\n"
420     "#define VAL 0\n"
421     "FM(VAL,0)\n"
422     "FM(0,VAL)\n"
423     "FM(FM(0,VAL),0)\n"
424     "#define CONCAT(X, Y) X##Y\n"
425     "CONCAT(1,1)\n";
426 
427   std::unique_ptr<llvm::MemoryBuffer> HeaderBuf =
428       llvm::MemoryBuffer::getMemBuffer(header);
429   std::unique_ptr<llvm::MemoryBuffer> MainBuf =
430       llvm::MemoryBuffer::getMemBuffer(main);
431   FileID mainFileID = SourceMgr.createFileID(std::move(MainBuf));
432   SourceMgr.setMainFileID(mainFileID);
433 
434   const FileEntry *headerFile = FileMgr.getVirtualFile("/test-header.h",
435                                                  HeaderBuf->getBufferSize(), 0);
436   SourceMgr.overrideFileContents(headerFile, std::move(HeaderBuf));
437 
438   TrivialModuleLoader ModLoader;
439   HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
440                           Diags, LangOpts, &*Target);
441 
442   Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
443                   SourceMgr, HeaderInfo, ModLoader,
444                   /*IILookup =*/nullptr,
445                   /*OwnsHeaderSearch =*/false);
446   // Ensure we can get expanded locations in presence of implicit includes.
447   // These are different than normal includes since predefines buffer doesn't
448   // have a valid insertion location.
449   PP.setPredefines("#include \"/implicit-header.h\"");
450   FileMgr.getVirtualFile("/implicit-header.h", 0, 0);
451   PP.Initialize(*Target);
452   PP.EnterMainSourceFile();
453 
454   std::vector<Token> toks;
455   while (1) {
456     Token tok;
457     PP.Lex(tok);
458     if (tok.is(tok::eof))
459       break;
460     toks.push_back(tok);
461   }
462 
463   // Make sure we got the tokens that we expected.
464   ASSERT_EQ(4U, toks.size());
465   ASSERT_EQ(tok::numeric_constant, toks[0].getKind());
466   ASSERT_EQ(tok::numeric_constant, toks[1].getKind());
467   ASSERT_EQ(tok::numeric_constant, toks[2].getKind());
468   ASSERT_EQ(tok::numeric_constant, toks[3].getKind());
469 
470   SourceLocation defLoc = SourceMgr.translateLineCol(mainFileID, 2, 13);
471   SourceLocation loc1 = SourceMgr.translateLineCol(mainFileID, 3, 8);
472   SourceLocation loc2 = SourceMgr.translateLineCol(mainFileID, 4, 4);
473   SourceLocation loc3 = SourceMgr.translateLineCol(mainFileID, 5, 7);
474   SourceLocation defLoc2 = SourceMgr.translateLineCol(mainFileID, 6, 22);
475   defLoc = SourceMgr.getMacroArgExpandedLocation(defLoc);
476   loc1 = SourceMgr.getMacroArgExpandedLocation(loc1);
477   loc2 = SourceMgr.getMacroArgExpandedLocation(loc2);
478   loc3 = SourceMgr.getMacroArgExpandedLocation(loc3);
479   defLoc2 = SourceMgr.getMacroArgExpandedLocation(defLoc2);
480 
481   EXPECT_TRUE(defLoc.isFileID());
482   EXPECT_TRUE(loc1.isFileID());
483   EXPECT_TRUE(SourceMgr.isMacroArgExpansion(loc2));
484   EXPECT_TRUE(SourceMgr.isMacroArgExpansion(loc3));
485   EXPECT_EQ(loc2, toks[1].getLocation());
486   EXPECT_EQ(loc3, toks[2].getLocation());
487   EXPECT_TRUE(defLoc2.isFileID());
488 }
489 
490 namespace {
491 
492 struct MacroAction {
493   enum Kind { kExpansion, kDefinition, kUnDefinition};
494 
495   SourceLocation Loc;
496   std::string Name;
497   unsigned MAKind : 3;
498 
499   MacroAction(SourceLocation Loc, StringRef Name, unsigned K)
500       : Loc(Loc), Name(std::string(Name)), MAKind(K) {}
501 
502   bool isExpansion() const { return MAKind == kExpansion; }
503   bool isDefinition() const { return MAKind & kDefinition; }
504   bool isUnDefinition() const { return MAKind & kUnDefinition; }
505 };
506 
507 class MacroTracker : public PPCallbacks {
508   std::vector<MacroAction> &Macros;
509 
510 public:
511   explicit MacroTracker(std::vector<MacroAction> &Macros) : Macros(Macros) { }
512 
513   void MacroDefined(const Token &MacroNameTok,
514                     const MacroDirective *MD) override {
515     Macros.push_back(MacroAction(MD->getLocation(),
516                                  MacroNameTok.getIdentifierInfo()->getName(),
517                                  MacroAction::kDefinition));
518   }
519   void MacroUndefined(const Token &MacroNameTok,
520                       const MacroDefinition &MD,
521                       const MacroDirective  *UD) override {
522     Macros.push_back(
523         MacroAction(UD ? UD->getLocation() : SourceLocation(),
524                     MacroNameTok.getIdentifierInfo()->getName(),
525                     UD ? MacroAction::kDefinition | MacroAction::kUnDefinition
526                        : MacroAction::kUnDefinition));
527   }
528   void MacroExpands(const Token &MacroNameTok, const MacroDefinition &MD,
529                     SourceRange Range, const MacroArgs *Args) override {
530     Macros.push_back(MacroAction(MacroNameTok.getLocation(),
531                                  MacroNameTok.getIdentifierInfo()->getName(),
532                                  MacroAction::kExpansion));
533   }
534 };
535 
536 }
537 
538 TEST_F(SourceManagerTest, isBeforeInTranslationUnitWithMacroInInclude) {
539   const char *header =
540     "#define MACRO_IN_INCLUDE 0\n"
541     "#define MACRO_DEFINED\n"
542     "#undef MACRO_DEFINED\n"
543     "#undef MACRO_UNDEFINED\n";
544 
545   const char *main =
546     "#define M(x) x\n"
547     "#define INC \"/test-header.h\"\n"
548     "#include M(INC)\n"
549     "#define INC2 </test-header.h>\n"
550     "#include M(INC2)\n";
551 
552   std::unique_ptr<llvm::MemoryBuffer> HeaderBuf =
553       llvm::MemoryBuffer::getMemBuffer(header);
554   std::unique_ptr<llvm::MemoryBuffer> MainBuf =
555       llvm::MemoryBuffer::getMemBuffer(main);
556   SourceMgr.setMainFileID(SourceMgr.createFileID(std::move(MainBuf)));
557 
558   const FileEntry *headerFile = FileMgr.getVirtualFile("/test-header.h",
559                                                  HeaderBuf->getBufferSize(), 0);
560   SourceMgr.overrideFileContents(headerFile, std::move(HeaderBuf));
561 
562   TrivialModuleLoader ModLoader;
563   HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
564                           Diags, LangOpts, &*Target);
565   Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
566                   SourceMgr, HeaderInfo, ModLoader,
567                   /*IILookup =*/nullptr,
568                   /*OwnsHeaderSearch =*/false);
569   PP.Initialize(*Target);
570 
571   std::vector<MacroAction> Macros;
572   PP.addPPCallbacks(std::make_unique<MacroTracker>(Macros));
573 
574   PP.EnterMainSourceFile();
575 
576   std::vector<Token> toks;
577   while (1) {
578     Token tok;
579     PP.Lex(tok);
580     if (tok.is(tok::eof))
581       break;
582     toks.push_back(tok);
583   }
584 
585   // Make sure we got the tokens that we expected.
586   ASSERT_EQ(0U, toks.size());
587 
588   ASSERT_EQ(15U, Macros.size());
589   // #define M(x) x
590   ASSERT_TRUE(Macros[0].isDefinition());
591   ASSERT_EQ("M", Macros[0].Name);
592   // #define INC "/test-header.h"
593   ASSERT_TRUE(Macros[1].isDefinition());
594   ASSERT_EQ("INC", Macros[1].Name);
595   // M expansion in #include M(INC)
596   ASSERT_FALSE(Macros[2].isDefinition());
597   ASSERT_EQ("M", Macros[2].Name);
598   // INC expansion in #include M(INC)
599   ASSERT_TRUE(Macros[3].isExpansion());
600   ASSERT_EQ("INC", Macros[3].Name);
601   // #define MACRO_IN_INCLUDE 0
602   ASSERT_TRUE(Macros[4].isDefinition());
603   ASSERT_EQ("MACRO_IN_INCLUDE", Macros[4].Name);
604   // #define MACRO_DEFINED
605   ASSERT_TRUE(Macros[5].isDefinition());
606   ASSERT_FALSE(Macros[5].isUnDefinition());
607   ASSERT_EQ("MACRO_DEFINED", Macros[5].Name);
608   // #undef MACRO_DEFINED
609   ASSERT_TRUE(Macros[6].isDefinition());
610   ASSERT_TRUE(Macros[6].isUnDefinition());
611   ASSERT_EQ("MACRO_DEFINED", Macros[6].Name);
612   // #undef MACRO_UNDEFINED
613   ASSERT_FALSE(Macros[7].isDefinition());
614   ASSERT_TRUE(Macros[7].isUnDefinition());
615   ASSERT_EQ("MACRO_UNDEFINED", Macros[7].Name);
616   // #define INC2 </test-header.h>
617   ASSERT_TRUE(Macros[8].isDefinition());
618   ASSERT_EQ("INC2", Macros[8].Name);
619   // M expansion in #include M(INC2)
620   ASSERT_FALSE(Macros[9].isDefinition());
621   ASSERT_EQ("M", Macros[9].Name);
622   // INC2 expansion in #include M(INC2)
623   ASSERT_TRUE(Macros[10].isExpansion());
624   ASSERT_EQ("INC2", Macros[10].Name);
625   // #define MACRO_IN_INCLUDE 0
626   ASSERT_TRUE(Macros[11].isDefinition());
627   ASSERT_EQ("MACRO_IN_INCLUDE", Macros[11].Name);
628 
629   // The INC expansion in #include M(INC) comes before the first
630   // MACRO_IN_INCLUDE definition of the included file.
631   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(Macros[3].Loc, Macros[4].Loc));
632 
633   // The INC2 expansion in #include M(INC2) comes before the second
634   // MACRO_IN_INCLUDE definition of the included file.
635   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(Macros[10].Loc, Macros[11].Loc));
636 }
637 
638 TEST_F(SourceManagerTest, isMainFile) {
639   const char *Source = "int x;";
640 
641   std::unique_ptr<llvm::MemoryBuffer> Buf =
642       llvm::MemoryBuffer::getMemBuffer(Source);
643   const FileEntry *SourceFile =
644       FileMgr.getVirtualFile("mainFile.cpp", Buf->getBufferSize(), 0);
645   SourceMgr.overrideFileContents(SourceFile, std::move(Buf));
646 
647   std::unique_ptr<llvm::MemoryBuffer> Buf2 =
648       llvm::MemoryBuffer::getMemBuffer(Source);
649   const FileEntry *SecondFile =
650       FileMgr.getVirtualFile("file2.cpp", Buf2->getBufferSize(), 0);
651   SourceMgr.overrideFileContents(SecondFile, std::move(Buf2));
652 
653   FileID MainFileID = SourceMgr.getOrCreateFileID(SourceFile, SrcMgr::C_User);
654   SourceMgr.setMainFileID(MainFileID);
655 
656   EXPECT_TRUE(SourceMgr.isMainFile(*SourceFile));
657   EXPECT_TRUE(SourceMgr.isMainFile(*SourceFile));
658   EXPECT_FALSE(SourceMgr.isMainFile(*SecondFile));
659 }
660 
661 #endif
662 
663 } // anonymous namespace
664