Undo/Redo transaction issues on C++ Spline Actor with component pool

I am trying to create a C++ Actor representing a Spline that gets spline components from a component pool while it is being adjusted in the Editor. The goal of the component pool is that components are created once, preventing the continuous creation and destruction of new spline components while changing the shape and length of the spline in the Editor.

The basic principle is simple, when the spline gets longer, check if there are components available in the pool. If so, get a component from the pool. Otherwise, add extra components to the pool and return the first of those. When the spline gets shorter, put the unnecessary components back into the pool.

This principle works simple enough, but I am getting into trouble when trying to make it play nice with the Editor Undo/Redo functionality. Especially when I Undo the Deletion of the Actor. In that case, the recreated Actor has many broken references in the pool and as children of the spline.

The Spline (to which all components are attached as children) still holds references to its children, but the children do not have a reference to the Spline anymore. This makes it impossible to call DetachFromComponent(), sine this is called on the child, which is unaware of its attachment. In order to fix this, I need to override the PostEditUndo() function of the Actor. There I first have to call ClearPendingKill() to be able to change the component’s attachment, then attach it again to the Spline, overwriting the original attachment. Then I can Detach and Destroy the component to clean it up. Finally I need to manually call MarkPendingKill() to ensure it can be distinguished from newly created components later on.

I added my code below. The USplineMeshComponentOver::TryUnregisterAndDestroy() function handles cleaning the Spline children. To prevent loosing references through transactions a TArray of Active components is required next to the pool, resulting in two reference for each component in use. This also adds the requirement to clean up references both in the pool and in the active component array. Another issue is that Undo will Register all components in my pool, spawning them at their last known locations. Even when they were pooled when the Actor was destroyed. I seem to need to manually UnRegister() all pooled components from the PostEditUndo() function.

While my code seems to be working now, I genuinely feel like this is not supposed to be done like this. Especially the point where I need to play with the kill flag to manually repair component references. I hope anyone can help me out here and tell me if there is a better way to do this.

Here follows my code. ATrackGenerator is the Actor holding the Spline. UActorComponentPool is the pool of components and USplineMeshComponentOver represents the components that are pooled and should be added to the spline.

TrackGenerator.h

#pragma once

#include "CoreMinimal.h"
#include "GameFramework/Actor.h"
#include "Components/SplineMeshComponent.h"
#include "TrackGenerator.generated.h"

class USplineComponent;
class USplineMeshComponentOver;
class UStaticMesh;

UCLASS()
class UActorComponentPool : public UActorComponent
{
	GENERATED_BODY()
public:
	void Init(UStaticMesh* mesh);
	USplineMeshComponentOver* GetFromPool(UPrimitiveComponent* Parent, UMaterialInterface* material);
	void PutInPool(USplineMeshComponentOver* component);

	void ResetNextComponent(UPrimitiveComponent* Parent, int32 NumOfActiveComponents);

	bool SanitizePool(UPrimitiveComponent* Parent);

private:
	void AddComponentsToPool(int32 Num);

	UPROPERTY()
	TArray<USplineMeshComponentOver*> pool;

	UPROPERTY()
	UStaticMesh* Mesh;

	UPROPERTY()
	int32 NextComponent = 0;
};

UCLASS()
class USplineMeshComponentOver : public USplineMeshComponent
{
	GENERATED_BODY()
public:
	void SetPool(UActorComponentPool* pool);
	void ReturnToPool();
	void OnReturnToPool();
	void OnGetFromPool(UPrimitiveComponent* Parent);
	void UnRegister();
	bool TryUnregisterAndDestroy(UPrimitiveComponent* Parent);

	static USplineMeshComponentOver* Create(UActorComponentPool* pool, UStaticMesh* mesh);

private:
	UPROPERTY()
	UActorComponentPool* Pool;

	UPROPERTY()
	bool bInPool = false;
};

UCLASS()
class RACETRACK_API ATrackGenerator : public AActor
{
	GENERATED_BODY()
	
public:
	ATrackGenerator();

#if WITH_EDITOR
	virtual void PostEditMove(bool bFinished) override;
	virtual void PostEditUndo() override;
#endif

protected:
	virtual void OnConstruction(const FTransform& Transform) override;
	virtual void PostActorCreated() override;

private:
	void SetupSpline();

