Prefabs UI | Finalize RootAliasPath as a standardized instance handle system, and refactor PrefabFocusHandler to take advantage of it. (#7365)

* Move the RootAliasPath class definition to the Prefab EOS. Introduce utility functions to enable using the path as an instance handler in UI.

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>

* Refactor the PrefabFocusHandler to use the utility functions that were introduced to the Prefab EOS. Remove the m_instanceFocusHierarchy from the Focus Handler, and simplify calls for performance.

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>

* Minor adjustments after PR discussion

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>

* Silly mistake in function simplification. Turns out == and != are different.

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>

* Fixed doxygen comments

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>

* Add numeric cast when storing result of AZStd::distance to int.

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>

* Restore additional checks in FocusOnparentOfFocusedPrefab that got removed during merge.

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>

Code/Framework/AzToolsFramework/AzToolsFramework/Entity/PrefabEditorEntityOwnershipInterface.h

@@ -19,45 +19,69 @@
 
 namespace AzToolsFramework
 {
+    namespace Prefab
+    {
+        //! A RootAliasPath can be used to store an alias path that starts from the Prefab EOS root instance.
+        //! The root instance itself is included in the path. These can be used as Instance handles across systems
+        //! that do not have visibility over InstanceOptionalReferences, or that need to store Instance handles
+        //! for longer than just the span of a function without the risk of them going out of scope.
+        using RootAliasPath = AliasPath;
+    }
 
     class PrefabEditorEntityOwnershipInterface
     {
     public:
         AZ_RTTI(PrefabEditorEntityOwnershipInterface, "{38E764BA-A089-49F3-848F-46018822CE2E}");
 
+        //! Returns whether the system has a root instance assigned.
+        //! @return True if a root prefab is assigned, false otherwise.
+        virtual bool IsRootPrefabAssigned() const = 0;
+
+        //! Returns an optional reference to the root prefab instance.
+        virtual Prefab::InstanceOptionalReference GetRootPrefabInstance() = 0;
+
+        //! Returns the template id for the root prefab instance.
+        virtual Prefab::TemplateId GetRootPrefabTemplateId() = 0;
+
+        virtual void CreateNewLevelPrefab(AZStd::string_view filename, const AZStd::string& templateFilename) = 0;
+
         //! Creates a prefab instance with the provided entities and nestedPrefabInstances.
-        //! /param entities The entities to put under the new prefab.
-        //! /param nestedPrefabInstances The nested prefab instances to put under the new prefab.
-        //! /param filePath The filepath corresponding to the prefab file to be created.
-        //! /param instanceToParentUnder The instance the newly created prefab instance is parented under.
-        //! /return The optional reference to the prefab created.
+        //! @param entities The entities to put under the new prefab.
+        //! @param nestedPrefabInstances The nested prefab instances to put under the new prefab.
+        //! @param filePath The filepath corresponding to the prefab file to be created.
+        //! @param instanceToParentUnder The instance the newly created prefab instance is parented under.
+        //! @return The optional reference to the prefab created.
         virtual Prefab::InstanceOptionalReference CreatePrefab(
             const AZStd::vector<AZ::Entity*>& entities, AZStd::vector<AZStd::unique_ptr<Prefab::Instance>>&& nestedPrefabInstances,
             AZ::IO::PathView filePath, Prefab::InstanceOptionalReference instanceToParentUnder = AZStd::nullopt) = 0;
 
         //! Instantiate the prefab file provided.
-        //! /param filePath The filepath for the prefab file the instance should be created from.
-        //! /param instanceToParentUnder The instance the newly instantiated prefab instance is parented under.
-        //! /return The optional reference to the prefab instance.
+        //! @param filePath The filepath for the prefab file the instance should be created from.
+        //! @param instanceToParentUnder The instance the newly instantiated prefab instance is parented under.
+        //! @return The optional reference to the prefab instance.
         virtual Prefab::InstanceOptionalReference InstantiatePrefab(
             AZ::IO::PathView filePath, Prefab::InstanceOptionalReference instanceToParentUnder = AZStd::nullopt) = 0;
 
-        virtual Prefab::InstanceOptionalReference GetRootPrefabInstance() = 0;
-
-        virtual Prefab::TemplateId GetRootPrefabTemplateId() = 0;
+        virtual void StartPlayInEditor() = 0;
+        virtual void StopPlayInEditor() = 0;
 
         //! Get all Assets generated by Prefab processing when entering Play-In Editor mode (Ctrl+G)
-        //! /return The vector of Assets generated by Prefab processing
+        //! @return The vector of Assets generated by Prefab processing
         virtual const Prefab::PrefabConversionUtils::InMemorySpawnableAssetContainer::SpawnableAssets& GetPlayInEditorAssetData() const = 0;
 
         virtual bool LoadFromStream(AZ::IO::GenericStream& stream, AZStd::string_view filename) = 0;
         virtual bool SaveToStream(AZ::IO::GenericStream& stream, AZStd::string_view filename) = 0;
 
-        virtual void StartPlayInEditor() = 0;
-        virtual void StopPlayInEditor() = 0;
+        //! Returns the reference to the instance corresponding to the RootAliasPath provided.
+        //! @param rootAliasPath The RootAliasPath to be queried.
+        //! @return A reference to the instance if valid, AZStd::nullopt otherwise.
+        virtual Prefab::InstanceOptionalReference GetInstanceReferenceFromRootAliasPath(Prefab::RootAliasPath rootAliasPath) const = 0;
 
-        virtual void CreateNewLevelPrefab(AZStd::string_view filename, const AZStd::string& templateFilename) = 0;
-
-        virtual bool IsRootPrefabAssigned() const = 0;
+        //! Allows to iterate through all instances referenced in the path, from the root down.
+        //! @param rootAliasPath The RootAliasPath to iterate through. If invalid, callback will not be called.
+        //! @param callback The function to call on each instance. If it returns true, it prevents the rest of the path from being called.
+        //! @return True if the iteration was halted by a callback returning true, false otherwise. Also returns false if the path is invalid.
+        virtual bool GetInstancesInRootAliasPath(
+            Prefab::RootAliasPath rootAliasPath, const AZStd::function<bool(const Prefab::InstanceOptionalReference)>& callback) const = 0;
     };
 }

Code/Framework/AzToolsFramework/AzToolsFramework/Entity/PrefabEditorEntityOwnershipService.cpp

@@ -510,6 +510,70 @@ namespace AzToolsFramework
         m_playInEditorData.m_isEnabled = false;
     }
 
