feat: collect embed/has_embed directive in PPCallbacks

Author: Guo-Shiyu

include/Compiler/Directive.h

@@ -127,13 +127,46 @@ struct Import {
     std::vector<clang::SourceLocation> name_locations;
 };
 
+struct Embed {
+    /// The file name in the embed directive, not including quotes or angle brackets.
+    llvm::StringRef file_name;
+
+    /// The actual file that may be embedded by this embed directive.
+    clang::OptionalFileEntryRef file;
+
+    /// Whether the file name is angled.
+    bool is_angled;
+
+    /// Location of the `#` token.
+    clang::SourceLocation loc;
+
+    /// TODO: Currently we do not store parameters of the embed directive.
+    /// See clang::LexEmbedParametersResult for details.
+};
+
+struct HasEmbed {
+    /// The file name in the embed directive, not including quotes or angle brackets.
+    llvm::StringRef file_name;
+
+    /// The actual file that may be embedded by this embed directive.
+    clang::OptionalFileEntryRef file;
+
+    /// Whether the file name is angled.
+    bool is_angled;
+
+    /// Location of the `__has_embed` token.
+    clang::SourceLocation loc;
+};
+
 struct Directive {
     std::vector<Include> includes;
     std::vector<HasInclude> has_includes;
     std::vector<Condition> conditions;
     std::vector<MacroRef> macros;
     std::vector<Pragma> pragmas;
     std::vector<Import> imports;
+    std::vector<Embed> embeds;
+    std::vector<HasEmbed> has_embeds;
 
     /// Tell preprocessor to collect directives information and store them in `directives`.
     static void attach(clang::Preprocessor& pp,

src/Compiler/Directive.cpp

@@ -57,14 +57,39 @@ class DirectiveCollector : public clang::PPCallbacks {
         }
 
         auto& directive = directives[sm.getFileID(loc)];
-        directive.macros.emplace_back(MacroRef{def, kind, loc});
+        directive.macros.push_back({def, kind, loc});
     }
 
 public:
     /// ============================================================================
     ///                         Rewritten Preprocessor Callbacks
     /// ============================================================================
 
+    void HasEmbed(clang::SourceLocation location,
+                  llvm::StringRef filename,
+                  bool is_angled,
+                  clang::OptionalFileEntryRef file) override {
+        directives[sm.getFileID(location)].has_embeds.push_back({
+            filename,
+            file,
+            is_angled,
+            location,
+        });
+    }
+
+    void EmbedDirective(clang::SourceLocation location,
+                        clang::StringRef filename,
+                        bool is_angled,
+                        clang::OptionalFileEntryRef file,
+                        const clang::LexEmbedParametersResult& Params) override {
+        directives[sm.getFileID(location)].embeds.push_back({
+            filename,
+            file,
+            is_angled,
+            location,
+        });
+    }
+
     void InclusionDirective(clang::SourceLocation hash_loc,
                             const clang::Token& include_tok,
                             llvm::StringRef,
@@ -80,7 +105,7 @@ class DirectiveCollector : public clang::PPCallbacks {
 
         /// An `IncludeDirective` call is always followed by either a `LexedFileChanged`
         /// or a `FileSkipped`. so we cannot get the file id of included file here.
-        directives[prev_fid].includes.emplace_back(Include{
+        directives[prev_fid].includes.push_back({
             .fid = {},
             .location = include_tok.getLocation(),
             .filename_range = filename_range.getAsRange(),
@@ -139,7 +164,10 @@ class DirectiveCollector : public clang::PPCallbacks {
             fid = sm.translateFile(*file);
         }
 
-        directives[sm.getFileID(location)].has_includes.emplace_back(fid, location);
+        directives[sm.getFileID(location)].has_includes.push_back({
+            fid,
+            location,
+        });
     }
 
     void PragmaDirective(clang::SourceLocation loc,
@@ -158,7 +186,7 @@ class DirectiveCollector : public clang::PPCallbacks {
                                                             : Pragma::Other;
 
         auto& directive = directives[fid];
-        directive.pragmas.emplace_back(Pragma{
+        directive.pragmas.push_back({
             that_line,
             kind,
             loc,

tests/unit/Compiler/Directive.cpp

@@ -11,6 +11,8 @@ suite<"Directive"> directive = [] {
     std::vector<Condition> conditions;
     std::vector<MacroRef> macros;
     std::vector<Pragma> pragmas;
+    std::vector<Embed> embeds;
+    std::vector<HasEmbed> has_embeds;
 
     using u32 = std::uint32_t;
 
@@ -24,6 +26,8 @@ suite<"Directive"> directive = [] {
         conditions = tester.unit->directives()[fid].conditions;
         macros = tester.unit->directives()[fid].macros;
         pragmas = tester.unit->directives()[fid].pragmas;
+        embeds = tester.unit->directives()[fid].embeds;
+        has_embeds = tester.unit->directives()[fid].has_embeds;
     };
 
     auto expect_include = [&](u32 index, llvm::StringRef position, llvm::StringRef path) {
@@ -51,6 +55,25 @@ suite<"Directive"> directive = [] {
         expect(that % target == path);
     };
 
+    auto expect_embed = [&](u32 index, llvm::StringRef position, llvm::StringRef filename) {
+        auto& embed = embeds[index];
+        auto [_, offset] = tester.unit->decompose_location(embed.loc);
+        expect(that % offset == tester.point(position));
+
+        expect(that % embed.file.has_value());
+        expect(that % embed.file_name == filename);
+    };
+
+    auto expect_has_embed =
+        [&](u32 index, llvm::StringRef position, llvm::StringRef filename, bool exists = true) {
+            auto& has_embed = has_embeds[index];
+            auto [_, offset] = tester.unit->decompose_location(has_embed.loc);
+            expect(that % offset == tester.point(position));
+
+            expect(that % has_embed.file.has_value() == exists);
+            expect(that % has_embed.file_name == filename);
+        };
+
     auto expect_con = [&](u32 index, Condition::BranchKind kind, llvm::StringRef pos) {
         auto& condition = conditions[index];
         auto [_, offset] = tester.unit->decompose_location(condition.loc);
@@ -124,6 +147,65 @@ suite<"Directive"> directive = [] {
         expect_has_inl(1, "1", "");
     };
 
+    test("Embed") = [&] {
+        run(R"cpp(
+#[bytes10.bin]
+0123456789
+
+#[bytes5.bin]
+ABCDE
+
+#[main.cpp]
+const char e0 = {
+$(0)#embed "bytes10.bin"
+};
+
+const char e1 = {
+$(1)#embed "bytes10.bin"
+};
+
+const char e2 = {
+$(2)#embed "bytes5.bin"
+};
+
+const char e3 = {
+$(3)#embed "bytes5.bin"
+};
+
+const char e4 = {
+$(4)#embed "non-existed.bin" 
+};
+
+)cpp");
+
+        // e4 will not be processed by clang::PPCallbacks::EmbedDirective(), so there is only 4
+        // embeds.
+        expect(that % embeds.size() == 4);
+        expect_embed(0, "0", "bytes10.bin");
+        expect_embed(1, "1", "bytes10.bin");
+        expect_embed(2, "2", "bytes5.bin");
+        expect_embed(3, "3", "bytes5.bin");
+    };
+
+    test("HasEmbed") = [&] {
+        run(R"cpp(
+#[test.bin]
+
+#[main.cpp]
+#embed "test.bin"
+
+#if __has_embed$(0)("test.bin")
+#endif
+
+#if __has_embed$(1)("non-existed.bin")
+#endif
+)cpp");
+
+        expect(that % has_embeds.size() == 2);
+        expect_has_embed(0, "0", "test.bin");
+        expect_has_embed(1, "1", "non-existed.bin", /*exists=*/false);
+    };
+
     test("Condition") = [&] {
         run(R"cpp(
 #[main.cpp]
@@ -196,7 +278,6 @@ int y = $(6)expr($(7)expr(1));
         expect_pragma(2, Pragma::Kind::EndRegion, "2", "#pragma endregion");
     };
 };
-
 }  // namespace
 
 }  // namespace clice::testing