ObjectPoolComponent.h

// Copyright (c) Tony Nguyen. All rights reserved.

// Note: Some code blocks are intentionally omitted.

#pragma once

using FPoolObjectCreationFunc = TFunction<TObjectPtr<UObject>()>;

USTRUCT()
struct FObjectPool
{
	GENERATED_BODY()

public:
	FObjectPool() { }

	/** Constructor that takes in a capacity parameter. Internally, the buffer leaves 1 item to detect full/empty states, so we add 1 to make up for it. */
	FObjectPool(const uint32 InCapacity, FPoolObjectCreationFunc* InCreationFunc)
		: Capacity(InCapacity)
		, CreationFunc(MakeShared<FPoolObjectCreationFunc>(*InCreationFunc))
	{
		// Initialize with set capacity. This cannot be changed.
		Objects.Reserve(InCapacity);
		AvailableObjects = MakeShared<TCircularQueue<TObjectPtr<UObject>>>(InCapacity + 1);

		Replenish();
	}

	/** Adds object to pool if pool is not at capacity. Do not use this for returning objects to pool. */
	bool AddObject(const TObjectPtr<UObject>& ObjectToAdd)
	{
		if (!ObjectToAdd || IsFull()) return false;

		const int32 AvailableIndex = Objects.Find(nullptr);

		if (AvailableIndex != INDEX_NONE)
		{
			Objects[AvailableIndex] = ObjectToAdd;
			AvailableObjects->Enqueue(ObjectToAdd);

			return true;
		}

		return false;
	}

	/** Returns the next available object from the pool. */
	TObjectPtr<UObject> AcquireObject()
	{
		TObjectPtr<UObject> Obj = nullptr;

		if (AvailableObjects->Dequeue(Obj))
		{
			if (IsValid(Obj))
			{
				return Obj;
			}
		}

		Replenish();
		return Obj;
	}

	/** Returns the object to the pool. */
	bool ReturnObject(const TObjectPtr<UObject>& ReturningObject) const
	{
		if (!IsValid(ReturningObject)) return false;

		return AvailableObjects->Enqueue(ReturningObject);
	}
	
private:
	/** Returns true if pool is full. */
	bool IsFull() const
	{
		// If index of nullptr cannot be found, this means the pool is completely filled.
		return Objects.Find(nullptr) == INDEX_NONE;
	}

	/** Fills the remaining space in the object pool with new objects. */
	void Replenish()
	{
		Clean();

		if (CreationFunc)
		{
			const int32 NextAvailableIndex = Objects.Find(nullptr);
			const int32 NumObjectsToCreate = Objects.Max() - (NextAvailableIndex != INDEX_NONE ? NextAvailableIndex : Objects.Max());

			for (int32 i = 0; i < NumObjectsToCreate; ++i)
			{
				auto Obj = (*CreationFunc)();
				AddObject(Obj);
			}
		}
	}

	/** Cleans the object pool. */
	void Clean()
	{
		Objects.RemoveAll([](const TObjectPtr<UObject>& Object)
		{
			return !IsValid(Object);
		});

		Objects.SetNum(Capacity);
	}

private:
	TArray<TObjectPtr<UObject>> Objects;

	TSharedPtr<TCircularQueue<TObjectPtr<UObject>>> AvailableObjects;

	/**
	 * The pool's buffer capacity.
	 * This will include the extra 1 item in the buffer that is used to detect full/empty states.
	 * So, this will appear 1 more than what the user inputs.
	 */
	uint32 Capacity = 0;

	TSharedPtr<FPoolObjectCreationFunc> CreationFunc = nullptr;
};

/*
 * The component that manages and stores the pooled objects.
 */
UCLASS(Blueprintable, ClassGroup=(Custom), meta=(BlueprintSpawnableComponent))
class OBJECTPOOLER_API UObjectPoolComponent : public UActorComponent
{
	GENERATED_BODY()

public:
	UObjectPoolComponent();

	virtual void BeginPlay() override;

	/**
	 * Acquires an object from the pool. Returns true if successful.
	 * @param OutObject The object acquired from the pool, if valid.
	 * @return True if acquiring the object was successful.
	 */
	UFUNCTION(BlueprintCallable, Category = "Object Pool", meta=(ExpandBoolAsExecs="ReturnValue"))
	bool AcquireObject(UObject*& OutObject);

	/**
	 * Returns an object to the pool.
	 * @param ObjectToReturn The object to return to the pool.
	 * @return True if return is successful.
	 */
	UFUNCTION(BlueprintCallable, Category = "Object Pool")
	bool ReturnObject(UObject* ObjectToReturn);

	/** Called when a pool object is created. Useful for object initialization. */
	UFUNCTION(BlueprintNativeEvent, Category = "Object Pool")
	void OnCreateObject(const UObject* InObject);
	virtual void OnCreateObject_Implementation(const UObject* InObject);

	/** Called when a pool object is acquired. */
	UFUNCTION(BlueprintNativeEvent, Category = "Object Pool")
	void OnAcquireObject(const UObject* InObject);
	virtual void OnAcquireObject_Implementation(const UObject* InObject);

	/** Called when an object returns to the pool. Any object re-initialization should be done here. */
	UFUNCTION(BlueprintNativeEvent, Category = "Object Pool")
	void OnReturnObject(const UObject* InObject);
	virtual void OnReturnObject_Implementation(const UObject* InObject);

	/** Called when an object fails to return to the pool. */
	UFUNCTION(BlueprintNativeEvent, Category = "Object Pool")
	void OnRejectObject(const UObject* InObject);
	virtual void OnRejectObject_Implementation(const UObject* InObject);

	bool IsPoolValid() const { return Pool.IsValid(); }

private:
	/** Initialize the object pool. Can only call this once successfully. */
	void InitializePool();

	/** Creates a new object to pooled. */
	TObjectPtr<UObject> CreateObject();

public:
	/** The class of the object to pool. */
	UPROPERTY(EditDefaultsOnly, BlueprintReadOnly)
	TSubclassOf<UObject> PoolObjectClass;

	/** The limit for the number of pooled objects. */
	UPROPERTY(EditDefaultsOnly, BlueprintReadOnly, meta=(ClampMin=0, UIMin=0))
	int32 PoolCapacity = 0;

private:
	TSharedPtr<FObjectPool> Pool;
};