refactor(document links): use Lexer for unified directive argument scanning

src/compile/directive.cpp

@@ -94,7 +94,7 @@ class DirectiveCollector : public clang::PPCallbacks {
                             const clang::Token& include_tok,
                             llvm::StringRef,
                             bool,
-                            clang::CharSourceRange filename_range,
+                            clang::CharSourceRange,
                             clang::OptionalFileEntryRef,
                             llvm::StringRef,
                             llvm::StringRef,
@@ -108,7 +108,6 @@ class DirectiveCollector : public clang::PPCallbacks {
         unit->directives[prev_fid].includes.emplace_back(Include{
             .fid = {},
             .location = include_tok.getLocation(),
-            .filename_range = filename_range.getAsRange(),
         });
     }
 

src/compile/directive.h

@@ -20,11 +20,8 @@ struct Include {
     /// The file id of included file.
     clang::FileID fid;
 
-    /// Location of the `include`.
+    /// Location of the `include` keyword.
     clang::SourceLocation location;
-
-    /// The range of filename(includes `""` or `<>`).
-    clang::SourceRange filename_range;
 };
 
 /// Information about `__has_include` directive.

src/feature/document_links.cpp

@@ -1,14 +1,12 @@
-#include <algorithm>
 #include <cstdint>
 #include <string>
 #include <vector>
 
 #include "feature/feature.h"
+#include "syntax/lexer.h"
 
 namespace clice::feature {
 
-namespace {}  // namespace
-
 auto document_links(CompilationUnitRef unit, PositionEncoding encoding)
     -> std::vector<protocol::DocumentLink> {
     std::vector<protocol::DocumentLink> links;
@@ -22,59 +20,41 @@ auto document_links(CompilationUnitRef unit, PositionEncoding encoding)
     auto content = unit.interested_content();
     PositionMapper converter(content, encoding);
     auto& directives = directives_it->second;
+    auto* lang_opts = &unit.lang_options();
 
-    // Scan forward from offset to find a quoted/angled filename range.
-    auto find_filename_range = [&](std::uint32_t offset) -> std::optional<LocalSourceRange> {
-        auto tail = content.substr(offset);
-        auto quote_pos = tail.find_first_of("<\"");
-        if(quote_pos == llvm::StringRef::npos) {
-            return std::nullopt;
-        }
-        char open = tail[quote_pos];
-        char close = open == '<' ? '>' : '"';
-        auto close_pos = tail.find(close, quote_pos + 1);
-        if(close_pos == llvm::StringRef::npos) {
-            return std::nullopt;
-        }
-        return LocalSourceRange(offset + static_cast<std::uint32_t>(quote_pos),
-                                offset + static_cast<std::uint32_t>(close_pos + 1));
-    };
-
-    auto add_link_by_location = [&](clang::SourceLocation loc, llvm::StringRef target) {
+    auto add_link = [&](clang::SourceLocation loc, llvm::StringRef target) {
         auto [fid, offset] = unit.decompose_location(loc);
-        if(fid != interested || offset >= content.size()) {
+        if(fid != interested || offset >= content.size())
             return;
-        }
-        auto range = find_filename_range(offset);
-        if(!range) {
+        auto range = find_directive_argument(content, offset, lang_opts);
+        if(!range)
             return;
-        }
         protocol::DocumentLink link{.range = to_range(converter, *range)};
         link.target = target.str();
         links.push_back(std::move(link));
     };
 
     for(const auto& include: directives.includes) {
         if(include.fid.isValid()) {
-            add_link_by_location(include.location, unit.file_path(include.fid));
+            add_link(include.location, unit.file_path(include.fid));
         }
     }
 
     for(const auto& has_include: directives.has_includes) {
         if(has_include.fid.isValid()) {
-            add_link_by_location(has_include.location, unit.file_path(has_include.fid));
+            add_link(has_include.location, unit.file_path(has_include.fid));
         }
     }
 
     for(const auto& embed: directives.embeds) {
         if(embed.file) {
-            add_link_by_location(embed.loc, embed.file->getName());
+            add_link(embed.loc, embed.file->getName());
         }
     }
 
     for(const auto& has_embed: directives.has_embeds) {
         if(has_embed.file) {
-            add_link_by_location(has_embed.loc, has_embed.file->getName());
+            add_link(has_embed.loc, has_embed.file->getName());
         }
     }
 

src/syntax/lexer.cpp

@@ -53,7 +53,8 @@ void Lexer::lex(Token& token) {
         }
     } else if(parse_pp_keyword) {
         parse_pp_keyword = false;
-        parse_header_name = token.text(content) == "include";
+        auto kw = token.text(content);
+        parse_header_name = kw == "include" || kw == "include_next" || kw == "embed";
     }
 }
 
@@ -105,4 +106,60 @@ Token Lexer::advance_until(TokenKind kind) {
     }
 }
 
+static bool is_directive_keyword(llvm::StringRef word) {
+    return word == "include" || word == "include_next" || word == "import" || word == "embed" ||
+           word == "__has_include" || word == "__has_include_next" || word == "__has_embed";
+}
+
+std::optional<LocalSourceRange> find_directive_argument(llvm::StringRef content,
+                                                        std::uint32_t offset,
+                                                        const clang::LangOptions* lang_opts) {
+    std::uint32_t line_start = 0;
+    if(auto nl = content.rfind('\n', offset); nl != llvm::StringRef::npos)
+        line_start = static_cast<std::uint32_t>(nl + 1);
+
+    auto line = content.substr(line_start);
+    Lexer lexer(line, true, lang_opts);
+    bool after_has_keyword = false;
+    bool ready = false;
+
+    while(true) {
+        auto tok = lexer.advance();
+        if(tok.is_eof() || tok.is_eod())
+            break;
+
+        auto abs_begin = line_start + tok.range.begin;
+        auto abs_end = line_start + tok.range.end;
+
+        if(tok.is_identifier()) {
+            auto text = tok.text(line);
+            if(text == "__has_include" || text == "__has_include_next" || text == "__has_embed") {
+                after_has_keyword = true;
+                continue;
+            }
+            if(text == "include" || text == "include_next" || text == "embed") {
+                ready = true;
+                continue;
+            }
+        }
+
+        if(tok.kind == clang::tok::l_paren && after_has_keyword) {
+            after_has_keyword = false;
+            ready = true;
+            lexer.set_header_name_mode();
+            continue;
+        }
+
+        if(abs_begin < offset || !ready)
+            continue;
+
+        if(tok.is_header_name() || tok.kind == clang::tok::string_literal)
+            return LocalSourceRange(abs_begin, abs_end);
+
+        if(tok.is_identifier())
+            return LocalSourceRange(abs_begin, abs_end);
+    }
+    return std::nullopt;
+}
+
 }  // namespace clice

