================
@@ -121,6 +181,219 @@ TEST(FeatureModulesTest, BeforeExecute) {
}
}
+TEST(FeatureModulesTest, AfterExecute) {
+ struct AfterExecuteState {
+ bool ReenterPreprocessorInConsumer = false;
+ bool HandledTranslationUnit = false;
+ bool ConsumerSawWholeTranslationUnit = false;
+ bool AfterExecuteCalled = false;
+ bool ReenteredPreprocessor = false;
+ std::vector<std::string> DeclNames;
+ };
+ struct AfterExecuteModule final : public FeatureModule {
+ struct Consumer : public ASTConsumer {
+ Consumer(AfterExecuteState &S, CompilerInstance &CI) : S(S), CI(CI) {}
+
+ void HandleTranslationUnit(ASTContext &Ctx) override {
+ S.HandledTranslationUnit = true;
+ auto Scope = Ctx.getTraversalScope();
+ S.ConsumerSawWholeTranslationUnit =
+ Scope.size() == 1 && Scope.front() == Ctx.getTranslationUnitDecl();
+ Pending = &Ctx;
+ if (S.ReenterPreprocessorInConsumer)
+ reenterAtEOF(CI);
+ }
+
+ void run(CompilerInstance &CI) {
+ if (!Pending)
+ return;
+ SourceLocation MainFileDeclLoc;
+ for (Decl *D : Pending->getTraversalScope()) {
+ if (const auto *ND = llvm::dyn_cast<NamedDecl>(D)) {
+ S.DeclNames.push_back(ND->getNameAsString());
+ MainFileDeclLoc = ND->getLocation();
+ }
+ }
+
+ if (MainFileDeclLoc.isValid())
+ S.ReenteredPreprocessor = reenterPreprocessor(CI, MainFileDeclLoc);
+ }
+
+ private:
+ static void reenterAtEOF(CompilerInstance &CI) {
+ Token End;
+ End.startToken();
+ auto &SM = CI.getSourceManager();
+ End.setLocation(SM.getLocForEndOfFile(SM.getMainFileID()));
+ End.setKind(tok::eof);
+ std::array<Token, 1> Stream{End};
+ auto &PP = CI.getPreprocessor();
+ PP.EnterTokenStream(Stream, /*DisableMacroExpansion=*/false,
+ /*IsReinject=*/false);
+ PP.Lex(End);
+ }
+
+ static bool reenterPreprocessor(CompilerInstance &CI,
+ SourceLocation MainFileDeclLoc) {
+ Token Reinjected;
+ if (Lexer::getRawToken(MainFileDeclLoc, Reinjected,
+ CI.getSourceManager(), CI.getLangOpts()))
+ return false;
+ auto &PP = CI.getPreprocessor();
+ PP.LookUpIdentifierInfo(Reinjected);
+ Token End;
+ End.startToken();
+ End.setKind(tok::eof);
+ std::array<Token, 2> Stream{Reinjected, End};
+ PP.EnterTokenStream(Stream, /*DisableMacroExpansion=*/false,
+ /*IsReinject=*/false);
+ do {
+ PP.Lex(Reinjected);
+ } while (Reinjected.isNot(tok::eof));
+ return true;
+ }
+
+ AfterExecuteState &S;
+ CompilerInstance &CI;
+ ASTContext *Pending = nullptr;
+ };
+
+ struct Listener : public FeatureModule::ASTListener {
+ Listener(AfterExecuteState &S) : S(S) {}
+
+ void beforeExecute(CompilerInstance &CI) override {
+ std::vector<std::unique_ptr<ASTConsumer>> Consumers;
+ Consumers.push_back(CI.takeASTConsumer());
+ auto Deferred = std::make_unique<Consumer>(S, CI);
+ DeferredConsumer = Deferred.get();
+ Consumers.push_back(std::move(Deferred));
+ CI.setASTConsumer(
+ std::make_unique<MultiplexConsumer>(std::move(Consumers)));
+ }
+
+ void afterExecute(CompilerInstance &CI) override {
+ S.AfterExecuteCalled = true;
+ if (DeferredConsumer)
+ DeferredConsumer->run(CI);
+ }
+
+ private:
+ AfterExecuteState &S;
+ Consumer *DeferredConsumer = nullptr;
+ };
+
+ AfterExecuteModule(AfterExecuteState &S) : S(S) {}
+ std::unique_ptr<ASTListener> astListeners() override {
+ return std::make_unique<Listener>(S);
+ };
+
+ private:
+ AfterExecuteState &S;
+ };
+
+ // HandleTranslationUnit is too early for work that re-enters the
+ // preprocessor: clangd's token collector is still installed, observes the
+ // extra token, and cannot build a valid TokenBuffer afterwards.
----------------
timon-ul wrote:
Thanks! Makes sense.
https://github.com/llvm/llvm-project/pull/221054
_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