	UPROPERTY(VisibleAnywhere, Category = "Spline")
	USplineComponent* Spline;

	UPROPERTY(EditAnywhere, Category = "Spline")
	UStaticMesh* StraightMesh;

	UPROPERTY(EditAnywhere, Category = "Spline")
	UMaterialInterface* CurbMaterial;

	UPROPERTY()
	UActorComponentPool* SplineMeshComponentPool;
	
	UPROPERTY()
	TArray<USplineMeshComponentOver*> ActiveComponents;
};

TrackGenerator.cpp

#include "TrackGenerator.h"
#include "Components/SplineComponent.h"
#include "Engine/StaticMesh.h"
#include "Components/SplineMeshComponent.h"

ATrackGenerator::ATrackGenerator() : Super()
{
	Spline = CreateDefaultSubobject<USplineComponent>("Spline");
	if (Spline)
	{
		SetRootComponent(Spline);
		Spline->SetMobility(EComponentMobility::Static);
	}
	
	SplineMeshComponentPool = CreateDefaultSubobject<UActorComponentPool>("ComponentPool");
}

void ATrackGenerator::PostActorCreated()
{
	SplineMeshComponentPool->Init(StraightMesh);
}

#if WITH_EDITOR
void ATrackGenerator::PostEditMove(bool bFinished)
{
	if (!bFinished)
	{
		if (ReregisterComponentsWhenModified())
		{
			SetupSpline();
		}
		Super::PostEditMove(bFinished);
	}
}

void ATrackGenerator::PostEditUndo()
{
	if (SplineMeshComponentPool != nullptr)
	{
		SplineMeshComponentPool->SanitizePool(Spline);

		bool SanitizedActiveComponents = false;
		for (int32 i = 0; i < ActiveComponents.Num(); ++i)
		{
			if (ActiveComponents[i]->TryUnregisterAndDestroy(Spline))
			{
				SanitizedActiveComponents = true;
			}
		}
		if (SanitizedActiveComponents)
		{
			ActiveComponents.Empty();
			SplineMeshComponentPool->ResetNextComponent(Spline, ActiveComponents.Num());
		}
	}
	
	Super::PostEditUndo();
	
	if (SplineMeshComponentPool != nullptr)
	{
		SplineMeshComponentPool->ResetNextComponent(Spline, ActiveComponents.Num());
	}
}
#endif

void ATrackGenerator::OnConstruction(const FTransform& Transform)
{
	SetupSpline();
}

void ATrackGenerator::SetupSpline()
{
	if (StraightMesh)
	{
		float SegmentLength = StraightMesh->GetBoundingBox().GetSize().X;
		float SplineLength = Spline->GetSplineLength();

		int32 NumberOfTrackSegments = static_cast<int32>(SplineLength / SegmentLength) + 1;
		int32 NumberOfAddedMeshComponents = ActiveComponents.Num();
		
		for (int32 Segment = NumberOfTrackSegments; Segment < NumberOfAddedMeshComponents; ++Segment)
		{
			ActiveComponents.Pop(true)->ReturnToPool();
		}
		
		for (int32 Segment = NumberOfAddedMeshComponents; Segment < NumberOfTrackSegments; ++Segment)
		{
			USplineMeshComponentOver* component = SplineMeshComponentPool->GetFromPool(Spline, CurbMaterial);
			ActiveComponents.Push(component);
		}

		for (int32 Segment = 0; Segment < NumberOfTrackSegments; ++Segment)
		{
			FVector MeshStartPoint = Spline->GetLocationAtDistanceAlongSpline(Segment * SegmentLength, ESplineCoordinateSpace::Type::Local);
			FVector MeshEndPoint = Spline->GetLocationAtDistanceAlongSpline((Segment + 1) * SegmentLength, ESplineCoordinateSpace::Type::Local);
			FVector MeshStartTangent = Spline->GetTangentAtDistanceAlongSpline(Segment * SegmentLength, ESplineCoordinateSpace::Type::Local);
			FVector MeshEndTangent = Spline->GetTangentAtDistanceAlongSpline((Segment + 1) * SegmentLength, ESplineCoordinateSpace::Type::Local);

			ActiveComponents[Segment]->SetStartAndEnd(MeshStartPoint, MeshStartTangent, MeshEndPoint, MeshEndTangent, true);
		}
	}
}