src/syntax/lexer.h

@@ -51,6 +51,15 @@ class Lexer {
 
     Token advance_until(TokenKind kind);
 
+    /// Force the lexer into header-name mode so the next token is lexed
+    /// via LexIncludeFilename (correctly handling both "..." and <...>).
+    /// Use this before lexing filename arguments in contexts like
+    /// __has_include() or __has_embed() where the lexer cannot detect
+    /// the mode automatically.
+    void set_header_name_mode() {
+        parse_header_name = true;
+    }
+
 private:
     bool ignore_end_of_directive = true;
     bool parse_pp_keyword = false;
@@ -64,4 +73,13 @@ class Lexer {
     std::unique_ptr<clang::Lexer> lexer;
 };
 
+/// Find the range of the filename argument in a preprocessor directive line.
+/// `content` is the full source text, `offset` points at or before the directive keyword.
+/// Returns the range of the first filename-like token (header name, string literal,
+/// or macro identifier) found on the same line, or nullopt if none.
+std::optional<LocalSourceRange>
+    find_directive_argument(llvm::StringRef content,
+                            std::uint32_t offset,
+                            const clang::LangOptions* lang_opts = nullptr);
+
 }  // namespace clice

tests/unit/feature/document_link_tests.cpp

@@ -89,6 +89,19 @@ TEST_CASE(HasInclude) {
     EXPECT_LINK(1, "1", TestVFS::path("test.h"));
 }
 
