terra-llm/Assets/LLM/source/utils/hierarchy_utils.py

from pathlib import Path

class HierarchyParser:
    """
    Parses a Unity scene's object_map to identify parent-child relationships
    and determine the root objects in the hierarchy.
    """
    def __init__(self, object_map):
        self.object_map = object_map
        self.transform_to_gameobject = {}
        self.gameobject_to_transform = {}
        self.transform_children = {}
        self.prefab_instances = {}
        self.stripped_transforms = {}
        self.all_child_ids = set()

        self._build_maps()
        self._find_all_children()

    def _build_maps(self):
        """
        First pass over the object map to build essential lookups for transforms,
        GameObjects, and PrefabInstances.
        """
        for file_id, obj_data in self.object_map.items():
            if 'Transform' in obj_data:
                transform_info = obj_data['Transform']
                gameobject_id = transform_info.get('m_GameObject', {}).get('fileID')
                if gameobject_id:
                    self.transform_to_gameobject[file_id] = gameobject_id
                    self.gameobject_to_transform[gameobject_id] = file_id
                
                parent_id = transform_info.get('m_Father', {}).get('fileID')
                if parent_id and parent_id != 0:
                    if parent_id not in self.transform_children:
                        self.transform_children[parent_id] = []
                    self.transform_children[parent_id].append(file_id)

                prefab_instance_id = transform_info.get('m_PrefabInstance', {}).get('fileID')
                if prefab_instance_id:
                    self.stripped_transforms[file_id] = prefab_instance_id

            elif 'PrefabInstance' in obj_data:
                prefab_info = obj_data['PrefabInstance']
                modifications = prefab_info.get('m_Modification', {})
                parent_id = modifications.get('m_TransformParent', {}).get('fileID')
                if parent_id and parent_id != 0:
                    self.prefab_instances[file_id] = {'parent': parent_id}

    def _find_all_children(self):
        """
        Iterates through the built maps to populate a set of all fileIDs that
        are children of another object.
        """
        # Children of regular GameObjects (via Transform)
        for parent_id, children in self.transform_children.items():
            for child_transform_id in children:
                child_go_id = self.transform_to_gameobject.get(child_transform_id)
                if child_go_id:
                    self.all_child_ids.add(child_go_id)
                # Also add stripped transforms that point to prefabs
                elif child_transform_id in self.stripped_transforms:
                    prefab_instance_id = self.stripped_transforms[child_transform_id]
                    self.all_child_ids.add(prefab_instance_id)

        # Children that are PrefabInstances
        for prefab_id, prefab_data in self.prefab_instances.items():
            self.all_child_ids.add(prefab_id)

    def get_root_object_ids(self):
        """
        Identifies all root objects (GameObjects and PrefabInstances) and returns
        them as a sorted list of tuples containing (fileID, m_RootOrder).
        """
        root_objects = []

        # Find root GameObjects
        for go_id, transform_id in self.gameobject_to_transform.items():
            if go_id not in self.all_child_ids:
                transform_info = self.object_map.get(transform_id, {}).get('Transform', {})
                root_order = transform_info.get('m_RootOrder', 999999)
                root_objects.append((go_id, root_order))

        # Find root PrefabInstances
        for prefab_id, obj_data in self.object_map.items():
            if 'PrefabInstance' not in obj_data:
                continue
            
            if prefab_id not in self.all_child_ids:
                modifications = obj_data['PrefabInstance'].get('m_Modification', {})
                root_order = 999999
                for mod in modifications.get('m_Modifications', []):
                    if mod.get('propertyPath') == 'm_RootOrder':
                        root_order = mod.get('value', 999999)
                        break
                root_objects.append((prefab_id, root_order))
        
        # Sort all root objects by their m_RootOrder
        root_objects.sort(key=lambda x: x[1])
        
        return root_objects