o3de/Gems/Atom/RPI/Code/Source/RPI.Reflect/Material/ShaderCollection.cpp

/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#include <AtomCore/std/containers/vector_set.h>
#include <AzCore/Asset/AssetSerializer.h>
#include <Atom/RPI.Reflect/Material/ShaderCollection.h>
#include <Atom/RHI/RHISystemInterface.h>
#include <Atom/RHI/DrawListTagRegistry.h>

namespace AZ
{
    namespace RPI
    {
        //! This allows ShaderCollection::Item to serialize only a ShaderVariantId rather than the ShaderOptionsGroup object,
        //! but still provide the corresponding ShaderOptionsGroup for use at runtime.
        //! RenderStates will be modified at runtime as well. It will be merged into the RenderStates stored in the corresponding ShaderVariant.
        class ShaderVariantReferenceSerializationEvents
            : public SerializeContext::IEventHandler
        {
            //! Called right before we start reading from the instance pointed by classPtr.
            void OnReadBegin(void* classPtr) override
            {
                ShaderCollection::Item* shaderVariantReference = reinterpret_cast<ShaderCollection::Item*>(classPtr);
                shaderVariantReference->m_shaderVariantId = shaderVariantReference->m_shaderOptionGroup.GetShaderVariantId();
            }

            //! Called right after we finish writing data to the instance pointed at by classPtr.
            void OnWriteEnd(void* classPtr) override
            {
                ShaderCollection::Item* shaderVariantReference = reinterpret_cast<ShaderCollection::Item*>(classPtr);

                if (shaderVariantReference->m_shaderAsset.IsReady())
                {
                    shaderVariantReference->m_shaderOptionGroup = ShaderOptionGroup{
                        shaderVariantReference->m_shaderAsset->GetShaderOptionGroupLayout(),
                        shaderVariantReference->m_shaderVariantId
                    };
                }
                else
                {
                    // No worries, eventually the Material::Init will end up
                    // calling InitializeShaderOptionGroup() and @m_shaderOptionGroup
                    // will get proper data.
                    shaderVariantReference->m_shaderOptionGroup = {};
                    shaderVariantReference->m_shaderAsset.QueueLoad(); // Not necessary to call QueueLoad, but doesn't hurt either.
                }

            }
        };

        void ShaderCollection::Reflect(AZ::ReflectContext* context)
        {
            ShaderCollection::Item::Reflect(context);
            NameReflectionMapForIndex::Reflect(context);

            if (auto* serializeContext = azrtti_cast<AZ::SerializeContext*>(context))
            {
                serializeContext->Class<ShaderCollection>()
                    ->Version(5)
                    ->Field("ShaderItems", &ShaderCollection::m_shaderItems)
                    ->Field("ShaderTagIndexMap", &ShaderCollection::m_shaderTagIndexMap)
                    ;
            }
        }

        void ShaderCollection::Item::Reflect(AZ::ReflectContext* context)
        {
            if (auto* serializeContext = azrtti_cast<AZ::SerializeContext*>(context))
            {
                serializeContext->Enum<ShaderCollection::Item::DrawItemType>()
                    ->Value("Raster", ShaderCollection::Item::DrawItemType::Raster)
                    ->Value("Dispatch", ShaderCollection::Item::DrawItemType::Dispatch)
                    ->Value("Deferred", ShaderCollection::Item::DrawItemType::Deferred)
                    ->Value("RayTracing", ShaderCollection::Item::DrawItemType::RayTracing)
                    ->Value("Custom", ShaderCollection::Item::DrawItemType::Custom);

                serializeContext->Class<ShaderCollection::Item>()
                    ->Version(7)
                    ->EventHandler<ShaderVariantReferenceSerializationEvents>()
                    ->Field("ShaderAsset", &ShaderCollection::Item::m_shaderAsset)
                    ->Field("ShaderVariantId", &ShaderCollection::Item::m_shaderVariantId)
                    ->Field("Enabled", &ShaderCollection::Item::m_enabled)
                    ->Field("OwnedShaderOptionIndices", &ShaderCollection::Item::m_ownedShaderOptionIndices)
                    ->Field("ShaderTag", &ShaderCollection::Item::m_shaderTag)
                    ->Field("DrawItemType", &ShaderCollection::Item::m_drawItemType);
            }

            if (BehaviorContext* behaviorContext = azrtti_cast<BehaviorContext*>(context))
            {
                behaviorContext->Class<Item>("ShaderCollectionItem")
                    ->Attribute(AZ::Script::Attributes::Scope, AZ::Script::Attributes::ScopeFlags::Automation)
                    ->Attribute(AZ::Script::Attributes::Category, "Shader")
                    ->Attribute(AZ::Script::Attributes::Module, "shader")
                    ->Method("GetShaderAsset", &Item::GetShaderAsset)
                    ->Method("GetShaderAssetId", &Item::GetShaderAssetId)
                    ->Method("GetShaderVariantId", &Item::GetShaderVariantId)
                    ->Method("GetShaderOptionGroup", &Item::GetShaderOptionGroup)
                    ->Method("GetDrawItemType", &Item::GetDrawItemType)
                    ->Method("MaterialOwnsShaderOption", static_cast<bool (Item::*)(const Name&) const>(&Item::MaterialOwnsShaderOption));
            }
        }

