Unity DOTS中的baking（五）prefabs

在DOTS的baking过程中，prefabs会被烘焙成entity prefabs。entity prefabs也是一个entity，可以在运行时实例化，就像是prefab一样。我们可以使用EntityPrefabReference这个struct，将prefab注册到baker中。

首先，我们需要定义一个包含EntityPrefabReference的ECS component：

public struct Config : IComponentData
{
    public EntityPrefabReference PrefabReference;
}
1
2
3
4

然后，再定义一个authoring component，用来给entity添加Config这个component：

public class ConfigAuthoring : MonoBehaviour
{
    public GameObject Prefab;

    class Baker : Baker
    {
        public override void Bake(ConfigAuthoring authoring)
        {
            var prefabEntity = new EntityPrefabReference(authoring.Prefab);

            var entity = GetEntity(TransformUsageFlags.None);
            AddComponent(entity, new Config
            {
                PrefabReference = prefabEntity,
            });
        }
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

此时我们使用传入的prefab新建了一个EntityPrefabReference类型的对象。自此，prefab的baking过程就结束了。接下来，就是要如何加载这个entity prefab了。

要加载 EntityPrefabReference 引用的prefab，还需要将 RequestEntityPrefabLoaded 这个component添加到entity中。当prefab加载结束时，Unity会将PrefabLoadResult这个component添加到包含RequestEntityPrefabLoaded的同一entity中。我们可以在统一的一个system中完成这项工作：

public partial struct LoadPrefabSystem : ISystem
{
    public void OnCreate(ref SystemState state)
    {
        state.RequireForUpdate();
    }