void USplineMeshComponentOver::ReturnToPool()
{
	if (Pool)
	{
		Pool->PutInPool(this);
	}
}

void USplineMeshComponentOver::UnRegister()
{
	DetachFromComponent(FDetachmentTransformRules::KeepWorldTransform);
	if(IsRegistered())
	{
		UnregisterComponent();
	}
}

void USplineMeshComponentOver::OnReturnToPool()
{
	if (!bInPool)
	{
		UnRegister();
		bInPool = true;
	}
}

void USplineMeshComponentOver::OnGetFromPool(UPrimitiveComponent* Parent)
{
	if (bInPool)
	{
		RegisterComponent();
		CreationMethod = EComponentCreationMethod::Instance;
		SetMobility(EComponentMobility::Static);
		ResetRelativeTransform();
		AttachToComponent(Parent, FAttachmentTransformRules::KeepRelativeTransform);
		bInPool = false;
	}
}

bool USplineMeshComponentOver::TryUnregisterAndDestroy(UPrimitiveComponent* Parent)
{
	if (IsPendingKill())
	{
		ClearPendingKill();
		AttachToComponent(Parent, FAttachmentTransformRules::KeepRelativeTransform);
		DetachFromComponent(FDetachmentTransformRules::KeepWorldTransform);
		DestroyComponent();
		MarkPendingKill();

		return true;
	}

	return false;
}

void USplineMeshComponentOver::SetPool(UActorComponentPool* pool)
{
	Pool = pool;
}

USplineMeshComponentOver* USplineMeshComponentOver::Create(UActorComponentPool* InPool, UStaticMesh* InMesh)
{
	USplineMeshComponentOver* component = NewObject<USplineMeshComponentOver>(InPool, USplineMeshComponentOver::StaticClass());
	component->SetStaticMesh(InMesh);
	component->SetCollisionEnabled(ECollisionEnabled::Type::QueryAndPhysics);
	component->SetPool(InPool);
	component->bInPool = true;

	return component;
}


void UActorComponentPool::Init(UStaticMesh* mesh)
{
	Mesh = mesh;
	AddComponentsToPool(1);
}

void UActorComponentPool::AddComponentsToPool(int32 Num)
{
	for (int i = 0; i < Num; ++i)
	{
		USplineMeshComponentOver* component = USplineMeshComponentOver::Create(this, Mesh);
		pool.Push(component);
	}
}

USplineMeshComponentOver* UActorComponentPool::GetFromPool(UPrimitiveComponent* Parent, UMaterialInterface* material)
{
	if (pool.Num() == NextComponent)
	{
		AddComponentsToPool(1);
	}

	USplineMeshComponentOver* component = pool[NextComponent++];
	component->SetMaterial(0, material);
	component->OnGetFromPool(Parent);
	
	return component;
}

void UActorComponentPool::PutInPool(USplineMeshComponentOver* component)
{
	if (component == nullptr)
	{
		if (NextComponent == 0)
		{
			return;
		}
		component = pool[NextComponent - 1];
	}

	--NextComponent;
	component->OnReturnToPool();
}

void UActorComponentPool::ResetNextComponent(UPrimitiveComponent* Parent, int32 NumOfActiveComponents)
{
	for (int32 component = NumOfActiveComponents; component < pool.Num(); ++component)
	{
		pool[component]->UnRegister();
	}
	NextComponent = NumOfActiveComponents;
}

bool UActorComponentPool::SanitizePool(UPrimitiveComponent* Parent)
{
	int32 NumOfSanitizedComponents = 0;
	for (int32 i = 0; i < pool.Num(); ++i)
	{
		if (pool[i]->TryUnregisterAndDestroy(Parent))
		{
			pool[i] = USplineMeshComponentOver::Create(this, Mesh);

			++NumOfSanitizedComponents;
		}
	}

	NextComponent = pool.Num() - NumOfSanitizedComponents;

	return NumOfSanitizedComponents > 0;
}

Try setting this property on your component
ActorComp->CreationMethod = EComponentCreationMethod::UserConstructionScript;