1 files changed, 98 insertions, 3 deletions

diff --git a/clang/lib/Frontend/FrontendActions.cpp b/clang/lib/Frontend/FrontendActions.cpp
index 685a9bbf2cde..7424958d4661 100644
--- a/clang/lib/Frontend/FrontendActions.cpp
+++ b/clang/lib/Frontend/FrontendActions.cpp
@@ -22,6 +22,7 @@
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
+#include "clang/Parse/ParseHLSLRootSignature.h"
 #include "clang/Sema/TemplateInstCallback.h"
 #include "clang/Serialization/ASTReader.h"
 #include "clang/Serialization/ASTWriter.h"
@@ -795,9 +796,10 @@ namespace {
     /// Indicates that the AST file contains particular input file.
     ///
     /// \returns true to continue receiving the next input file, false to stop.
-    bool visitInputFile(StringRef FilenameAsRequested, StringRef Filename,
-                        bool isSystem, bool isOverridden,
-                        bool isExplicitModule) override {
+    bool visitInputFileAsRequested(StringRef FilenameAsRequested,
+                                   StringRef Filename, bool isSystem,
+                                   bool isOverridden,
+                                   bool isExplicitModule) override {
 
       Out.indent(2) << "Input file: " << FilenameAsRequested;
 
@@ -1241,3 +1243,96 @@ void GetDependenciesByModuleNameAction::ExecuteAction() {
   PPCallbacks *CB = PP.getPPCallbacks();
   CB->moduleImport(SourceLocation(), Path, ModResult);
 }
+
+//===----------------------------------------------------------------------===//
+// HLSL Specific Actions
+//===----------------------------------------------------------------------===//
+
+class InjectRootSignatureCallback : public PPCallbacks {
+private:
+  Sema &Actions;
+  StringRef RootSigName;
+  llvm::dxbc::RootSignatureVersion Version;
+
+  std::optional<StringLiteral *> processStringLiteral(ArrayRef<Token> Tokens) {
+    for (Token Tok : Tokens)
+      if (!tok::isStringLiteral(Tok.getKind()))
+        return std::nullopt;
+
+    ExprResult StringResult = Actions.ActOnUnevaluatedStringLiteral(Tokens);
+    if (StringResult.isInvalid())
+      return std::nullopt;
+
+    if (auto Signature = dyn_cast<StringLiteral>(StringResult.get()))
+      return Signature;
+
+    return std::nullopt;
+  }
+
+public:
+  void MacroDefined(const Token &MacroNameTok,
+                    const MacroDirective *MD) override {
+    if (RootSigName != MacroNameTok.getIdentifierInfo()->getName())
+      return;
+
+    const MacroInfo *MI = MD->getMacroInfo();
+    auto Signature = processStringLiteral(MI->tokens());
+    if (!Signature.has_value()) {
+      Actions.getDiagnostics().Report(MI->getDefinitionLoc(),
+                                      diag::err_expected_string_literal)
+          << /*in attributes...*/ 4 << "RootSignature";
+      return;
+    }
+
+    IdentifierInfo *DeclIdent =
+        hlsl::ParseHLSLRootSignature(Actions, Version, *Signature);
+    Actions.HLSL().SetRootSignatureOverride(DeclIdent);
+  }
+
+  InjectRootSignatureCallback(Sema &Actions, StringRef RootSigName,
+                              llvm::dxbc::RootSignatureVersion Version)
+      : PPCallbacks(), Actions(Actions), RootSigName(RootSigName),
+        Version(Version) {}
+};
+
+void HLSLFrontendAction::ExecuteAction() {
+  // Pre-requisites to invoke
+  CompilerInstance &CI = getCompilerInstance();
+  if (!CI.hasASTContext() || !CI.hasPreprocessor())
+    return WrapperFrontendAction::ExecuteAction();
+
+  // InjectRootSignatureCallback requires access to invoke Sema to lookup/
+  // register a root signature declaration. The wrapped action is required to
+  // account for this by only creating a Sema if one doesn't already exist
+  // (like we have done, and, ASTFrontendAction::ExecuteAction)
+  if (!CI.hasSema())
+    CI.createSema(getTranslationUnitKind(),
+                  /*CodeCompleteConsumer=*/nullptr);
+  Sema &S = CI.getSema();
+
+  auto &TargetInfo = CI.getASTContext().getTargetInfo();
+  bool IsRootSignatureTarget =
+      TargetInfo.getTriple().getEnvironment() == llvm::Triple::RootSignature;
+  StringRef HLSLEntry = TargetInfo.getTargetOpts().HLSLEntry;
+
+  // Register HLSL specific callbacks
+  auto LangOpts = CI.getLangOpts();
+  StringRef RootSigName =
+      IsRootSignatureTarget ? HLSLEntry : LangOpts.HLSLRootSigOverride;
+
+  auto MacroCallback = std::make_unique<InjectRootSignatureCallback>(
+      S, RootSigName, LangOpts.HLSLRootSigVer);
+
+  Preprocessor &PP = CI.getPreprocessor();
+  PP.addPPCallbacks(std::move(MacroCallback));
+
+  // If we are targeting a root signature, invoke custom handling
+  if (IsRootSignatureTarget)
+    return hlsl::HandleRootSignatureTarget(S, HLSLEntry);
+  else // otherwise, invoke as normal
+    return WrapperFrontendAction::ExecuteAction();
+}
+
+HLSLFrontendAction::HLSLFrontendAction(
+    std::unique_ptr<FrontendAction> WrappedAction)
+    : WrapperFrontendAction(std::move(WrappedAction)) {}