1 files changed, 138 insertions, 1 deletions

diff --git a/source/slang/slang-preprocessor.cpp b/source/slang/slang-preprocessor.cpp
index fceee27f9..5a54bfcb0 100644
--- a/source/slang/slang-preprocessor.cpp
+++ b/source/slang/slang-preprocessor.cpp
@@ -222,6 +222,9 @@ struct Preprocessor
         /// The module, if any, that the preprocessed result will belong to
     Module*                                 parentModule = nullptr;
 
+        /// The AST builder that should be used when creating AST nodes for `parentModule`
+    ASTBuilder*                             astBuilder = nullptr;
+
     // The unique identities of any paths that have issued `#pragma once` directives to
     // stop them from being included again.
     HashSet<String>                         pragmaOnceUniqueIdentities;
@@ -2447,18 +2450,141 @@ static TokenList ReadAllTokens(
     return tokens;
 }
 
+    /// Try to look up a macro with the given `macroName` and produce its value as a string
+static bool _findMacroValue(
+    Preprocessor*   preprocessor,
+    char const*     macroName,
+    String&         outValue,
+    SourceLoc&      outLoc)
+{
+    auto namePool = preprocessor->linkage->getNamePool();
+    auto macro = LookupMacro(preprocessor, namePool->getName(macroName));
+    if(!macro)
+        return false;
+    if(macro->flavor != PreprocessorMacroFlavor::ObjectLike)
+        return false;
+
+    MacroExpansion* expansion = new MacroExpansion();
+    initializeMacroExpansion(preprocessor, expansion, macro);
+    pushMacroExpansion(preprocessor, expansion);
+
+    String value;
+    for(bool first = true;;first = false)
+    {
+        Token token = ReadToken(preprocessor);
+        if(token.type == TokenType::EndOfFile)
+            break;
+
+        if(!first && (token.flags & TokenFlag::AfterWhitespace))
+            value.append(" ");
+        value.append(token.getContent());
+    }
+
+    outValue = value;
+    outLoc = macro->getLoc();
+    return true;
+}
+
+    /// Validate that a re-defintion of an NVAPI-related macro matches any previous definition
+static void _validateNVAPIMacroMatch(
+    Preprocessor*   preprocessor,
+    char const*     macroName,
+    String const&   existingValue,
+    String const&   newValue,
+    SourceLoc       loc)
+{
+    if( existingValue != newValue )
+    {
+        preprocessor->sink->diagnose(loc, Diagnostics::nvapiMacroMismatch, macroName, existingValue, newValue);
+    }
+}
+
+    /// Collect macro definitions that are relevant to subsequent compilation steps, and store them
+static void _collectDownstreamRelevantMacros(
+    Preprocessor*   preprocessor)
+{
+    // For now, the only case of semantically-relevant macros we need to worrry
+    // about are the NVAPI macros used to establish the register/space to use.
+    //
+    static const char* kNVAPIRegisterMacroName = "NV_SHADER_EXTN_SLOT";
+    static const char* kNVAPISpaceMacroName = "NV_SHADER_EXTN_REGISTER_SPACE";
+
+    // For NVAPI use, the `NV_SHADER_EXTN_SLOT` macro is required to be defined.
+    //
+    String nvapiRegister;
+    SourceLoc nvapiRegisterLoc;
+    if(_findMacroValue(preprocessor, kNVAPIRegisterMacroName, nvapiRegister, nvapiRegisterLoc))
+    {
+        // In contrast, NVAPI can be used without defining `NV_SHADER_EXTN_REGISTER_SPACE`,
+        // which effectively defaults to `space0`.
+        //
+        String nvapiSpace = "space0";
+        SourceLoc nvapiSpaceLoc;
+        _findMacroValue(preprocessor, kNVAPISpaceMacroName, nvapiSpace, nvapiSpaceLoc);
+
+        // We are going to store the values of these macros on the AST-level `ModuleDecl`
+        // so that they will be available to later processing stages.
+        //
+        // Note: An alternative design here would be to either put the data directly
+        // on the `Module`, or to define some kind of side-channel output that the
+        // preprocessor can use to communicate these macro values back (and then
+        // allow another system to create the AST nodes). In practice, all of these
+        // alternatives would actually increase the amount of code/complexity,
+        // so we stick with the simple-but-hacky option of having the
+        // preprocessor create AST nodes directly.
+        //
+        auto module = preprocessor->parentModule;
+        if(!module) return;
+        auto moduleDecl = module->getModuleDecl();
+
+        // We need to make sure that the AST nodes we create will have the right
+        // lifetime (it should match the module we are adding to).
+        //
+        auto astBuilder = preprocessor->astBuilder;
+        if(!astBuilder) return;
+
+        if(auto existingModifier = moduleDecl->findModifier<NVAPISlotModifier>())
+        {
+            // If there is already a modifier attached to the module (perhaps
+            // because of preprocessing a different source file, or because
+            // of settings established via command-line options), then we
+            // need to validate that the values being set in this file
+            // match those already set (or else there is likely to be
+            // some kind of error in the user's code).
+            //
+            _validateNVAPIMacroMatch(preprocessor, kNVAPIRegisterMacroName, existingModifier->registerName, nvapiRegister,  nvapiRegisterLoc);
+            _validateNVAPIMacroMatch(preprocessor, kNVAPISpaceMacroName,    existingModifier->spaceName,    nvapiSpace,     nvapiSpaceLoc);
+        }
+        else
+        {
+            // If there is no existing modifier on the module, then we
+            // take responsibility for adding one, based on the macro
+            // values we saw.
+            //
+            auto modifier = astBuilder->create<NVAPISlotModifier>();
+            modifier->loc = nvapiRegisterLoc;
+            modifier->registerName = nvapiRegister;
+            modifier->spaceName = nvapiSpace;
+
+            addModifier(moduleDecl, modifier);
+        }
+    }
+}
+
 TokenList preprocessSource(
     SourceFile*                 file,
     DiagnosticSink*             sink,
     IncludeSystem*              includeSystem,
     Dictionary<String, String>  defines,
     Linkage*                    linkage,
-    Module*                     parentModule)
+    Module*                     parentModule,
+    ASTBuilder*                 astBuilder)
 {
     Preprocessor preprocessor;
     InitializePreprocessor(&preprocessor, sink);
     preprocessor.linkage = linkage;
     preprocessor.parentModule = parentModule;
+    preprocessor.astBuilder = astBuilder;
 
     preprocessor.includeSystem = includeSystem;
     for (auto p : defines)
@@ -2475,6 +2601,17 @@ TokenList preprocessSource(
 
     TokenList tokens = ReadAllTokens(&preprocessor);
 
+    // We look at the preprocessor state after reading the entire
+    // source file/string, in order to see if any macros have been
+    // set that should be considered semantically relevant for
+    // later stages of compilation.
+    //
+    // Note: Checking the macro environment *after* preprocessing is complete
+    // means that we can treat macros introduced via `-D` options or the API
+    // equivalently to macros introduced via `#define`s in user code.
+    //
+    _collectDownstreamRelevantMacros(&preprocessor);
+
     FinalizePreprocessor(&preprocessor);
 
     // debugging: build the pre-processed source back together