EnTT Snapshot Creation & Restoration Unit Test

SnapshotTests.h

#include "pch.h" // #include <gtest/gtest.h>

#include "TestGameManager.h"

namespace SnapshotTests {
    class SnapshotTests : public ::testing::Test {
    protected:
        entt::registry registry;
        TestGameManager toTest;

        SnapshotTests() : toTest(registry) {}

        entt::entity getTrackedEntityFromSpawnComponent() {
            entt::entity result = entt::null;

            bool componentFound = false; // Only one spawn component expected
            auto view = registry.view<TestSpawnComponent>();
            for (auto &&[entity, testComp] : view.each()) {
                if (componentFound) {
                    ADD_FAILURE(); // "Only one spawn component expected"
                    return entt::null;
                }

                componentFound = true;
                result = testComp.trackedEntity;
            }

            if (!componentFound) {
                ADD_FAILURE(); // "No spawn component found"
            }
            else if (result == entt::null) {
                ADD_FAILURE(); // "Tracked entity is still null"
            }
            
            return result;
        }

        void assertThatEntityIsValid(entt::entity entityId) {
            ASSERT_TRUE(registry.valid(entityId));
        }
    };
    
    TEST_F(SnapshotTests, whenSnapshotRestored_thenRegistryDoesNotMarkExistingEntityAsInvalid) {
        // Initial data population:
        registry.emplace<TestSpawnComponent>(registry.create());
        toTest.update();
        entt::entity initialTrackedEntity = getTrackedEntityFromSpawnComponent();

        // Sanity check: Verify registry believes that the tracked entity is still valid
        assertThatEntityIsValid(initialTrackedEntity);

        // Sanity check #2: Call update several times to verify spawn system behavior
        toTest.update();
        toTest.update();
        toTest.update();
        ASSERT_EQ(initialTrackedEntity, getTrackedEntityFromSpawnComponent());
        
        // Actual test: Create and restore snapshot
        ASSERT_TRUE(toTest.createSnapshot());
        ASSERT_TRUE(toTest.restoreSnapshot());
        assertThatEntityIsValid(initialTrackedEntity); // Quick and easy verification
        
        // Post-restoration verification
        toTest.update(); // Spawn system is expected to NOT change tracked entity as registry should not treat entity as invalid
        assertThatEntityIsValid(initialTrackedEntity);
        entt::entity postSnapshotRestorationTrackedEntity = getTrackedEntityFromSpawnComponent();
        assertThatEntityIsValid(postSnapshotRestorationTrackedEntity);
        ASSERT_EQ(initialTrackedEntity, postSnapshotRestorationTrackedEntity);
    }

}

TestComponents.h

#pragma once

#include <EnTT/entt.hpp>

struct TestSpawnComponent {
    entt::entity trackedEntity;
};

struct SimpleComponent {
    bool someData;
};

TestGameManager.h

#pragma once

#include <EnTT/entt.hpp>

#include "TestComponents.h"
#include "TestSpawnSystem.h"

class TestGameManager {
    entt::registry& registry;
    TestSpawnSystem testSystem;

    /// Snapshot data
    bool isSnapshotStored = false;
    // ECS Snapshot: Create vector of entities and components per component class, as EnTT stores data on a per-component-class basis
    std::vector<entt::entity> entitySnapshot_TestSpawnComponent;
    std::vector<TestSpawnComponent> componentSnapshot_TestSpawnComponent;
    std::vector<entt::entity> entitySnapshot_SimpleComponent;
    std::vector<SimpleComponent> componentSnapshot_SimpleComponent;
    // Additional snapshot data for proper entity restoration
    entt::entity snapshotReleased = entt::null;
    std::size_t snapshotRegistrySize = {};
    std::vector<entt::entity> snapshotRegistryEntities;

public:
    TestGameManager(entt::registry& registry) : registry(registry) {}

    void update() {
        testSystem.update(registry);
    }

    // Returns true if succeeds
    bool createSnapshot() {
        if (isSnapshotStored) {
            return false;
        }

        snapshotRegistrySize = registry.size();
        snapshotReleased = registry.released();
        snapshotRegistryEntities.insert(snapshotRegistryEntities.end(), registry.data(),
                                        registry.data() + snapshotRegistrySize);

        // Store the actual snapshot data per each component class
        // This copy-per-class approach is necessary due to how EnTT is built
        createSnapshotForComponent<TestSpawnComponent>(entitySnapshot_TestSpawnComponent, componentSnapshot_TestSpawnComponent);
        createSnapshotForComponent<SimpleComponent>(entitySnapshot_SimpleComponent, componentSnapshot_SimpleComponent);

        isSnapshotStored = true;
        return true;
    }

