Use inference of shader stage from #pragma

deprilula28 · deprilula28 · commit 01334a402320 · 2022-12-20T10:06:35.000-03:00
diff --git a/include/nbl/asset/utils/CHLSLCompiler.h b/include/nbl/asset/utils/CHLSLCompiler.h
@@ -42,7 +42,7 @@ class NBL_API2 CHLSLCompiler final : public IShaderCompiler
 		//	return "";
 		//}
 
-		std::string preprocessShader(std::string&& code, IShader::E_SHADER_STAGE stage, const SPreprocessorOptions& preprocessOptions) const override;
+		std::string preprocessShader(std::string&& code, IShader::E_SHADER_STAGE& stage, const SPreprocessorOptions& preprocessOptions) const override;
 
 	protected:
 
diff --git a/include/nbl/asset/utils/IShaderCompiler.h b/include/nbl/asset/utils/IShaderCompiler.h
@@ -178,7 +178,7 @@ class NBL_API2 IShaderCompiler : public core::IReferenceCounted
 
 		@returns Shader containing logically same High Level code as input but with #include directives resolved.
 		*/
-		virtual std::string preprocessShader(std::string&& code, IShader::E_SHADER_STAGE stage, const SPreprocessorOptions& preprocessOptions) const;
+		virtual std::string preprocessShader(std::string&& code, IShader::E_SHADER_STAGE& stage, const SPreprocessorOptions& preprocessOptions) const;
 
 		std::string preprocessShader(system::IFile* sourcefile, IShader::E_SHADER_STAGE stage, const SPreprocessorOptions& preprocessOptions) const;
 		
diff --git a/src/nbl/asset/interchange/CHLSLLoader.cpp b/src/nbl/asset/interchange/CHLSLLoader.cpp
@@ -28,22 +28,22 @@ SAssetBundle CHLSLLoader::loadAsset(system::IFile* _file, const IAssetLoader::SA
 	const auto filename = _file->getFileName();
 	std::filesystem::path extension = filename.extension();
 
-	//core::unordered_map<std::string,IShader::E_SHADER_STAGE> typeFromExt =	{	
-	//																						{".vert",IShader::ESS_VERTEX},
-	//																						{".tesc",IShader::ESS_TESSELLATION_CONTROL},
-	//																						{".tese",IShader::ESS_TESSELLATION_EVALUATION},
-	//																						{".geom",IShader::ESS_GEOMETRY},
-	//																						{".frag",IShader::ESS_FRAGMENT},
-	//																						{".comp",IShader::ESS_COMPUTE}
-	//																					};
-	//auto found = typeFromExt.find(extension.string());
-	//if (found == typeFromExt.end())
-	//{
-	//	_NBL_ALIGNED_FREE(source);
-	//	return {};
-	//}
-
-	auto shader = core::make_smart_refctd_ptr<ICPUShader>(reinterpret_cast<char*>(source), IShader::ESS_COMPUTE, IShader::E_CONTENT_TYPE::ECT_HLSL, filename.string());
+	core::unordered_map<std::string,IShader::E_SHADER_STAGE> typeFromExt =	{	
+		{".vert.hlsl",IShader::ESS_VERTEX},
+		{".tesc.hlsl",IShader::ESS_TESSELLATION_CONTROL},
+		{".tese.hlsl",IShader::ESS_TESSELLATION_EVALUATION},
+		{".geom.hlsl",IShader::ESS_GEOMETRY},
+		{".frag.hlsl",IShader::ESS_FRAGMENT},
+		{".comp.hlsl",IShader::ESS_COMPUTE}
+	};
+	auto found = typeFromExt.find(extension.string());
+	auto shaderStage = IShader::ESS_UNKNOWN;
+	if (found != typeFromExt.end())
+	{
+		shaderStage = found->second;
+	}
+
+	auto shader = core::make_smart_refctd_ptr<ICPUShader>(reinterpret_cast<char*>(source), shaderStage, IShader::E_CONTENT_TYPE::ECT_HLSL, filename.string());
 	_NBL_ALIGNED_FREE(source);
 
 	return SAssetBundle(nullptr,{ core::make_smart_refctd_ptr<ICPUSpecializedShader>(std::move(shader),ISpecializedShader::SInfo({},nullptr,"main")) });
diff --git a/src/nbl/asset/utils/CCompilerSet.cpp b/src/nbl/asset/utils/CCompilerSet.cpp
@@ -44,15 +44,17 @@ core::smart_refctd_ptr<ICPUShader> CCompilerSet::preprocessShader(const ICPUShad
 		case IShader::E_CONTENT_TYPE::ECT_HLSL:
 		{
 			const char* code = reinterpret_cast<const char*>(shader->getContent()->getPointer());
-			auto resolvedCode = m_HLSLCompiler->preprocessShader(code, shader->getStage(), preprocessOptions);
-			return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), shader->getStage(), IShader::E_CONTENT_TYPE::ECT_HLSL, std::string(shader->getFilepathHint()));
+			auto stage = shader->getStage();
+			auto resolvedCode = m_HLSLCompiler->preprocessShader(code, stage, preprocessOptions);
+			return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), stage, IShader::E_CONTENT_TYPE::ECT_HLSL, std::string(shader->getFilepathHint()));
 		}
 		break;
 		case IShader::E_CONTENT_TYPE::ECT_GLSL:
 		{
 			const char* code = reinterpret_cast<const char*>(shader->getContent()->getPointer());
-			auto resolvedCode = m_GLSLCompiler->preprocessShader(code, shader->getStage(), preprocessOptions);
-			return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), shader->getStage(), IShader::E_CONTENT_TYPE::ECT_GLSL, std::string(shader->getFilepathHint()));
+			auto stage = shader->getStage();
+			auto resolvedCode = m_GLSLCompiler->preprocessShader(code, stage, preprocessOptions);
+			return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), stage, IShader::E_CONTENT_TYPE::ECT_GLSL, std::string(shader->getFilepathHint()));
 		}
 		break;
 		default:
diff --git a/src/nbl/asset/utils/CHLSLCompiler.cpp b/src/nbl/asset/utils/CHLSLCompiler.cpp
@@ -74,11 +74,6 @@ static tcpp::IInputStream* getInputStreamInclude(
     if (!res_str.size()) {
         return nullptr;
     }
-    else {
-        //employ encloseWithinExtraInclGuards() in order to prevent infinite loop of (not necesarilly direct) self-inclusions while other # directives (incl guards among them) are disabled
-        IShaderCompiler::disableAllDirectivesExceptIncludes(res_str);
-        res_str = IShaderCompiler::encloseWithinExtraInclGuards(std::move(res_str), _maxInclCnt, name.string().c_str());
-    }
 
     return new tcpp::StringInputStream(std::move(res_str));
 }
@@ -143,7 +138,7 @@ CHLSLCompiler::DxcCompilationResult CHLSLCompiler::dxcCompile(std::string& sourc
     return result;
 }
 
-std::string CHLSLCompiler::preprocessShader(std::string&& code, IShader::E_SHADER_STAGE stage, const SPreprocessorOptions& preprocessOptions) const
+std::string CHLSLCompiler::preprocessShader(std::string&& code, IShader::E_SHADER_STAGE& stage, const SPreprocessorOptions& preprocessOptions) const
 {
     if (preprocessOptions.extraDefines.size())
     {
@@ -166,6 +161,52 @@ std::string CHLSLCompiler::preprocessShader(std::string&& code, IShader::E_SHADE
             }
         );
 
+        auto pragmaShaderStageCallback = [&](IShader::E_SHADER_STAGE _stage) {
+            return [&](tcpp::Preprocessor& preprocessor, tcpp::Lexer& lexer, const std::string& text) {
+                stage = _stage;
+                return std::string("");
+            };
+        };
+
+        proc.AddCustomDirectiveHandler(std::string("pragma shader_stage"), [&](tcpp::Preprocessor& preprocessor, tcpp::Lexer& lexer, const std::string& text) {
+            if (!lexer.HasNextToken()) return std::string("#error Malformed shader_stage pragma");
+            auto token = lexer.GetNextToken();
+            if (token.mType != tcpp::E_TOKEN_TYPE::OPEN_BRACKET) return std::string("#error Malformed shader_stage pragma");
+
+            if (!lexer.HasNextToken()) return std::string("#error Malformed shader_stage pragma");
+            token = lexer.GetNextToken();
+            if (token.mType != tcpp::E_TOKEN_TYPE::IDENTIFIER) return std::string("#error Malformed shader_stage pragma");
+
+            auto& shaderStageIdentifier = token.mRawView;
+            core::unordered_map<std::string, IShader::E_SHADER_STAGE> stageFromIdent = {
+                { "vertex", IShader::ESS_VERTEX },
+                { "fragment", IShader::ESS_FRAGMENT },
+                { "tesscontrol", IShader::ESS_TESSELLATION_CONTROL },
+                { "tesseval", IShader::ESS_TESSELLATION_EVALUATION },
+                { "geometry", IShader::ESS_GEOMETRY },
+                { "compute", IShader::ESS_COMPUTE }
+            };
+
+            auto found = stageFromIdent.find(shaderStageIdentifier);
+            if (found == stageFromIdent.end())
+            {
+                return std::string("#error Malformed shader_stage pragma, unknown stage");
+            }
+            stage = found->second;
+
+            if (!lexer.HasNextToken()) return std::string("#error Malformed shader_stage pragma");
+            token = lexer.GetNextToken();
+            if (token.mType != tcpp::E_TOKEN_TYPE::CLOSE_BRACKET) return std::string("#error Malformed shader_stage pragma");
+
+            while (lexer.HasNextToken()) {
+                auto token = lexer.GetNextToken();
+                if (token.mType == tcpp::E_TOKEN_TYPE::NEWLINE) break;
+                if (token.mType != tcpp::E_TOKEN_TYPE::SPACE) return std::string("#error Malformed shader_stage pragma");
+            }
+
+            return std::string("");
+        });
+
         return proc.Process();
     }
     else