+    bool PrefabEditorEntityOwnershipService::IsValidRootAliasPath(Prefab::RootAliasPath rootAliasPath) const
+    {
+        return GetInstanceReferenceFromRootAliasPath(rootAliasPath) != AZStd::nullopt;
+    }
+
+    Prefab::InstanceOptionalReference PrefabEditorEntityOwnershipService::GetInstanceReferenceFromRootAliasPath(
+        Prefab::RootAliasPath rootAliasPath) const
+    {
+        Prefab::InstanceOptionalReference instance = *m_rootInstance;
+
+        for (const auto& pathElement : rootAliasPath)
+        {
+            if (pathElement.Native() == rootAliasPath.begin()->Native())
+            {
+                // If the root is not the root Instance, the rootAliasPath is invalid.
+                if (pathElement.Native() != instance->get().GetInstanceAlias())
+                {
+                    return Prefab::InstanceOptionalReference();
+                }
+            }
+            else
+            {
+                // If the instance alias can't be found, the rootAliasPath is invalid.
+                instance = instance->get().FindNestedInstance(pathElement.Native());
+                if (!instance.has_value())
+                {
+                    return Prefab::InstanceOptionalReference();
+                }
+            }
+        }
+
+        return instance;
+    }
+
+    bool PrefabEditorEntityOwnershipService::GetInstancesInRootAliasPath(
+        Prefab::RootAliasPath rootAliasPath, const AZStd::function<bool(const Prefab::InstanceOptionalReference)>& callback) const
+    {
+        if (!IsValidRootAliasPath(rootAliasPath))
+        {
+            return false;
+        }
+
+        Prefab::InstanceOptionalReference instance;
+
+        for (const auto& pathElement : rootAliasPath)
+        {
+            if (!instance.has_value())
+            {
+                instance = *m_rootInstance;
+            }
+            else
+            {
+                instance = instance->get().FindNestedInstance(pathElement.Native());
+            }
+
+            if(callback(instance))
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     //////////////////////////////////////////////////////////////////////////
     // Slice Buses implementation with Assert(false), this will exist only during Slice->Prefab
     // development to pinpoint and replace specific calls to Slice system

Code/Framework/AzToolsFramework/AzToolsFramework/Entity/PrefabEditorEntityOwnershipService.h

@@ -169,11 +169,17 @@ namespace AzToolsFramework
         void CreateNewLevelPrefab(AZStd::string_view filename, const AZStd::string& templateFilename) override;
         bool IsRootPrefabAssigned() const override;
 
+        Prefab::InstanceOptionalReference GetInstanceReferenceFromRootAliasPath(Prefab::RootAliasPath rootAliasPath) const override;
+        bool GetInstancesInRootAliasPath(
+            Prefab::RootAliasPath rootAliasPath, const AZStd::function<bool(const Prefab::InstanceOptionalReference)>& callback) const override;
+
     protected:
 
         AZ::SliceComponent::SliceInstanceAddress GetOwningSlice() override;
 
     private:
+        bool IsValidRootAliasPath(Prefab::RootAliasPath rootAliasPath) const;
+
         struct PlayInEditorData
         {
             AzToolsFramework::Prefab::PrefabConversionUtils::InMemorySpawnableAssetContainer m_assetsCache;

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/Instance/Instance.h

@@ -37,7 +37,6 @@ namespace AzToolsFramework
 
         using AliasPath = AZ::IO::Path;
         using AliasPathView = AZ::IO::PathView;
-        using RootAliasPath = AliasPath;
         using EntityAlias = AZStd::string;
         using EntityAliasView = AZStd::string_view;
         using InstanceAlias = AZStd::string;

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusHandler.cpp

@@ -11,7 +11,6 @@
 #include <AzToolsFramework/Commands/SelectionCommand.h>
 #include <AzToolsFramework/ContainerEntity/ContainerEntityInterface.h>
 #include <AzToolsFramework/Entity/EditorEntityHelpers.h>
-#include <AzToolsFramework/Entity/PrefabEditorEntityOwnershipInterface.h>
 #include <AzToolsFramework/Entity/ReadOnly/ReadOnlyEntityInterface.h>
 #include <AzToolsFramework/Prefab/Instance/Instance.h>
 #include <AzToolsFramework/Prefab/Instance/InstanceEntityMapperInterface.h>
@@ -112,15 +111,24 @@ namespace AzToolsFramework::Prefab
         [[maybe_unused]] AzFramework::EntityContextId entityContextId)
     {
         // If only one instance is in the hierarchy, this operation is invalid
-        size_t hierarchySize = m_instanceFocusHierarchy.size();
-        if (hierarchySize <= 1)
+        if (m_rootAliasFocusPathLength <= 1)
         {
-            return AZ::Failure(
-                AZStd::string("Prefab Focus Handler: Could not complete FocusOnParentOfFocusedPrefab operation while focusing on the root."));
+            return AZ::Failure(AZStd::string(
+                "Prefab Focus Handler: Could not complete FocusOnParentOfFocusedPrefab operation while focusing on the root."));
         }
 
+        RootAliasPath parentPath = m_rootAliasFocusPath;
+        parentPath.RemoveFilename();
+
         // Retrieve parent of currently focused prefab.
-        InstanceOptionalReference parentInstance = GetInstanceReferenceFromRootAliasPath(m_instanceFocusHierarchy[hierarchySize - 2]);
+        InstanceOptionalReference parentInstance = GetInstanceReference(parentPath);
+
+        // If only one instance is in the hierarchy, this operation is invalid
+        if (!parentInstance.has_value())
+        {
+            return AZ::Failure(AZStd::string(
+                "Prefab Focus Handler: Could not retrieve parent of current focus in FocusOnParentOfFocusedPrefab."));
+        }
 
         // Use container entity of parent Instance for focus operations.
         AZ::EntityId entityId = parentInstance->get().GetContainerEntityId();
@@ -149,12 +157,26 @@ namespace AzToolsFramework::Prefab
 
     PrefabFocusOperationResult PrefabFocusHandler::FocusOnPathIndex([[maybe_unused]] AzFramework::EntityContextId entityContextId, int index)
     {
-        if (index < 0 || index >= m_instanceFocusHierarchy.size())
+        if (index < 0 || index >= m_rootAliasFocusPathLength)
         {
             return AZ::Failure(AZStd::string("Prefab Focus Handler: Invalid index on FocusOnPathIndex."));
         }
 
-        InstanceOptionalReference focusedInstance = GetInstanceReferenceFromRootAliasPath(m_instanceFocusHierarchy[index]);
+        int i = 0;
+        RootAliasPath indexedPath;
+        for (const auto& pathElement : m_rootAliasFocusPath)
+        {
+            indexedPath.Append(pathElement);
+
+            if (i == index)
+            {
+                break;
+            }
+
+            ++i;
+        }
+
+        InstanceOptionalReference focusedInstance = GetInstanceReference(indexedPath);
 
         if (!focusedInstance.has_value())
         {
@@ -197,13 +219,14 @@ namespace AzToolsFramework::Prefab
         }
 
         // Close all container entities in the old path.
-        SetInstanceContainersOpenState(m_instanceFocusHierarchy, false);
+        SetInstanceContainersOpenState(m_rootAliasFocusPath, false);
 
-        const RootAliasPath previousContainerRootAliasPath = m_focusedInstanceRootAliasPath;
-        const InstanceOptionalConstReference previousFocusedInstance = GetInstanceReferenceFromRootAliasPath(previousContainerRootAliasPath);
+        const RootAliasPath previousContainerRootAliasPath = m_rootAliasFocusPath;
+        const InstanceOptionalConstReference previousFocusedInstance = GetInstanceReference(previousContainerRootAliasPath);
 
-        m_focusedInstanceRootAliasPath = focusedInstance->get().GetAbsoluteInstanceAliasPath();
+        m_rootAliasFocusPath = focusedInstance->get().GetAbsoluteInstanceAliasPath();
         m_focusedTemplateId = focusedInstance->get().GetTemplateId();
+        m_rootAliasFocusPathLength = aznumeric_cast<int>(AZStd::distance(m_rootAliasFocusPath.begin(), m_rootAliasFocusPath.end()));
 
         // Focus on the descendants of the container entity in the Editor, if the interface is initialized.
         if (m_focusModeInterface)
@@ -231,11 +254,10 @@ namespace AzToolsFramework::Prefab
         }
 
         // Refresh path variables.
-        RefreshInstanceFocusList();
         RefreshInstanceFocusPath();
 
         // Open all container entities in the new path.
-        SetInstanceContainersOpenState(m_instanceFocusHierarchy, true);
+        SetInstanceContainersOpenState(m_rootAliasFocusPath, true);
 
         PrefabFocusNotificationBus::Broadcast(&PrefabFocusNotifications::OnPrefabFocusChanged);
 
@@ -250,12 +272,12 @@ namespace AzToolsFramework::Prefab
     InstanceOptionalReference PrefabFocusHandler::GetFocusedPrefabInstance(
         [[maybe_unused]] AzFramework::EntityContextId entityContextId) const
     {
-        return GetInstanceReferenceFromRootAliasPath(m_focusedInstanceRootAliasPath);
+        return GetInstanceReference(m_rootAliasFocusPath);
     }
 
     AZ::EntityId PrefabFocusHandler::GetFocusedPrefabContainerEntityId([[maybe_unused]] AzFramework::EntityContextId entityContextId) const
     {
-        if (const InstanceOptionalConstReference instance = GetInstanceReferenceFromRootAliasPath(m_focusedInstanceRootAliasPath); instance.has_value())
+        if (const InstanceOptionalConstReference instance = GetInstanceReference(m_rootAliasFocusPath); instance.has_value())
         {
             return instance->get().GetContainerEntityId();
         }
@@ -276,7 +298,7 @@ namespace AzToolsFramework::Prefab
             return false;
         }
 
-        return (instance->get().GetAbsoluteInstanceAliasPath() == m_focusedInstanceRootAliasPath);
+        return (instance->get().GetAbsoluteInstanceAliasPath() == m_rootAliasFocusPath);
     }
 
     bool PrefabFocusHandler::IsOwningPrefabInFocusHierarchy(AZ::EntityId entityId) const
@@ -289,7 +311,7 @@ namespace AzToolsFramework::Prefab
         InstanceOptionalConstReference instance = m_instanceEntityMapperInterface->FindOwningInstance(entityId);
         while (instance.has_value())
         {
-            if (instance->get().GetAbsoluteInstanceAliasPath() == m_focusedInstanceRootAliasPath)
+            if (instance->get().GetAbsoluteInstanceAliasPath() == m_rootAliasFocusPath)
             {
                 return true;
             }
@@ -302,42 +324,49 @@ namespace AzToolsFramework::Prefab
 
     const AZ::IO::Path& PrefabFocusHandler::GetPrefabFocusPath([[maybe_unused]] AzFramework::EntityContextId entityContextId) const
     {
-        return m_instanceFocusPath;
+        return m_filenameFocusPath;
     }
 
     const int PrefabFocusHandler::GetPrefabFocusPathLength([[maybe_unused]] AzFramework::EntityContextId entityContextId) const
     {
-        return aznumeric_cast<int>(m_instanceFocusHierarchy.size());
+        return m_rootAliasFocusPathLength;
     }
 
     void PrefabFocusHandler::OnContextReset()
     {
-        // Clear the old focus vector
-        m_instanceFocusHierarchy.clear();
-
         // Focus on the root prefab (AZ::EntityId() will default to it)
         FocusOnPrefabInstanceOwningEntityId(AZ::EntityId());
     }
 
     void PrefabFocusHandler::OnEntityInfoUpdatedName(AZ::EntityId entityId, [[maybe_unused]]const AZStd::string& name)
+    {
+        PrefabEditorEntityOwnershipInterface* prefabEditorEntityOwnershipInterface =
+            AZ::Interface<PrefabEditorEntityOwnershipInterface>::Get();
+
+        if (prefabEditorEntityOwnershipInterface)
         {
             // Determine if the entityId is the container for any of the instances in the vector.
-        auto result = AZStd::find_if(
-            m_instanceFocusHierarchy.begin(), m_instanceFocusHierarchy.end(),
-            [&, entityId](const RootAliasPath& rootAliasPath)
+            bool match = prefabEditorEntityOwnershipInterface->GetInstancesInRootAliasPath(
+                m_rootAliasFocusPath,
+                [&](const Prefab::InstanceOptionalReference instance)
                 {
-                InstanceOptionalReference instance = GetInstanceReferenceFromRootAliasPath(rootAliasPath);
-                return (instance->get().GetContainerEntityId() == entityId);
+                    if (instance->get().GetContainerEntityId() == entityId)
+                    {
+                        return true;
+                    }
+
+                    return false;
                 }
             );
 
-        if (result != m_instanceFocusHierarchy.end())
+            if (match)
             {
                 // Refresh the path and notify changes.
                 RefreshInstanceFocusPath();
                 PrefabFocusNotificationBus::Broadcast(&PrefabFocusNotifications::OnPrefabFocusChanged);
             }
         }
+    }
 
     void PrefabFocusHandler::OnPrefabInstancePropagationEnd()
     {
@@ -347,65 +376,64 @@ namespace AzToolsFramework::Prefab
     }
 
     void PrefabFocusHandler::OnPrefabTemplateDirtyFlagUpdated(TemplateId templateId, [[maybe_unused]] bool status)
+    {
+        PrefabEditorEntityOwnershipInterface* prefabEditorEntityOwnershipInterface =
+            AZ::Interface<PrefabEditorEntityOwnershipInterface>::Get();
+
+        if (prefabEditorEntityOwnershipInterface)
         {
             // Determine if the templateId matches any of the instances in the vector.
-        auto result = AZStd::find_if(
-            m_instanceFocusHierarchy.begin(), m_instanceFocusHierarchy.end(),
-            [&, templateId](const RootAliasPath& rootAliasPath)
+            bool match = prefabEditorEntityOwnershipInterface->GetInstancesInRootAliasPath(
+                m_rootAliasFocusPath,
+                [&](const Prefab::InstanceOptionalReference instance)
+                {
+                    if (instance->get().GetTemplateId() == templateId)
                     {
-                InstanceOptionalReference instance = GetInstanceReferenceFromRootAliasPath(rootAliasPath);
-                return (instance->get().GetTemplateId() == templateId);
+                        return true;
+                    }
+
+                    return false;
                 }
             );
 
-        if (result != m_instanceFocusHierarchy.end())
+            if (match)
             {
                 // Refresh the path and notify changes.
                 RefreshInstanceFocusPath();
                 PrefabFocusNotificationBus::Broadcast(&PrefabFocusNotifications::OnPrefabFocusChanged);
             }
         }
-
-    void PrefabFocusHandler::RefreshInstanceFocusList()
-    {
-        m_instanceFocusHierarchy.clear();
-
-        InstanceOptionalConstReference currentInstance = GetInstanceReferenceFromRootAliasPath(m_focusedInstanceRootAliasPath);
-        while (currentInstance.has_value())
-        {
-            m_instanceFocusHierarchy.emplace_back(currentInstance->get().GetAbsoluteInstanceAliasPath());
-            currentInstance = currentInstance->get().GetParentInstance();
-        }
-
-        // Invert the vector, since we need the top instance to be at index 0.
-        AZStd::reverse(m_instanceFocusHierarchy.begin(), m_instanceFocusHierarchy.end());
     }
 
     void PrefabFocusHandler::RefreshInstanceFocusPath()
     {
-        auto prefabSystemComponentInterface = AZ::Interface<PrefabSystemComponentInterface>::Get();
+        m_filenameFocusPath.clear();
         
-        m_instanceFocusPath.clear();
+        PrefabEditorEntityOwnershipInterface* prefabEditorEntityOwnershipInterface =
+            AZ::Interface<PrefabEditorEntityOwnershipInterface>::Get();
+        PrefabSystemComponentInterface* prefabSystemComponentInterface = AZ::Interface<PrefabSystemComponentInterface>::Get();
 
-        size_t index = 0;
-        size_t maxIndex = m_instanceFocusHierarchy.size() - 1;
+        if (prefabEditorEntityOwnershipInterface && prefabSystemComponentInterface)
+        {
+            int i = 0;
 
-        for (const RootAliasPath& rootAliasPath : m_instanceFocusHierarchy)
+            prefabEditorEntityOwnershipInterface->GetInstancesInRootAliasPath(
+                m_rootAliasFocusPath,
+                [&](const Prefab::InstanceOptionalReference instance)
                 {
-            const InstanceOptionalConstReference instance = GetInstanceReferenceFromRootAliasPath(rootAliasPath);
                     if (instance.has_value())
                     {
                         AZStd::string prefabName;
 
-                if (index < maxIndex)
+                        if (i == m_rootAliasFocusPathLength - 1)
                         {
-                    // Get the filename without the extension (stem).
-                    prefabName = instance->get().GetTemplateSourcePath().Stem().Native();
+                            // Get the full filename.
+                            prefabName = instance->get().GetTemplateSourcePath().Filename().Native();
                         }
                         else
                         {
-                    // Get the full filename.
-                    prefabName = instance->get().GetTemplateSourcePath().Filename().Native();
+                            // Get the filename without the extension (stem).
+                            prefabName = instance->get().GetTemplateSourcePath().Stem().Native();
                         }
 
                         if (prefabSystemComponentInterface->IsTemplateDirty(instance->get().GetTemplateId()))
@@ -413,14 +441,17 @@ namespace AzToolsFramework::Prefab
                             prefabName += "*";
                         }
 
-                m_instanceFocusPath.Append(prefabName);
+                        m_filenameFocusPath.Append(prefabName);
                     }
 
-            ++index;
+                    ++i;
+                    return false;
+                }
+            );
         }
     }
 
-    void PrefabFocusHandler::SetInstanceContainersOpenState(const AZStd::vector<RootAliasPath>& instances, bool openState) const
+    void PrefabFocusHandler::SetInstanceContainersOpenState(const RootAliasPath& rootAliasPath, bool openState) const
     {
         // If this is called outside the Editor, this interface won't be initialized.
         if (!m_containerEntityInterface)
@@ -428,51 +459,34 @@ namespace AzToolsFramework::Prefab
             return;
         }
 
-        for (const RootAliasPath& rootAliasPath : instances)
-        {
-            InstanceOptionalReference instance = GetInstanceReferenceFromRootAliasPath(rootAliasPath);
+        PrefabEditorEntityOwnershipInterface* prefabEditorEntityOwnershipInterface =
+            AZ::Interface<PrefabEditorEntityOwnershipInterface>::Get();
 
-            if (instance.has_value())
+        if (prefabEditorEntityOwnershipInterface)
+        {
+            prefabEditorEntityOwnershipInterface->GetInstancesInRootAliasPath(
+                rootAliasPath,
+                [&](const Prefab::InstanceOptionalReference instance)
                 {
                     m_containerEntityInterface->SetContainerOpen(instance->get().GetContainerEntityId(), openState);
+
+                    return false;
                 }
+            );
         }
     }
 
-    InstanceOptionalReference PrefabFocusHandler::GetInstanceReferenceFromRootAliasPath(RootAliasPath rootAliasPath) const
+    InstanceOptionalReference PrefabFocusHandler::GetInstanceReference(RootAliasPath rootAliasPath) const
     {
         PrefabEditorEntityOwnershipInterface* prefabEditorEntityOwnershipInterface =
             AZ::Interface<PrefabEditorEntityOwnershipInterface>::Get();
 
         if (prefabEditorEntityOwnershipInterface)
         {
-            InstanceOptionalReference instance = prefabEditorEntityOwnershipInterface->GetRootPrefabInstance();
-
-            for (const auto& pathElement : rootAliasPath)
-            {
-                if (pathElement.Native() == rootAliasPath.begin()->Native())
-                {
-                    // If the root is not the root Instance, the rootAliasPath is invalid.
-                    if (pathElement.Native() != instance->get().GetInstanceAlias())
-                    {
-                        return InstanceOptionalReference();
-                    }
-                }
-                else
-                {
-                    // If the instance alias can't be found, the rootAliasPath is invalid.
-                    instance = instance->get().FindNestedInstance(pathElement.Native());
-                    if (!instance.has_value())
-                    {
-                        return InstanceOptionalReference();
-                    }
-                }
-            }
-
-            return instance;
+            return prefabEditorEntityOwnershipInterface->GetInstanceReferenceFromRootAliasPath(rootAliasPath);
         }
 
-        return InstanceOptionalReference();
+        return AZStd::nullopt;
     }
 
 } // namespace AzToolsFramework::Prefab

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusHandler.h

@@ -12,6 +12,7 @@
 
 #include <AzToolsFramework/Entity/EditorEntityContextBus.h>
 #include <AzToolsFramework/Entity/EditorEntityInfoBus.h>
+#include <AzToolsFramework/Entity/PrefabEditorEntityOwnershipInterface.h>
 #include <AzToolsFramework/FocusMode/FocusModeInterface.h>
 #include <AzToolsFramework/Prefab/PrefabFocusInterface.h>
 #include <AzToolsFramework/Prefab/PrefabFocusPublicInterface.h>
@@ -28,6 +29,7 @@ namespace AzToolsFramework
 namespace AzToolsFramework::Prefab
 {
     class InstanceEntityMapperInterface;
+    class PrefabSystemComponentInterface;
 
     //! Handles Prefab Focus mode, determining which prefab file entity changes will target.
     class PrefabFocusHandler final
@@ -73,22 +75,20 @@ namespace AzToolsFramework::Prefab
         
     private:
         PrefabFocusOperationResult FocusOnPrefabInstance(InstanceOptionalReference focusedInstance);
-        void RefreshInstanceFocusList();
         void RefreshInstanceFocusPath();
 
-        void SetInstanceContainersOpenState(const AZStd::vector<RootAliasPath>& instances, bool openState) const;
+        void SetInstanceContainersOpenState(const RootAliasPath& rootAliasPath, bool openState) const;
 
-        InstanceOptionalReference GetInstanceReferenceFromRootAliasPath(RootAliasPath rootAliasPath) const;
+        InstanceOptionalReference GetInstanceReference(RootAliasPath rootAliasPath) const;
 
         //! The alias path for the instance the editor is currently focusing on, starting from the root instance.
-        RootAliasPath m_focusedInstanceRootAliasPath = RootAliasPath();
+        RootAliasPath m_rootAliasFocusPath = RootAliasPath();
         //! The templateId of the focused instance.
         TemplateId m_focusedTemplateId;
-        //! The list of instances going from the root (index 0) to the focused instance,
-        //! referenced by their alias path from the root instance.
-        AZStd::vector<RootAliasPath> m_instanceFocusHierarchy;
         //! A path containing the filenames of the instances in the focus hierarchy, separated with a /.
-        AZ::IO::Path m_instanceFocusPath;
+        AZ::IO::Path m_filenameFocusPath;
+        //! The length of the current focus path. Stored to simplify internal checks.
+        int m_rootAliasFocusPathLength = 0;
 
         ContainerEntityInterface* m_containerEntityInterface = nullptr;
         FocusModeInterface* m_focusModeInterface = nullptr;