    // Returns true if succeeds
    bool restoreSnapshot() {
        if (!isSnapshotStored) {
            return false;
        }

        registry = {};
        registry.assign(snapshotRegistryEntities.data(), snapshotRegistryEntities.data() + snapshotRegistrySize, snapshotReleased);

        // Restore the actual snapshot data per each component class
        // This copy-per-class approach is necessary due to how EnTT is built
        restoreSnapshotForComponent<TestSpawnComponent>(entitySnapshot_TestSpawnComponent, componentSnapshot_TestSpawnComponent);
        restoreSnapshotForComponent<SimpleComponent>(entitySnapshot_SimpleComponent, componentSnapshot_SimpleComponent);

        cleanUpSnapshotData();

        return true;
    }

private:
    template <typename ComponentType>
    void createSnapshotForComponent(std::vector<entt::entity>& entitySnapshot,
                                    std::vector<ComponentType>& componentSnapshot) {
        // Get reference to unique storage for this component type
        auto&& storage = registry.storage<ComponentType>();

        // Create entity snapshot
        // For why not using memcpy (speed) and examples of different methods for copying data into a vector: https://stackoverflow.com/a/261607/3735890
        entitySnapshot.insert(entitySnapshot.end(), storage.data(), storage.data() + storage.size());

        // Create component snapshot
        // Note that data may cross multiple pages in memory. Thus, need to copy the data into the output vector one page at a time
        const std::size_t pageSize = entt::component_traits<ComponentType>::page_size;
        const std::size_t totalPages = (storage.size() + pageSize - 1u) / pageSize;
        for (std::size_t pageIndex{}; pageIndex < totalPages; pageIndex++) {
            // Calculate number of elements to copy so only copyyig over the necessary number of elements
            // Truly necessary as using ComponentType vector here instead of, say, vector of arbitrary bytes
            const std::size_t offset = pageIndex * pageSize;
            const std::size_t numberOfElementsToCopy = std::min(pageSize, storage.size() - offset);

            // Do the actual copying
            ComponentType* pageStartPtr = storage.raw()[pageIndex];
            componentSnapshot.insert(componentSnapshot.end(), pageStartPtr, pageStartPtr + numberOfElementsToCopy);
        }
    }

    template<typename ComponentType>
    void restoreSnapshotForComponent(std::vector<entt::entity>& entitySnapshot, std::vector<ComponentType>& componentSnapshot) {
        // Get reference to unique storage for this component type
        auto&& storage = registry.storage<ComponentType>();

        // Restore entities
        storage.insert(entitySnapshot.begin(), entitySnapshot.end());

        // Restore components
        // Note that data may cross multiple pages in memory. Thus, need to copy the data into the storage one page at a time
        const std::size_t pageSize = entt::component_traits<ComponentType>::page_size;
        const std::size_t totalPages = (storage.size() + pageSize - 1u) / pageSize;
        for(std::size_t pageIndex{}; pageIndex < totalPages; pageIndex++) {
            const std::size_t offset = pageIndex * pageSize;
            const std::size_t numberOfElementsToCopy = std::min(pageSize, componentSnapshot.size() - offset);

            ComponentType* pageStartPtr = storage.raw()[pageIndex];
            memcpy(pageStartPtr, componentSnapshot.data() + offset, sizeof(ComponentType) * numberOfElementsToCopy);
        }
    }
        
    void cleanUpSnapshotData() {
        snapshotRegistryEntities.clear();
        
        // Clear all existing snapshot data
        cleanUpSnapshotData<TestSpawnComponent>(entitySnapshot_TestSpawnComponent, componentSnapshot_TestSpawnComponent);
        cleanUpSnapshotData<SimpleComponent>(entitySnapshot_SimpleComponent, componentSnapshot_SimpleComponent);
    }
        
    template<typename ComponentType>
    void cleanUpSnapshotData(std::vector<entt::entity>& entitySnapshot, std::vector<ComponentType>& componentSnapshot) {
        entitySnapshot.clear();
        componentSnapshot.clear();
    }
};

TestSpawnSystem.h

#pragma once

#include <EnTT/entt.hpp>

#include "TestComponents.h"

class TestSpawnSystem {
public:
    void update(entt::registry& registry) {
        auto view = registry.view<TestSpawnComponent>();
        for (auto &&[entity, testComp] : view.each()) {
            removeInvalidEntity(registry, testComp);
            spawnEntityIfNecessary(registry, testComp);
        }
    }

private:
    void removeInvalidEntity(entt::registry& registry, TestSpawnComponent& testComp) {
        if (testComp.trackedEntity != entt::null && !registry.valid(testComp.trackedEntity)) {
            testComp.trackedEntity = entt::null;
        }
    }

    void spawnEntityIfNecessary(entt::registry& registry, TestSpawnComponent& testComp) {
        if (testComp.trackedEntity != entt::null) {
            return;
        }

        entt::entity newEntity = createEntity(registry);
        testComp.trackedEntity = newEntity;
    }

    entt::entity createEntity(entt::registry& registry) {
        auto entityId = registry.create();
        registry.emplace<SimpleComponent>(entityId);
        return entityId;
    }
};