    [BurstCompile]
    public void OnUpdate(ref SystemState state)
    {
        state.Enabled = false;

        var config = SystemAPI.GetSingleton();
        var configEntity = SystemAPI.GetSingletonEntity();

        state.EntityManager.AddComponentData(configEntity, new RequestEntityPrefabLoaded
        {
            Prefab = config.PrefabReference
        });
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

最后，在使用这个prefab的system中，只有当获取到PrefabLoadResult这个component时，system才会执行：

    public partial struct SpawnSystem : ISystem
    {
        public void OnCreate(ref SystemState state)
        {
            state.RequireForUpdate();
            state.RequireForUpdate();
        }

        [BurstCompile]
        public void OnUpdate(ref SystemState state)
        {
            var config = SystemAPI.GetSingleton();

            var configEntity = SystemAPI.GetSingletonEntity();
            if (!SystemAPI.HasComponent(configEntity))
            {
                return;
            }

            var prefabLoadResult = SystemAPI.GetComponent(configEntity);
            var entity = state.EntityManager.Instantiate(prefabLoadResult.PrefabRoot);
            var random = Random.CreateFromIndex((uint) state.GlobalSystemVersion);
            state.EntityManager.SetComponentData(entity,
                LocalTransform.FromPosition(random.NextFloat(-5, 5), random.NextFloat(-5, 5), 0));
        }
    }

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

同样地，我们来窥探一下Unity内部的实现机制。首先是构造EntityPrefabReference时，调用的构造函数为：

public EntityPrefabReference(GameObject prefab) : this(
    AssetDatabase.GUIDFromAssetPath(AssetDatabase.GetAssetPath(prefab)))
{
}
1
2
3
4

可以看到，这里保存了prefab对应的guid，后续是通过guid来进行prefab加载的。我们再来看看请求加载prefab时使用的RequestEntityPrefabLoaded，它只包含了EntityPrefabReference一个成员：

/// 

/// Component to signal the request to convert the referenced prefab.
/// 
public struct RequestEntityPrefabLoaded : IComponentData
{
    /// 

    /// The reference of the prefab to be converted.
    /// 
    public EntityPrefabReference Prefab;
}
1
2
3
4
5
6
7
8
9
10

容易猜到，Unity内部有一个System，会遍历所有包含RequestEntityPrefabLoaded的Entity。它就是WeakAssetReferenceLoadingSystem，这个System创建时拥有一个WeakAssetReferenceLoadingData类型的Component：

struct WeakAssetReferenceLoadingData : IComponentData
{
    public struct LoadedPrefab
    {
        public int RefCount;
        public Entity SceneEntity;
        public Entity PrefabRoot;
    }

    public NativeParallelMultiHashMap InProgressLoads;
    public NativeParallelHashMap LoadedPrefabs;
}
1
2
3
4
5
6
7
8
9
10
11
12

InProgressLoads是一个一对多的HashMap，它记录了当前有哪些Entity发起了加载同一个prefab的请求。LoadedPrefabs则是一对一的HashMap，它保存了当前已经加载完成的prefab信息，包含引用计数，prefab对应的Entity。这个Component在System的OnCreate方法执行时就被添加到System上，直到System销毁时在OnDestroy方法中移除。

System的OnUpdate方法才是重头戏。它负责prefab引用计数的管理，加载和卸载prefab的操作。

[BurstCompile]
public void OnUpdate(ref SystemState state)
{
    var ecb = new EntityCommandBuffer(Allocator.Temp);
    foreach(var (req, e) in SystemAPI.Query().WithEntityAccess().WithNone())
    {
        StartLoadRequest(ref state, e, req.Prefab, ref ecb);
    }
    ecb.Playback(state.EntityManager);

    var sceneEntitiesToUnload = new NativeList(Allocator.Temp);
    var prefabIdsToRemove = new NativeList(Allocator.Temp);
    var inProgressLoadsToRemove = new NativeParallelHashMap(16, Allocator.Temp);
    var PARSToRemove = new NativeList(Allocator.Temp);
    ref var loadingData = ref state.EntityManager.GetComponentDataRW(state.SystemHandle).ValueRW;
    foreach (var (prefabReference, e) in SystemAPI.Query().WithEntityAccess().WithNone())
    {
        if (loadingData.LoadedPrefabs.TryGetValue(prefabReference._ReferenceId, out var loadedPrefab))
        {
            if (loadedPrefab.RefCount > 1)
            {
                loadedPrefab.RefCount--;
                loadingData.LoadedPrefabs[prefabReference._ReferenceId] = loadedPrefab;
            }
            else
            {
                prefabIdsToRemove.Add(prefabReference._ReferenceId);

                sceneEntitiesToUnload.Add(loadedPrefab.SceneEntity);
            }
        }
        else
        {
            inProgressLoadsToRemove.Add(prefabReference._ReferenceId, e);
        }
        PARSToRemove.Add(e);
    }

    for (int i = 0; i < sceneEntitiesToUnload.Length; i++)
    {
        SceneSystem.UnloadScene(state.WorldUnmanaged, sceneEntitiesToUnload[i], SceneSystem.UnloadParameters.DestroyMetaEntities);
    }

    for (int i = 0; i < prefabIdsToRemove.Length; i++)
    {
        loadingData.LoadedPrefabs.Remove(prefabIdsToRemove[i]);
    }

    foreach (var k in inProgressLoadsToRemove.GetKeyArray(Allocator.Temp))
    {
        loadingData.InProgressLoads.Remove(k, inProgressLoadsToRemove[k]);
    }

    for (int i = 0; i < PARSToRemove.Length; i++)
    {
        state.EntityManager.RemoveComponent(PARSToRemove[i]);
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

方法稍微有点长，但可以划分为几个部分。第一块是查询带有RequestEntityPrefabLoaded但是又没有PrefabAssetReference的Entity，前者表示请求加载prefab，后者表示该请求System是否已经在处理中。这个Component可以避免prefab的重复加载。

对于满足条件的query，则会调用StartLoadRequest方法：

void StartLoadRequest(ref SystemState state, Entity entity, EntityPrefabReference loadRequestWeakReferenceId, ref EntityCommandBuffer ecb)
{
    ref var loadingData = ref state.EntityManager.GetComponentDataRW(state.SystemHandle).ValueRW;
    if (loadingData.LoadedPrefabs.TryGetValue(loadRequestWeakReferenceId, out var loadedPrefab))
    {
        ecb.AddComponent(entity, new PrefabAssetReference { _ReferenceId = loadRequestWeakReferenceId});
        ecb.AddComponent(entity, new PrefabLoadResult { PrefabRoot = loadedPrefab.PrefabRoot});
        loadedPrefab.RefCount++;
        loadingData.LoadedPrefabs.Remove(loadRequestWeakReferenceId);
        loadingData.LoadedPrefabs.Add(loadRequestWeakReferenceId, loadedPrefab);
        return;
    }

    ecb.AddComponent(entity, new PrefabAssetReference { _ReferenceId = loadRequestWeakReferenceId });

    if (!loadingData.InProgressLoads.ContainsKey(loadRequestWeakReferenceId))
    {
        var loadParameters = new SceneSystem.LoadParameters { Flags = SceneLoadFlags.NewInstance };
        var sceneEntity = SceneSystem.LoadPrefabAsync(state.WorldUnmanaged, loadRequestWeakReferenceId, loadParameters);
        ecb.AddComponent(sceneEntity, new WeakAssetPrefabLoadRequest { WeakReferenceId = loadRequestWeakReferenceId });
    }

    loadingData.InProgressLoads.Add(loadRequestWeakReferenceId, entity);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

方法首先检查prefab是否已经加载完成，如果完成了则直接往发起请求的Entity上添加PrefabAssetReference和PrefabLoadResult这两个Component，前者表示System已经处理了，后者表示prefab加载完成的结果。然后还要更新下prefab的引用计数。

如果还未加载完成，则要判断是已经在加载中，还是还没发起加载请求。如果还不在加载中，则需要调用SceneSystem.LoadPrefabAsync发起加载请求，该方法返回一个表示prefab当前加载状态的Entity，如果prefab加载完成，Entity上会添加一个PrefabRoot类型的Component，它包含了加载完成的prefab。

/// 

/// Load a prefab by its weak reference id.
/// A PrefabRoot component is added to the returned entity when the load completes.
/// 
/// The  in which the prefab is loaded.
/// The weak asset reference to the prefab.
/// The load parameters for the prefab.
/// An entity representing the loading state of the prefab.
public static Entity LoadPrefabAsync(WorldUnmanaged world, EntityPrefabReference prefabReferenceId, LoadParameters parameters = default)
{
    return LoadSceneAsync(world, prefabReferenceId.Id.GlobalId.AssetGUID, parameters);
}
1
2
3
4
5
6
7
8
9
10
11
12

发起加载请求后，还要为Scene Entity添加WeakAssetPrefabLoadRequest类型的Component。这个component的作用是为了让prefab加载完成时通知到当前System，类似注册了一个回调函数。只有带有该Component的Scene Entity完成prefab加载时，才会调用System的CompleteLoad方法：

if (prefabRoot != Entity.Null)
{
    var sceneEntity = EntityManager.GetComponentData(sectionEntity).SceneEntity;
    EntityManager.AddComponentData(sceneEntity, new PrefabRoot {Root = prefabRoot});
    if (EntityManager.HasComponent(sceneEntity))
    {
        WeakAssetReferenceLoadingSystem.CompleteLoad(ref systemState,
            sceneEntity,
            prefabRoot,
            EntityManager.GetComponentData(sceneEntity)
                .WeakReferenceId);
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13

CompleteLoad方法更新System的WeakAssetReferenceLoadingData信息，所有请求加载该prefab的Entity都会添加PrefabAssetReference和PrefabLoadResult这两个Component，并且Entity的数量就是该prefab的引用计数。如果没有这样的Entity，则直接卸载这个prefab。

/// 

/// Marks a prefab as loaded and cleans up the in progress state.
/// 
/// The entity system state.
/// The entity representing the loading state of the scene.
/// The root entity of a converted prefab.
/// The prefab reference used to initiate the load.
public static void CompleteLoad(ref SystemState state, Entity sceneEntity, Entity prefabRoot, EntityPrefabReference weakReferenceId)
{
    ref var loadingData = ref state.EntityManager.GetComponentDataRW(state.WorldUnmanaged.GetExistingUnmanagedSystem()).ValueRW;
    if (!loadingData.InProgressLoads.TryGetFirstValue(weakReferenceId, out var entity, out var it))
    {
#if UNITY_EDITOR
        if (loadingData.LoadedPrefabs.TryGetValue(weakReferenceId, out var loadedPrefab))
        {
            // Prefab was reloaded, patch all references to point to the new prefab root
            loadedPrefab.PrefabRoot = prefabRoot;
            loadedPrefab.SceneEntity = sceneEntity;
            loadingData.LoadedPrefabs[weakReferenceId] = loadedPrefab;
            var query = state.GetEntityQuery(ComponentType.ReadOnly());
            foreach (var req in query.ToComponentDataArray(Allocator.Temp))
            {
                if (req.Prefab == weakReferenceId)
                    loadedPrefab.PrefabRoot = prefabRoot;
            }
            return;
        }
#endif
        //No one was waiting for this load, unload it immediately
        SceneSystem.UnloadScene(state.WorldUnmanaged, sceneEntity, SceneSystem.UnloadParameters.DestroyMetaEntities);
        return;
    }

    state.EntityManager.AddComponentData(prefabRoot, new RequestEntityPrefabLoaded() {Prefab =  weakReferenceId});
    int count = 0;
    do
    {
        state.EntityManager.AddComponentData(entity, new PrefabAssetReference { _ReferenceId = weakReferenceId});
        state.EntityManager.AddComponentData(entity, new PrefabLoadResult { PrefabRoot =  prefabRoot});
        count++;
    } while (loadingData.InProgressLoads.TryGetNextValue(out entity, ref it));
    loadingData.InProgressLoads.Remove(weakReferenceId);
    loadingData.LoadedPrefabs.Add(weakReferenceId, new WeakAssetReferenceLoadingData.LoadedPrefab { SceneEntity = sceneEntity, PrefabRoot = prefabRoot, RefCount = count});
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

看完加载的部分，再来看看卸载相关的逻辑。我们再次回到OnUpdate方法。接下来查询所有包含PrefabAssetReference并且不包含RequestEntityPrefabLoaded的Entity，这与前面恰好相反，它表示所有需要取消加载prefab的Entity。如果prefab已经加载完成了，此时需要将其引用计数减一；如果引用计数为0了，则需要将其卸载。如果prefab还在加载中，则需要从请求加载的hashmap中把Entity移除。最后，为了避免重复处理这类取消加载的Entity，处理完毕后需要移除它们身上的PrefabAssetReference。

当然，实际上另一种情形更常见，就是Entity在其生命周期内请求加载了某个prefab之后一直在使用，直到它被destroy。这种情况外部并不需要做额外的操作，Entity所持有的prefab引用计数也会自动正确地更新。这是怎么做到的呢？

答案就在RequestEntityPrefabLoaded和PrefabAssetReference这两个Component上。默认Entity是同时拥有这两个Component的，但是PrefabAssetReference是ICleanupComponentData，意味着Entity销毁时它也不会被销毁，这样被销毁的Entity身上，只有一个PrefabAssetReference，就会在System的OnUpdate时触发prefab的清理工作。

Reference

[1] Prefabs in baking