+TEST_CASE(MacroInclude) {
+    run(R"cpp(
+#[test.h]
+
+#[main.cpp]
+#define HEADER "test.h"
+#include @0[HEADER$]
+)cpp");
+
+    ASSERT_EQ(links.size(), 1U);
+    EXPECT_LINK(0, "0", TestVFS::path("test.h"));
+}
+
 TEST_CASE(Embed) {
     run(R"cpp(
 #[bytes.bin]

tests/unit/syntax/lexer_tests.cpp

@@ -83,5 +83,87 @@ int x = 1;
 }
 
 };  // TEST_SUITE(SourceText)
+
+TEST_SUITE(DirectiveArgument) {
+
+void EXPECT_RANGE(llvm::StringRef content, std::uint32_t offset, llvm::StringRef expected) {
+    auto result = find_directive_argument(content, offset);
+    ASSERT_TRUE(result.has_value());
+    ASSERT_EQ(content.substr(result->begin, result->length()), expected);
+}
+
+void EXPECT_NONE(llvm::StringRef content, std::uint32_t offset) {
+    auto result = find_directive_argument(content, offset);
+    ASSERT_FALSE(result.has_value());
+}
+
+TEST_CASE(IncludeQuoted) {
+    llvm::StringRef src = R"(#include "foo.h")";
+    EXPECT_RANGE(src, 0, R"("foo.h")");
+}
+
+TEST_CASE(IncludeAngled) {
+    llvm::StringRef src = "#include <iostream>";
+    EXPECT_RANGE(src, 0, "<iostream>");
+}
+
+TEST_CASE(IncludeMacro) {
+    llvm::StringRef src = "#include HEADER";
+    EXPECT_RANGE(src, 0, "HEADER");
+}
+
+TEST_CASE(HasIncludeQuoted) {
+    llvm::StringRef src = R"(#if __has_include("foo.h"))";
+    // offset at __has_include
+    auto pos = src.find("__has_include");
+    EXPECT_RANGE(src, static_cast<std::uint32_t>(pos), R"("foo.h")");
+}
+
+TEST_CASE(HasIncludeAngled) {
+    llvm::StringRef src = "#if __has_include(<vector>)";
+    auto pos = src.find("__has_include");
+    EXPECT_RANGE(src, static_cast<std::uint32_t>(pos), "<vector>");
+}
+
+TEST_CASE(EmbedQuoted) {
+    llvm::StringRef src = R"(#embed "data.bin")";
+    EXPECT_RANGE(src, 0, R"("data.bin")");
+}
+
+TEST_CASE(HasEmbedQuoted) {
+    llvm::StringRef src = R"(#if __has_embed("data.bin"))";
+    auto pos = src.find("__has_embed");
+    EXPECT_RANGE(src, static_cast<std::uint32_t>(pos), R"("data.bin")");
+}
+
+TEST_CASE(MultilineOffset) {
+    llvm::StringRef src = "#include \"a.h\"\n#include \"b.h\"";
+    // offset pointing into the second line
+    auto pos = src.find("#include \"b.h\"");
+    EXPECT_RANGE(src, static_cast<std::uint32_t>(pos), R"("b.h")");
+}
+
+TEST_CASE(EmptyDirective) {
+    llvm::StringRef src = "#include \n";
+    EXPECT_NONE(src, 0);
+}
+
+TEST_CASE(HasIncludeFromLineStart) {
+    llvm::StringRef src = "#if __has_include(<vector>)";
+    EXPECT_RANGE(src, 0, "<vector>");
+}
+
+TEST_CASE(HasEmbedFromLineStart) {
+    llvm::StringRef src = R"(#if __has_embed("data.bin"))";
+    EXPECT_RANGE(src, 0, R"("data.bin")");
+}
+
+TEST_CASE(IncludeNext) {
+    llvm::StringRef src = "#include_next <stdlib.h>";
+    EXPECT_RANGE(src, 0, "<stdlib.h>");
+}
+
+};  // TEST_SUITE(DirectiveArgument)
+
 }  // namespace
 }  // namespace clice::testing