        size_t ShaderCollection::size() const
        {
            return m_shaderItems.size();
        }

        ShaderCollection::iterator ShaderCollection::begin()
        {
            return m_shaderItems.begin();
        }

        ShaderCollection::const_iterator ShaderCollection::begin() const
        {
            return m_shaderItems.begin();
        }

        ShaderCollection::iterator ShaderCollection::end()
        {
            return m_shaderItems.end();
        }

        ShaderCollection::const_iterator ShaderCollection::end() const
        {
            return m_shaderItems.end();
        }

        ShaderCollection::Item::Item()
            : m_renderStatesOverlay(RHI::GetInvalidRenderStates())
        {
        }

        ShaderCollection::Item& ShaderCollection::operator[](size_t i)
        {
            return m_shaderItems[i];
        }

        const ShaderCollection::Item& ShaderCollection::operator[](size_t i) const
        {
            return m_shaderItems[i];
        }

        ShaderCollection::Item& ShaderCollection::operator[](const AZ::Name& shaderTag)
        {
            return m_shaderItems[m_shaderTagIndexMap.Find(shaderTag).GetIndex()];
        }

        const ShaderCollection::Item& ShaderCollection::operator[](const AZ::Name& shaderTag) const
        {
            return m_shaderItems[m_shaderTagIndexMap.Find(shaderTag).GetIndex()];
        }

        bool ShaderCollection::HasShaderTag(const AZ::Name& shaderTag) const
        {
            return (m_shaderTagIndexMap.Find(shaderTag).IsValid());
        }

        void ShaderCollection::TryReplaceShaderAsset(const Data::Asset<ShaderAsset>& newShaderAsset)
        {
            for (auto& shaderItem : m_shaderItems)
            {
                shaderItem.TryReplaceShaderAsset(newShaderAsset);
            }
        }

        bool ShaderCollection::InitializeShaderOptionGroups()
        {
            for (auto& shaderItem : m_shaderItems)
            {
                if (!shaderItem.InitializeShaderOptionGroup())
                {
                    return false;
                }
            }
            return true;
        }

        ShaderCollection::Item::Item(
            const Data::Asset<ShaderAsset>& shaderAsset, const AZ::Name& shaderTag, DrawItemType drawItemType, ShaderVariantId variantId)
            : m_renderStatesOverlay(RHI::GetInvalidRenderStates())
            , m_shaderAsset(shaderAsset)
            , m_shaderVariantId(variantId)
            , m_shaderTag(shaderTag)
            , m_drawItemType(drawItemType)
            , m_shaderOptionGroup(shaderAsset->GetShaderOptionGroupLayout(), variantId)
        {
        }

        ShaderCollection::Item::Item(
            Data::Asset<ShaderAsset>&& shaderAsset, const AZ::Name& shaderTag, DrawItemType drawItemType, ShaderVariantId variantId)
            : m_renderStatesOverlay(RHI::GetInvalidRenderStates())
            , m_shaderAsset(AZStd::move(shaderAsset))
            , m_shaderVariantId(variantId)
            , m_shaderTag(shaderTag)
            , m_drawItemType(drawItemType)
            , m_shaderOptionGroup(shaderAsset->GetShaderOptionGroupLayout(), variantId)
        {
        }