@@ -184,13 +225,14 @@ core::smart_refctd_ptr<ICPUShader> CHLSLCompiler::compileToSPIRV(const char* cod
         return nullptr;
     }
 
-    auto newCode = preprocessShader(code, hlslOptions.stage, hlslOptions.preprocessorOptions);
+    auto stage = hlslOptions.stage;
+    auto newCode = preprocessShader(code, stage, hlslOptions.preprocessorOptions);
 
     // Suffix is the shader model version
     std::wstring targetProfile(L"XX_6_2");
 
     // Set profile two letter prefix based on stage
-    switch (options.stage) {
+    switch (stage) {
     case asset::IShader::ESS_VERTEX:
         targetProfile.replace(0, 2, L"vs");
         break;
@@ -216,7 +258,7 @@ core::smart_refctd_ptr<ICPUShader> CHLSLCompiler::compileToSPIRV(const char* cod
         targetProfile.replace(0, 2, L"ms");
         break;
     default:
-        hlslOptions.preprocessorOptions.logger.log("invalid shader stage %i", system::ILogger::ELL_ERROR, options.stage);
+        hlslOptions.preprocessorOptions.logger.log("invalid shader stage %i", system::ILogger::ELL_ERROR, stage);
         return nullptr;
     };
 
@@ -246,7 +288,7 @@ core::smart_refctd_ptr<ICPUShader> CHLSLCompiler::compileToSPIRV(const char* cod
 
     compileResult.release();
 
-    return core::make_smart_refctd_ptr<asset::ICPUShader>(std::move(outSpirv), hlslOptions.stage, IShader::E_CONTENT_TYPE::ECT_SPIRV, hlslOptions.preprocessorOptions.sourceIdentifier.data());
+    return core::make_smart_refctd_ptr<asset::ICPUShader>(std::move(outSpirv), stage, IShader::E_CONTENT_TYPE::ECT_SPIRV, hlslOptions.preprocessorOptions.sourceIdentifier.data());
 }
 
 void CHLSLCompiler::insertIntoStart(std::string& code, std::ostringstream&& ins) const
diff --git a/src/nbl/asset/utils/IShaderCompiler.cpp b/src/nbl/asset/utils/IShaderCompiler.cpp
@@ -101,7 +101,7 @@ namespace nbl::asset::impl
 
 std::string IShaderCompiler::preprocessShader(
     std::string&& code,
-    IShader::E_SHADER_STAGE stage,
+    IShader::E_SHADER_STAGE& stage,
     const SPreprocessorOptions& preprocessOptions) const
 {
     if (preprocessOptions.extraDefines.size())
diff --git a/src/nbl/video/CVulkanLogicalDevice.h b/src/nbl/video/CVulkanLogicalDevice.h
@@ -570,7 +570,7 @@ class CVulkanLogicalDevice final : public ILogicalDevice
         if (m_devf.vk.vkCreateShaderModule(m_vkdev, &vk_createInfo, nullptr, &vk_shaderModule) == VK_SUCCESS)
         {
             return core::make_smart_refctd_ptr<video::CVulkanShader>(
-                core::smart_refctd_ptr<CVulkanLogicalDevice>(this), cpushader->getStage(), std::string(cpushader->getFilepathHint()), vk_shaderModule);
+                core::smart_refctd_ptr<CVulkanLogicalDevice>(this), spirvShader->getStage(), std::string(cpushader->getFilepathHint()), vk_shaderModule);
         }
         else
         {

Original file line number	Diff line number	Diff line change
`@@ -44,15 +44,17 @@ core::smart_refctd_ptr<ICPUShader> CCompilerSet::preprocessShader(const ICPUShad`
`44`	`44`	`case IShader::E_CONTENT_TYPE::ECT_HLSL:`
`45`	`45`	`{`
`46`	`46`	`const char* code = reinterpret_cast<const char*>(shader->getContent()->getPointer());`
`47`		`- auto resolvedCode = m_HLSLCompiler->preprocessShader(code, shader->getStage(), preprocessOptions);`
`48`		`- return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), shader->getStage(), IShader::E_CONTENT_TYPE::ECT_HLSL, std::string(shader->getFilepathHint()));`
	`47`	`+ auto stage = shader->getStage();`
	`48`	`+ auto resolvedCode = m_HLSLCompiler->preprocessShader(code, stage, preprocessOptions);`
	`49`	`+ return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), stage, IShader::E_CONTENT_TYPE::ECT_HLSL, std::string(shader->getFilepathHint()));`
`49`	`50`	`}`
`50`	`51`	`break;`
`51`	`52`	`case IShader::E_CONTENT_TYPE::ECT_GLSL:`
`52`	`53`	`{`
`53`	`54`	`const char* code = reinterpret_cast<const char*>(shader->getContent()->getPointer());`
`54`		`- auto resolvedCode = m_GLSLCompiler->preprocessShader(code, shader->getStage(), preprocessOptions);`
`55`		`- return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), shader->getStage(), IShader::E_CONTENT_TYPE::ECT_GLSL, std::string(shader->getFilepathHint()));`
	`55`	`+ auto stage = shader->getStage();`
	`56`	`+ auto resolvedCode = m_GLSLCompiler->preprocessShader(code, stage, preprocessOptions);`
	`57`	`+ return core::make_smart_refctd_ptr<ICPUShader>(resolvedCode.c_str(), stage, IShader::E_CONTENT_TYPE::ECT_GLSL, std::string(shader->getFilepathHint()));`
`56`	`58`	`}`
`57`	`59`	`break;`
`58`	`60`	`default:`
Original file line number	Diff line number	Diff line change
`@@ -101,7 +101,7 @@ namespace nbl::asset::impl`
`101`	`101`
`102`	`102`	`std::string IShaderCompiler::preprocessShader(`
`103`	`103`	`std::string&& code,`
`104`		`- IShader::E_SHADER_STAGE stage,`
	`104`	`+ IShader::E_SHADER_STAGE& stage,`
`105`	`105`	`const SPreprocessorOptions& preprocessOptions) const`
`106`	`106`	`{`
`107`	`107`	`if (preprocessOptions.extraDefines.size())`
Original file line number	Diff line number	Diff line change
`@@ -570,7 +570,7 @@ class CVulkanLogicalDevice final : public ILogicalDevice`
`570`	`570`	`if (m_devf.vk.vkCreateShaderModule(m_vkdev, &vk_createInfo, nullptr, &vk_shaderModule) == VK_SUCCESS)`
`571`	`571`	`{`
`572`	`572`	`return core::make_smart_refctd_ptr<video::CVulkanShader>(`
`573`		`- core::smart_refctd_ptr<CVulkanLogicalDevice>(this), cpushader->getStage(), std::string(cpushader->getFilepathHint()), vk_shaderModule);`
	`573`	`+ core::smart_refctd_ptr<CVulkanLogicalDevice>(this), spirvShader->getStage(), std::string(cpushader->getFilepathHint()), vk_shaderModule);`
`574`	`574`	`}`
`575`	`575`	`else`
`576`	`576`	`{`