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