From 56cf345dba9328e4319c6b7367ac72a8f4339e74 Mon Sep 17 00:00:00 2001
From: Chris Santora <santorac@amazon.com>
Date: Tue, 15 Jun 2021 00:11:43 -0700
Subject: [PATCH] Fixed an issue where vulkan shaders didn't work in
 AtomSampleViewer because the correct API data wasn't being selected.

---
 .../Atom/RPI.Reflect/Shader/ShaderAsset.h     |  5 ++++
 .../Source/RPI.Reflect/Shader/ShaderAsset.cpp | 23 ++++++++++++++++---
 .../RPI.Reflect/Shader/ShaderAssetCreator.cpp |  2 +-
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Shader/ShaderAsset.h b/Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Shader/ShaderAsset.h
index 55d10717a7..e2926c2eab 100644
--- a/Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Shader/ShaderAsset.h
+++ b/Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Shader/ShaderAsset.h
@@ -172,6 +172,11 @@ namespace AZ
 
             bool PostLoadInit() override;
             void SetReady();
+
+            //! SelectShaderApiData() must be called before most other ShaderAsset functions.
+            bool SelectShaderApiData();
+
+            //! Returns the active ShaderApiDataContainer which was selected in SelectShaderApiData().
             ShaderApiDataContainer& GetCurrentShaderApiData();
             const ShaderApiDataContainer& GetCurrentShaderApiData() const;
 
diff --git a/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAsset.cpp b/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAsset.cpp
index e08cf425d7..b4c041b763 100644
--- a/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAsset.cpp
+++ b/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAsset.cpp
@@ -332,8 +332,8 @@ namespace AZ
             // We may only endup here when running in a Builder context.
             return m_perAPIShaderData[0];
         }
-
-        bool ShaderAsset::PostLoadInit()
+        
+        bool ShaderAsset::SelectShaderApiData()
         {
             // Use the current RHI that is active to select which shader data to use.
             // We don't assert if the Factory is not available because this method could be called during build time,
@@ -371,6 +371,11 @@ namespace AZ
                 }
             }
 
+            return true;
+        }
+
+        bool ShaderAsset::PostLoadInit()
+        {
             // Once the ShaderAsset is loaded, it is necessary to listen for changes in the Root Variant Asset.
             Data::AssetBus::Handler::BusConnect(GetRootVariant().GetId());
             ShaderVariantFinderNotificationBus::Handler::BusConnect(GetId());
@@ -444,7 +449,19 @@ namespace AZ
         {
             if (Base::LoadAssetData(asset, stream, assetLoadFilterCB) == Data::AssetHandler::LoadResult::LoadComplete)
             {
-                asset.GetAs<ShaderAsset>()->AssetInitBus::Handler::BusConnect();
+                ShaderAsset* shaderAsset = asset.GetAs<ShaderAsset>();
+
+                // The shader API selection must occur immediately ofter loading, on the same thread, rather than
+                // deferring to AssetInitBus::PostLoadInit. Many functions in the ShaderAsset class are invalid
+                // until after SelectShaderApiData() is called and some client code may need to access data in
+                // the ShaderAsset before then.
+                if (!shaderAsset->SelectShaderApiData())
+                {
+                    return Data::AssetHandler::LoadResult::Error;
+                }
+
+                shaderAsset->AssetInitBus::Handler::BusConnect();
+
                 return Data::AssetHandler::LoadResult::LoadComplete;
             }
 
diff --git a/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAssetCreator.cpp b/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAssetCreator.cpp
index 3f698854df..ce912c2a10 100644
--- a/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAssetCreator.cpp
+++ b/Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderAssetCreator.cpp
@@ -218,7 +218,7 @@ namespace AZ
                 return false;
             }
             
-            if (!m_asset->PostLoadInit())
+            if (!m_asset->SelectShaderApiData())
             {
                 ReportError("Failed to finalize the ShaderAsset.");
                 return false;