        const Data::Asset<ShaderAsset>& ShaderCollection::Item::GetShaderAsset() const
        {
            return m_shaderAsset;
        }

        const ShaderVariantId& ShaderCollection::Item::GetShaderVariantId() const
        {
            return m_shaderOptionGroup.GetShaderVariantId();
        }

        const ShaderOptionGroup* ShaderCollection::Item::GetShaderOptions() const
        {
            return &m_shaderOptionGroup;
        }

        ShaderOptionGroup* ShaderCollection::Item::GetShaderOptions()
        {
            return &m_shaderOptionGroup;
        }

        bool ShaderCollection::Item::MaterialOwnsShaderOption(const AZ::Name& shaderOptionName) const
        {
            return m_ownedShaderOptionIndices.contains(m_shaderOptionGroup.FindShaderOptionIndex(shaderOptionName));
        }

        bool ShaderCollection::Item::MaterialOwnsShaderOption(ShaderOptionIndex shaderOptionIndex) const
        {
            return m_ownedShaderOptionIndices.contains(shaderOptionIndex);
        }

        const RHI::RenderStates* ShaderCollection::Item::GetRenderStatesOverlay() const
        {
            return &m_renderStatesOverlay;
        }

        RHI::RenderStates* ShaderCollection::Item::GetRenderStatesOverlay()
        {
            return &m_renderStatesOverlay;
        }

        RHI::DrawListTag ShaderCollection::Item::GetDrawListTagOverride() const
        {
            return m_drawListTagOverride;
        }

        void ShaderCollection::Item::SetDrawListTagOverride(RHI::DrawListTag drawList)
        {
            m_drawListTagOverride = drawList;
        }

        void ShaderCollection::Item::SetDrawListTagOverride(const AZ::Name& drawListName)
        {
            if (drawListName.IsEmpty())
            {
                m_drawListTagOverride.Reset();
                return;
            }

            RHI::DrawListTagRegistry* drawListTagRegistry = RHI::RHISystemInterface::Get()->GetDrawListTagRegistry();
            // Note: we should use FindTag instead of AcquireTag to avoid occupy DrawListTag entries. 
            RHI::DrawListTag newTag = drawListTagRegistry->FindTag(drawListName);
            if (newTag.IsNull())
            {
                AZ_Error("ShaderCollection", false, "Failed to set draw list with name: %s.", drawListName.GetCStr());
                return;
            }
            m_drawListTagOverride = newTag;
        }

        void ShaderCollection::Item::SetEnabled(bool enabled)
        {
            m_enabled = enabled;
        }

        bool ShaderCollection::Item::IsEnabled() const
        {
            return m_enabled;
        }

        const AZ::Name& ShaderCollection::Item::GetShaderTag() const
        {
            return m_shaderTag;
        }

        const Data::AssetId& ShaderCollection::Item::GetShaderAssetId() const
        {
            return m_shaderAsset->GetId();
        }

        const AZ::RPI::ShaderOptionGroup& ShaderCollection::Item::GetShaderOptionGroup() const
        {
            return m_shaderOptionGroup;
        }

        bool ShaderCollection::Item::InitializeShaderOptionGroup()
        {
            if (!m_shaderAsset.IsReady())
            {
                return false;
            }
            m_shaderOptionGroup = ShaderOptionGroup{
                m_shaderAsset->GetShaderOptionGroupLayout(),
                m_shaderVariantId };
            return true;
        }

        void ShaderCollection::Item::TryReplaceShaderAsset(const Data::Asset<ShaderAsset>& newShaderAsset)
        {
            if (newShaderAsset.GetId() != m_shaderAsset.GetId())
            {
                return;
            }
            m_shaderAsset = newShaderAsset;
            [[maybe_unused]] bool success = InitializeShaderOptionGroup();
            AZ_Assert(success, "Failed to InitializeShaderOptionGroup using shaderAsset with uuid=%s and hint=%s"
                , newShaderAsset.GetId().ToFixedString().c_str(), newShaderAsset.GetHint().c_str());
        }

    } // namespace RPI
} // namespace AZ