[HLSL] Add utility to convert environment to stage

We had a bunch of places in the code where we were translating triple
environment enum cases to shader stage enum cases. The order of these
enums needs to be kept in sync for the translation to be simple, but we
were not properly handling out-of-bounds cases.

In normal compilation out-of-bounds cases shouldn't be possible because
the driver errors if you don't have a valid shader environment set, but
in clang tooling that error doesn't get treated as fatal and parsing
continues. This can result in crashes in clang tooling for out-of-range
shader stages.

To address this, this patch adds a constexpr method to handle the
conversion which handles out-of-range values by converting them to
`Invalid`.

Since this new method is a constexpr, the tests for this are a group of
static_asserts in the implementation file that verifies the correct
conversion for each valid enum case and validates that other cases are
converted to `Invalid`.

Reviewed By: bogner

Differential Revision: https://reviews.llvm.org/D135595

1 files changed, 3 insertions, 2 deletions

diff --git a/clang/lib/Frontend/InitPreprocessor.cpp b/clang/lib/Frontend/InitPreprocessor.cpp
index e9bfab9..2273fb1 100644
--- a/clang/lib/Frontend/InitPreprocessor.cpp
+++ b/clang/lib/Frontend/InitPreprocessor.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Basic/FileManager.h"
+#include "clang/Basic/HLSLRuntime.h"
 #include "clang/Basic/MacroBuilder.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/SyncScope.h"
@@ -402,8 +403,8 @@ static void InitializeStandardPredefinedMacros(const TargetInfo &TI,
     Builder.defineMacro("__SHADER_STAGE_LIBRARY",
                         Twine((uint32_t)ShaderStage::Library));
     // The current shader stage itself
-    uint32_t StageInteger = (uint32_t)TI.getTriple().getEnvironment() -
-                            (uint32_t)llvm::Triple::Pixel;
+    uint32_t StageInteger = static_cast<uint32_t>(
+        hlsl::getStageFromEnvironment(TI.getTriple().getEnvironment()));
 
     Builder.defineMacro("__SHADER_TARGET_STAGE", Twine(StageInteger));
     // Add target versions