<commit_before>#include "RTSPluginPCH.h"
#include "RTSProductionComponent.h"

#include "GameFramework/Actor.h"
#include "GameFramework/Controller.h"
#include "Engine/BlueprintGeneratedClass.h"
#include "Engine/SCS_Node.h"
#include "Kismet/GameplayStatics.h"

#include "RTSGameMode.h"
#include "RTSProductionCostComponent.h"
#include "RTSPlayerController.h"
#include "RTSPlayerAdvantageComponent.h"
#include "RTSPlayerResourcesComponent.h"
#include "RTSUtilities.h"

URTSProductionComponent::URTSProductionComponent(const FObjectInitializer& ObjectInitializer)
	: Super(ObjectInitializer)
{
	PrimaryComponentTick.bCanEverTick = true;

	SetIsReplicated(true);
}

void URTSProductionComponent::GetLifetimeReplicatedProps(TArray<FLifetimeProperty>& OutLifetimeProps) const
{
	Super::GetLifetimeReplicatedProps(OutLifetimeProps);

	DOREPLIFETIME(URTSProductionComponent, ProductionQueues);
}

void URTSProductionComponent::BeginPlay()
{
	UActorComponent::BeginPlay();

	// Setup queues.
	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		FRTSProductionQueue NewQueue;
		ProductionQueues.Add(NewQueue);
	}
}

void URTSProductionComponent::TickComponent(float DeltaTime, enum ELevelTick TickType, FActorComponentTickFunction *ThisTickFunction)
{
	UActorComponent::TickComponent(DeltaTime, TickType, ThisTickFunction);

    // Check for speed boosts.
    float SpeedBoostFactor = 1.0f;
    AActor* OwningActor = GetOwner();

    if (OwningActor)
    {
        AActor* OwningPlayer = OwningActor->GetOwner();

        if (OwningPlayer)
        {
            URTSPlayerAdvantageComponent* PlayerAdvantageComponent = OwningPlayer->FindComponentByClass<URTSPlayerAdvantageComponent>();

            if (PlayerAdvantageComponent)
            {
                SpeedBoostFactor = PlayerAdvantageComponent->SpeedBoostFactor;
            }
        }
    }

	// Process all queues.
	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		if (ProductionQueues[QueueIndex].Num() <= 0)
		{
			continue;
		}

		// Check production costs.
		auto ProductClass = GetCurrentProduction(QueueIndex);
		auto ProductionCostComponent = URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);

		bool bProductionCostPaid = false;

		if (ProductionCostComponent && ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayOverTime)
		{
			auto Owner = GetOwner()->GetOwner();

			if (!Owner)
			{
				UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no owner."), *GetOwner()->GetName());
				continue;
			}

            auto PlayerResourcesComponent = Owner->FindComponentByClass<URTSPlayerResourcesComponent>();

            if (!PlayerResourcesComponent)
            {
                UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no PlayerResourcesComponent."), *Owner->GetName());
                continue;
            }

			bool bCanPayAllProductionCosts = true;

			for (auto& Resource : ProductionCostComponent->Resources)
			{
				float ResourceAmount = Resource.Value * SpeedBoostFactor * DeltaTime / ProductionCostComponent->ProductionTime;

				if (!PlayerResourcesComponent->CanPayResources(Resource.Key, ResourceAmount))
				{
					// Production stopped until resources become available again.
					bCanPayAllProductionCosts = false;
					break;
				}
			}

			if (bCanPayAllProductionCosts)
			{
				// Pay production costs.
				for (auto& Resource : ProductionCostComponent->Resources)
				{
					float ResourceAmount = Resource.Value * SpeedBoostFactor * DeltaTime / ProductionCostComponent->ProductionTime;
                    PlayerResourcesComponent->PayResources(Resource.Key, ResourceAmount);
				}

				bProductionCostPaid = true;
			}
		}
		else
		{
			bProductionCostPaid = true;
		}

		if (!bProductionCostPaid)
		{
			continue;
		}

		// Update production progress.
		ProductionQueues[QueueIndex].RemainingProductionTime -= SpeedBoostFactor * DeltaTime;

		// Check if finished.
		if (ProductionQueues[QueueIndex].RemainingProductionTime <= 0)
		{
			FinishProduction(QueueIndex);
		}
	}
}

bool URTSProductionComponent::CanAssignProduction(TSubclassOf<AActor> ProductClass) const
{
	return URTSUtilities::IsReadyToUse(GetOwner()) && FindQueueForProduct(ProductClass) >= 0;
}

int32 URTSProductionComponent::FindQueueForProduct(TSubclassOf<AActor> ProductClass) const
{
	// Find queue with least products that is not at capacity limit.
	int32 QueueWithLeastProducts = -1;
	int32 ProductsInShortestQueue = CapacityPerQueue;

	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		const FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

		if (Queue.Num() < ProductsInShortestQueue)
		{
			QueueWithLeastProducts = QueueIndex;
			ProductsInShortestQueue = Queue.Num();
		}
	}

	return QueueWithLeastProducts;
}

TSubclassOf<AActor> URTSProductionComponent::GetCurrentProduction(int32 QueueIndex /*= 0*/) const
{
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return nullptr;
	}

	const FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	return Queue.Num() > 0 ? Queue[0] : nullptr;
}

float URTSProductionComponent::GetProductionTime(int32 QueueIndex /*= 0*/) const
{
	TSubclassOf<AActor> CurrentProduction = GetCurrentProduction(QueueIndex);

	if (!CurrentProduction)
	{
		return 0.0f;
	}

	return GetProductionTimeForProduct(CurrentProduction);
}

float URTSProductionComponent::GetProductionTimeForProduct(TSubclassOf<AActor> ProductClass) const
{
	URTSProductionCostComponent* ProductionCostComponent =
		URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);
	return ProductionCostComponent ? ProductionCostComponent->ProductionTime : 0.0f;
}

float URTSProductionComponent::GetProgressPercentage(int32 QueueIndex /*= 0*/) const
{
	float TotalProductionTime = GetProductionTime(QueueIndex);

	if (TotalProductionTime <= 0.0f)
	{
		return 1.0f;
	}

	float RemainingProductionTime = GetRemainingProductionTime(QueueIndex);

	if (RemainingProductionTime <= 0.0f)
	{
		return 1.0f;
	}

	return 1 - (RemainingProductionTime / TotalProductionTime);
}

float URTSProductionComponent::GetRemainingProductionTime(int32 QueueIndex /*= 0*/) const
{
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return 0.0f;
	}

	return ProductionQueues[QueueIndex].RemainingProductionTime;
}

bool URTSProductionComponent::IsProducing() const
{
	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		const FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

		if (Queue.Num() > 0)
		{
			return true;
		}
	}

	return false;
}

void URTSProductionComponent::StartProduction(TSubclassOf<AActor> ProductClass)
{
	// Check production state.
	if (!CanAssignProduction(ProductClass))
	{
		return;
	}

	int32 QueueIndex = FindQueueForProduct(ProductClass);

	if (QueueIndex < 0)
	{
		return;
	}

    // Check requirements.
    TSubclassOf<AActor> MissingRequirement;

    if (URTSUtilities::GetMissingRequirementFor(this, GetOwner(), ProductClass, MissingRequirement))
    {
        UE_LOG(LogRTS, Error, TEXT("%s wants to produce %s, but is missing requirement %s."), *GetOwner()->GetName(), *ProductClass->GetName(), *MissingRequirement->GetName());
        return;
    }

	// Check production cost.
	URTSProductionCostComponent* ProductionCostComponent =
		URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);

	if (ProductionCostComponent && ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayImmediately)
	{
		auto Owner = GetOwner()->GetOwner();

		if (!Owner)
		{
			UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no owner."), *GetOwner()->GetName());
			return;
		}

        auto PlayerResourcesComponent = Owner->FindComponentByClass<URTSPlayerResourcesComponent>();

        if (!PlayerResourcesComponent)
        {
            UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no PlayerResourcesComponent."), *Owner->GetName());
            return;
        }

		if (!PlayerResourcesComponent->CanPayAllResources(ProductionCostComponent->Resources))
		{
			UE_LOG(LogRTS, Error, TEXT("%s needs to pay for producing %s, but does not have enough resources."),
				*Owner->GetName(),
				*ProductClass->GetName());
			return;
		}

		// Pay production costs.
        PlayerResourcesComponent->PayAllResources(ProductionCostComponent->Resources);
	}

	// Insert into queue.
	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];
	Queue.Add(ProductClass);

	UE_LOG(LogRTS, Log, TEXT("%s queued %s for production in queue %i."), *GetOwner()->GetName(), *ProductClass->GetName(), QueueIndex);

	// Notify listeners.
	OnProductQueued.Broadcast(ProductClass, QueueIndex);

	if (Queue.Num() == 1)
	{
		// Start production.
		StartProductionInQueue(QueueIndex);
	}
}

void URTSProductionComponent::FinishProduction(int32 QueueIndex /*= 0*/)
{
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];
	Queue.RemainingProductionTime = 0.0f;

	// Get game.
	ARTSGameMode* GameMode = Cast<ARTSGameMode>(UGameplayStatics::GetGameMode(this));

	if (!GameMode)
	{
		return;
	}

    TSubclassOf<AActor> ProductClass = Queue[0];

    // Determine spawn location: Start at producing actor location.
    FVector SpawnLocation = GetOwner()->GetActorLocation();

    // Spawn next to production actor.
    float SpawnOffset = 0.0f;
    SpawnOffset += URTSUtilities::GetActorCollisionSize(GetOwner()) / 2;
    SpawnOffset += URTSUtilities::GetCollisionSize(ProductClass) / 2;
    SpawnOffset *= 1.05f;
    SpawnLocation.X -= SpawnOffset;

    // Calculate location on the ground.
    SpawnLocation = URTSUtilities::GetGroundLocation(this, SpawnLocation);

    // Prevent spawn collision or spawning at wrong side of the world.
    SpawnLocation.Z += URTSUtilities::GetCollisionHeight(ProductClass) + 1.0f;

	// Spawn product.
	AActor* Product = GameMode->SpawnActorForPlayer(
		ProductClass,
		Cast<AController>(GetOwner()->GetOwner()),
		FTransform(FRotator::ZeroRotator, SpawnLocation));

	if (!Product)
	{
		return;
	}

	UE_LOG(LogRTS, Log, TEXT("%s finished producing %s in queue %i."), *GetOwner()->GetName(), *Product->GetName(), QueueIndex);

	// Notify listeners.
	OnProductionFinished.Broadcast(Product, QueueIndex);

	// Remove product from queue.
	DequeueProduct(QueueIndex);
}

void URTSProductionComponent::CancelProduction(int32 QueueIndex /*= 0*/, int32 ProductIndex /*= 0*/)
{
	// Get queue.
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	// Get product.
	if (ProductIndex < 0 || ProductIndex >= Queue.Num())
	{
		return;
	}

	TSubclassOf<AActor> ProductClass = Queue[ProductIndex];

	// Get elapsed production time.
	float TotalProductionTime = GetProductionTimeForProduct(ProductClass);
	float RemainingProductionTime = ProductIndex == 0 ? ProductionQueues[QueueIndex].RemainingProductionTime : TotalProductionTime;
	float ElapsedProductionTime = TotalProductionTime - RemainingProductionTime;

	UE_LOG(LogRTS, Log, TEXT("%s canceled producing product %i of class %s in queue %i after %f seconds."),
		*GetOwner()->GetName(),
		ProductIndex,
		*ProductClass->GetName(),
		QueueIndex,
		ElapsedProductionTime);

	// Notify listeners.
	OnProductionCanceled.Broadcast(ProductClass, QueueIndex, ElapsedProductionTime);

	// Remove product from queue.
	DequeueProduct(QueueIndex, ProductIndex);

	// Refund resources.
	URTSProductionCostComponent* ProductionCostComponent =
		URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);

	if (ProductionCostComponent)
	{
		auto Owner = GetOwner()->GetOwner();

		if (!Owner)
		{
			return;
		}

        auto PlayerResourcesComponent = Owner->FindComponentByClass<URTSPlayerResourcesComponent>();

        if (!PlayerResourcesComponent)
        {
            return;
        }

		float TimeRefundFactor = 0.0f;

		if (ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayImmediately)
		{
			TimeRefundFactor = 1.0f;
		}
		else if (ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayOverTime)
		{
			TimeRefundFactor = ElapsedProductionTime / TotalProductionTime;
		}

		float ActualRefundFactor = ProductionCostComponent->RefundFactor * TimeRefundFactor;

		// Refund production costs.
		for (auto& Resource : ProductionCostComponent->Resources)
		{
			TSubclassOf<URTSResourceType> ResourceType = Resource.Key;
			float ResourceAmount = Resource.Value * ActualRefundFactor;

            PlayerResourcesComponent->AddResources(ResourceType, ResourceAmount);

			UE_LOG(LogRTS, Log, TEXT("%f %s of production costs refunded."), ResourceAmount, *ResourceType->GetName());

			// Notify listeners.
			OnProductionCostRefunded.Broadcast(ResourceType, ResourceAmount);
		}
	}
}

void URTSProductionComponent::DequeueProduct(int32 QueueIndex /*= 0*/, int32 ProductIndex /*= 0*/)
{
	// Get queue.
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	if (ProductIndex < 0 || ProductIndex >= Queue.Num())
	{
		return;
	}

	Queue.RemoveAt(ProductIndex);

	// Check if need to start next production.
	if (ProductIndex == 0 && Queue.Num() > 0)
	{
		StartProductionInQueue(QueueIndex);
	}
}

void URTSProductionComponent::StartProductionInQueue(int32 QueueIndex /*= 0*/)
{
	// Get queue.
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	// Get product.
	if (Queue.Num() <= 0)
	{
		return;
	}

	TSubclassOf<AActor> ProductClass = Queue[0];

	// Start production.
	float ProductionTime = GetProductionTimeForProduct(ProductClass);
	Queue.RemainingProductionTime = ProductionTime;

	UE_LOG(LogRTS, Log, TEXT("%s started producing %s in queue %i."), *GetOwner()->GetName(), *ProductClass->GetName(), QueueIndex);

	// Notify listeners.
	OnProductionStarted.Broadcast(ProductClass, QueueIndex, ProductionTime);
}
<commit_msg>Prevent spawn collision or spawning at wrong side of the world for produced actors.<commit_after>#include "RTSPluginPCH.h"
#include "RTSProductionComponent.h"

#include "GameFramework/Actor.h"
#include "GameFramework/Controller.h"
#include "Engine/BlueprintGeneratedClass.h"
#include "Engine/SCS_Node.h"
#include "Kismet/GameplayStatics.h"

#include "RTSGameMode.h"
#include "RTSProductionCostComponent.h"
#include "RTSPlayerController.h"
#include "RTSPlayerAdvantageComponent.h"
#include "RTSPlayerResourcesComponent.h"
#include "RTSUtilities.h"

URTSProductionComponent::URTSProductionComponent(const FObjectInitializer& ObjectInitializer)
	: Super(ObjectInitializer)
{
	PrimaryComponentTick.bCanEverTick = true;

	SetIsReplicated(true);
}

void URTSProductionComponent::GetLifetimeReplicatedProps(TArray<FLifetimeProperty>& OutLifetimeProps) const
{
	Super::GetLifetimeReplicatedProps(OutLifetimeProps);

	DOREPLIFETIME(URTSProductionComponent, ProductionQueues);
}

void URTSProductionComponent::BeginPlay()
{
	UActorComponent::BeginPlay();

	// Setup queues.
	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		FRTSProductionQueue NewQueue;
		ProductionQueues.Add(NewQueue);
	}
}

void URTSProductionComponent::TickComponent(float DeltaTime, enum ELevelTick TickType, FActorComponentTickFunction *ThisTickFunction)
{
	UActorComponent::TickComponent(DeltaTime, TickType, ThisTickFunction);

    // Check for speed boosts.
    float SpeedBoostFactor = 1.0f;
    AActor* OwningActor = GetOwner();

    if (OwningActor)
    {
        AActor* OwningPlayer = OwningActor->GetOwner();

        if (OwningPlayer)
        {
            URTSPlayerAdvantageComponent* PlayerAdvantageComponent = OwningPlayer->FindComponentByClass<URTSPlayerAdvantageComponent>();

            if (PlayerAdvantageComponent)
            {
                SpeedBoostFactor = PlayerAdvantageComponent->SpeedBoostFactor;
            }
        }
    }

	// Process all queues.
	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		if (ProductionQueues[QueueIndex].Num() <= 0)
		{
			continue;
		}

		// Check production costs.
		auto ProductClass = GetCurrentProduction(QueueIndex);
		auto ProductionCostComponent = URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);

		bool bProductionCostPaid = false;

		if (ProductionCostComponent && ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayOverTime)
		{
			auto Owner = GetOwner()->GetOwner();

			if (!Owner)
			{
				UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no owner."), *GetOwner()->GetName());
				continue;
			}

            auto PlayerResourcesComponent = Owner->FindComponentByClass<URTSPlayerResourcesComponent>();

            if (!PlayerResourcesComponent)
            {
                UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no PlayerResourcesComponent."), *Owner->GetName());
                continue;
            }

			bool bCanPayAllProductionCosts = true;

			for (auto& Resource : ProductionCostComponent->Resources)
			{
				float ResourceAmount = Resource.Value * SpeedBoostFactor * DeltaTime / ProductionCostComponent->ProductionTime;

				if (!PlayerResourcesComponent->CanPayResources(Resource.Key, ResourceAmount))
				{
					// Production stopped until resources become available again.
					bCanPayAllProductionCosts = false;
					break;
				}
			}

			if (bCanPayAllProductionCosts)
			{
				// Pay production costs.
				for (auto& Resource : ProductionCostComponent->Resources)
				{
					float ResourceAmount = Resource.Value * SpeedBoostFactor * DeltaTime / ProductionCostComponent->ProductionTime;
                    PlayerResourcesComponent->PayResources(Resource.Key, ResourceAmount);
				}

				bProductionCostPaid = true;
			}
		}
		else
		{
			bProductionCostPaid = true;
		}

		if (!bProductionCostPaid)
		{
			continue;
		}

		// Update production progress.
		ProductionQueues[QueueIndex].RemainingProductionTime -= SpeedBoostFactor * DeltaTime;

		// Check if finished.
		if (ProductionQueues[QueueIndex].RemainingProductionTime <= 0)
		{
			FinishProduction(QueueIndex);
		}
	}
}

bool URTSProductionComponent::CanAssignProduction(TSubclassOf<AActor> ProductClass) const
{
	return URTSUtilities::IsReadyToUse(GetOwner()) && FindQueueForProduct(ProductClass) >= 0;
}

int32 URTSProductionComponent::FindQueueForProduct(TSubclassOf<AActor> ProductClass) const
{
	// Find queue with least products that is not at capacity limit.
	int32 QueueWithLeastProducts = -1;
	int32 ProductsInShortestQueue = CapacityPerQueue;

	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		const FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

		if (Queue.Num() < ProductsInShortestQueue)
		{
			QueueWithLeastProducts = QueueIndex;
			ProductsInShortestQueue = Queue.Num();
		}
	}

	return QueueWithLeastProducts;
}

TSubclassOf<AActor> URTSProductionComponent::GetCurrentProduction(int32 QueueIndex /*= 0*/) const
{
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return nullptr;
	}

	const FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	return Queue.Num() > 0 ? Queue[0] : nullptr;
}

float URTSProductionComponent::GetProductionTime(int32 QueueIndex /*= 0*/) const
{
	TSubclassOf<AActor> CurrentProduction = GetCurrentProduction(QueueIndex);

	if (!CurrentProduction)
	{
		return 0.0f;
	}

	return GetProductionTimeForProduct(CurrentProduction);
}

float URTSProductionComponent::GetProductionTimeForProduct(TSubclassOf<AActor> ProductClass) const
{
	URTSProductionCostComponent* ProductionCostComponent =
		URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);
	return ProductionCostComponent ? ProductionCostComponent->ProductionTime : 0.0f;
}

float URTSProductionComponent::GetProgressPercentage(int32 QueueIndex /*= 0*/) const
{
	float TotalProductionTime = GetProductionTime(QueueIndex);

	if (TotalProductionTime <= 0.0f)
	{
		return 1.0f;
	}

	float RemainingProductionTime = GetRemainingProductionTime(QueueIndex);

	if (RemainingProductionTime <= 0.0f)
	{
		return 1.0f;
	}

	return 1 - (RemainingProductionTime / TotalProductionTime);
}

float URTSProductionComponent::GetRemainingProductionTime(int32 QueueIndex /*= 0*/) const
{
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return 0.0f;
	}

	return ProductionQueues[QueueIndex].RemainingProductionTime;
}

bool URTSProductionComponent::IsProducing() const
{
	for (int32 QueueIndex = 0; QueueIndex < QueueCount; ++QueueIndex)
	{
		const FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

		if (Queue.Num() > 0)
		{
			return true;
		}
	}

	return false;
}

void URTSProductionComponent::StartProduction(TSubclassOf<AActor> ProductClass)
{
	// Check production state.
	if (!CanAssignProduction(ProductClass))
	{
		return;
	}

	int32 QueueIndex = FindQueueForProduct(ProductClass);

	if (QueueIndex < 0)
	{
		return;
	}

    // Check requirements.
    TSubclassOf<AActor> MissingRequirement;

    if (URTSUtilities::GetMissingRequirementFor(this, GetOwner(), ProductClass, MissingRequirement))
    {
        UE_LOG(LogRTS, Error, TEXT("%s wants to produce %s, but is missing requirement %s."), *GetOwner()->GetName(), *ProductClass->GetName(), *MissingRequirement->GetName());
        return;
    }

	// Check production cost.
	URTSProductionCostComponent* ProductionCostComponent =
		URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);

	if (ProductionCostComponent && ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayImmediately)
	{
		auto Owner = GetOwner()->GetOwner();

		if (!Owner)
		{
			UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no owner."), *GetOwner()->GetName());
			return;
		}

        auto PlayerResourcesComponent = Owner->FindComponentByClass<URTSPlayerResourcesComponent>();

        if (!PlayerResourcesComponent)
        {
            UE_LOG(LogRTS, Error, TEXT("%s needs to pay for production, but has no PlayerResourcesComponent."), *Owner->GetName());
            return;
        }

		if (!PlayerResourcesComponent->CanPayAllResources(ProductionCostComponent->Resources))
		{
			UE_LOG(LogRTS, Error, TEXT("%s needs to pay for producing %s, but does not have enough resources."),
				*Owner->GetName(),
				*ProductClass->GetName());
			return;
		}

		// Pay production costs.
        PlayerResourcesComponent->PayAllResources(ProductionCostComponent->Resources);
	}

	// Insert into queue.
	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];
	Queue.Add(ProductClass);

	UE_LOG(LogRTS, Log, TEXT("%s queued %s for production in queue %i."), *GetOwner()->GetName(), *ProductClass->GetName(), QueueIndex);

	// Notify listeners.
	OnProductQueued.Broadcast(ProductClass, QueueIndex);

	if (Queue.Num() == 1)
	{
		// Start production.
		StartProductionInQueue(QueueIndex);
	}
}

void URTSProductionComponent::FinishProduction(int32 QueueIndex /*= 0*/)
{
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];
	Queue.RemainingProductionTime = 0.0f;

	// Get game.
	ARTSGameMode* GameMode = Cast<ARTSGameMode>(UGameplayStatics::GetGameMode(this));

	if (!GameMode)
	{
		return;
	}

    TSubclassOf<AActor> ProductClass = Queue[0];

    // Determine spawn location: Start at producing actor location.
    FVector SpawnLocation = GetOwner()->GetActorLocation();

    // Spawn next to production actor.
    float SpawnOffset = 0.0f;
    SpawnOffset += URTSUtilities::GetActorCollisionSize(GetOwner()) / 2;
    SpawnOffset += URTSUtilities::GetCollisionSize(ProductClass) / 2;
    SpawnOffset *= 1.05f;
    SpawnLocation.X -= SpawnOffset;

    // Calculate location on the ground.
    SpawnLocation = URTSUtilities::GetGroundLocation(this, SpawnLocation);

    // Prevent spawn collision or spawning at wrong side of the world.
    SpawnLocation.Z += URTSUtilities::GetCollisionHeight(ProductClass) * 1.1f;

	// Spawn product.
	AActor* Product = GameMode->SpawnActorForPlayer(
		ProductClass,
		Cast<AController>(GetOwner()->GetOwner()),
		FTransform(FRotator::ZeroRotator, SpawnLocation));

	if (!Product)
	{
		return;
	}

	UE_LOG(LogRTS, Log, TEXT("%s finished producing %s in queue %i."), *GetOwner()->GetName(), *Product->GetName(), QueueIndex);

	// Notify listeners.
	OnProductionFinished.Broadcast(Product, QueueIndex);

	// Remove product from queue.
	DequeueProduct(QueueIndex);
}

void URTSProductionComponent::CancelProduction(int32 QueueIndex /*= 0*/, int32 ProductIndex /*= 0*/)
{
	// Get queue.
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	// Get product.
	if (ProductIndex < 0 || ProductIndex >= Queue.Num())
	{
		return;
	}

	TSubclassOf<AActor> ProductClass = Queue[ProductIndex];

	// Get elapsed production time.
	float TotalProductionTime = GetProductionTimeForProduct(ProductClass);
	float RemainingProductionTime = ProductIndex == 0 ? ProductionQueues[QueueIndex].RemainingProductionTime : TotalProductionTime;
	float ElapsedProductionTime = TotalProductionTime - RemainingProductionTime;

	UE_LOG(LogRTS, Log, TEXT("%s canceled producing product %i of class %s in queue %i after %f seconds."),
		*GetOwner()->GetName(),
		ProductIndex,
		*ProductClass->GetName(),
		QueueIndex,
		ElapsedProductionTime);

	// Notify listeners.
	OnProductionCanceled.Broadcast(ProductClass, QueueIndex, ElapsedProductionTime);

	// Remove product from queue.
	DequeueProduct(QueueIndex, ProductIndex);

	// Refund resources.
	URTSProductionCostComponent* ProductionCostComponent =
		URTSUtilities::FindDefaultComponentByClass<URTSProductionCostComponent>(ProductClass);

	if (ProductionCostComponent)
	{
		auto Owner = GetOwner()->GetOwner();

		if (!Owner)
		{
			return;
		}

        auto PlayerResourcesComponent = Owner->FindComponentByClass<URTSPlayerResourcesComponent>();

        if (!PlayerResourcesComponent)
        {
            return;
        }

		float TimeRefundFactor = 0.0f;

		if (ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayImmediately)
		{
			TimeRefundFactor = 1.0f;
		}
		else if (ProductionCostComponent->ProductionCostType == ERTSProductionCostType::COST_PayOverTime)
		{
			TimeRefundFactor = ElapsedProductionTime / TotalProductionTime;
		}

		float ActualRefundFactor = ProductionCostComponent->RefundFactor * TimeRefundFactor;

		// Refund production costs.
		for (auto& Resource : ProductionCostComponent->Resources)
		{
			TSubclassOf<URTSResourceType> ResourceType = Resource.Key;
			float ResourceAmount = Resource.Value * ActualRefundFactor;

            PlayerResourcesComponent->AddResources(ResourceType, ResourceAmount);

			UE_LOG(LogRTS, Log, TEXT("%f %s of production costs refunded."), ResourceAmount, *ResourceType->GetName());

			// Notify listeners.
			OnProductionCostRefunded.Broadcast(ResourceType, ResourceAmount);
		}
	}
}

void URTSProductionComponent::DequeueProduct(int32 QueueIndex /*= 0*/, int32 ProductIndex /*= 0*/)
{
	// Get queue.
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	if (ProductIndex < 0 || ProductIndex >= Queue.Num())
	{
		return;
	}

	Queue.RemoveAt(ProductIndex);

	// Check if need to start next production.
	if (ProductIndex == 0 && Queue.Num() > 0)
	{
		StartProductionInQueue(QueueIndex);
	}
}

void URTSProductionComponent::StartProductionInQueue(int32 QueueIndex /*= 0*/)
{
	// Get queue.
	if (QueueIndex < 0 || QueueIndex >= QueueCount)
	{
		return;
	}

	FRTSProductionQueue& Queue = ProductionQueues[QueueIndex];

	// Get product.
	if (Queue.Num() <= 0)
	{
		return;
	}

	TSubclassOf<AActor> ProductClass = Queue[0];

	// Start production.
	float ProductionTime = GetProductionTimeForProduct(ProductClass);
	Queue.RemainingProductionTime = ProductionTime;

	UE_LOG(LogRTS, Log, TEXT("%s started producing %s in queue %i."), *GetOwner()->GetName(), *ProductClass->GetName(), QueueIndex);

	// Notify listeners.
	OnProductionStarted.Broadcast(ProductClass, QueueIndex, ProductionTime);
}
<|endoftext|>


<commit_before>//
//	AudioListener.cpp
//	Created by Matt Kaufmann on 04/10/14.
//

#include "Vajra/Common/Messages/Message.h"
#include "Vajra/Common/Objects/Object.h"
#include "Vajra/Engine/AudioManager/AudioManager.h"
#include "Vajra/Engine/Components/ComponentTypes/ComponentTypeIds.h"
#include "Vajra/Engine/Components/DerivedComponents/Audio/AudioListener.h"
#include "Vajra/Engine/Components/DerivedComponents/Transform/Transform.h"
#include "Vajra/Engine/Core/Engine.h"

unsigned int AudioListener::componentTypeId = COMPONENT_TYPE_ID_AUDIO_LISTENER;
ObjectIdType AudioListener::activeListener  = OBJECT_ID_INVALID;

// Constructors
AudioListener::AudioListener() : Component() {
	this->init();
}

AudioListener::AudioListener(Object* object_) : Component(object_) {
	this->init();
}

// Destructor
AudioListener::~AudioListener() {
	this->destroy();
}

bool AudioListener::Is3DSoundEnabled() {
	return this->is3D;
}

void AudioListener::Enable3DSound() {
	this->is3D = true;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->Enable3DSound();
	}
}

void AudioListener::Disable3DSound() {
	this->is3D = false;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->Disable3DSound();
	}
}

void AudioListener::SetVelocity(glm::vec3 vel) {
	this->velocity = vel;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->SetListenerVelocity(vel);
	}
}

void AudioListener::SetVelocity(float x, float y, float z) {
	this->velocity = glm::vec3(x, y, z);
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->SetListenerVelocity(x, y, z);
	}
}

void AudioListener::SetVolume(float vol) {
	this->volume = vol;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->SetListenerVolume(vol);
	}
}

void AudioListener::HandleMessage(MessageChunk messageChunk) {
	Component::HandleMessage(messageChunk);

	switch (messageChunk->GetMessageType()) {
		case MESSAGE_TYPE_TRANSFORM_CHANGED_EVENT:
			onTransformChanged();
			break;
	}
}

void AudioListener::SetAsActiveListener() {
	if (AudioListener::activeListener != this->GetObject()->GetId()) {
		if (this->is3D) {
			ENGINE->GetAudioManager()->Enable3DSound();
		}
		else {
			ENGINE->GetAudioManager()->Disable3DSound();
		}

		Transform* trans = this->GetObject()->GetComponent<Transform>();
		ENGINE->GetAudioManager()->SetListenerPosition(trans->GetPositionWorld());
		ENGINE->GetAudioManager()->SetListenerOrientation(trans->GetOrientationWorld());

		ENGINE->GetAudioManager()->SetListenerVelocity(this->velocity);
		ENGINE->GetAudioManager()->SetListenerVolume(this->volume);

		AudioListener::activeListener = this->GetObject()->GetId();
	}
}

void AudioListener::init() {
	this->is3D = false;
	this->velocity = ZERO_VEC3;
	this->volume = 0.0f;

	this->addSubscriptionToMessageType(MESSAGE_TYPE_TRANSFORM_CHANGED_EVENT, this->GetTypeId(), true);
}

void AudioListener::destroy() {

}

void AudioListener::onTransformChanged() {
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		Transform* trans = this->GetObject()->GetComponent<Transform>();
		ENGINE->GetAudioManager()->SetListenerPosition(trans->GetPositionWorld());
		ENGINE->GetAudioManager()->SetListenerOrientation(trans->GetOrientationWorld());
	}
}
<commit_msg>unmuted audio listener<commit_after>//
//	AudioListener.cpp
//	Created by Matt Kaufmann on 04/10/14.
//

#include "Vajra/Common/Messages/Message.h"
#include "Vajra/Common/Objects/Object.h"
#include "Vajra/Engine/AudioManager/AudioManager.h"
#include "Vajra/Engine/Components/ComponentTypes/ComponentTypeIds.h"
#include "Vajra/Engine/Components/DerivedComponents/Audio/AudioListener.h"
#include "Vajra/Engine/Components/DerivedComponents/Transform/Transform.h"
#include "Vajra/Engine/Core/Engine.h"

unsigned int AudioListener::componentTypeId = COMPONENT_TYPE_ID_AUDIO_LISTENER;
ObjectIdType AudioListener::activeListener  = OBJECT_ID_INVALID;

// Constructors
AudioListener::AudioListener() : Component() {
	this->init();
}

AudioListener::AudioListener(Object* object_) : Component(object_) {
	this->init();
}

// Destructor
AudioListener::~AudioListener() {
	this->destroy();
}

bool AudioListener::Is3DSoundEnabled() {
	return this->is3D;
}

void AudioListener::Enable3DSound() {
	this->is3D = true;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->Enable3DSound();
	}
}

void AudioListener::Disable3DSound() {
	this->is3D = false;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->Disable3DSound();
	}
}

void AudioListener::SetVelocity(glm::vec3 vel) {
	this->velocity = vel;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->SetListenerVelocity(vel);
	}
}

void AudioListener::SetVelocity(float x, float y, float z) {
	this->velocity = glm::vec3(x, y, z);
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->SetListenerVelocity(x, y, z);
	}
}

void AudioListener::SetVolume(float vol) {
	this->volume = vol;
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		ENGINE->GetAudioManager()->SetListenerVolume(vol);
	}
}

void AudioListener::HandleMessage(MessageChunk messageChunk) {
	Component::HandleMessage(messageChunk);

	switch (messageChunk->GetMessageType()) {
		case MESSAGE_TYPE_TRANSFORM_CHANGED_EVENT:
			onTransformChanged();
			break;
	}
}

void AudioListener::SetAsActiveListener() {
	if (AudioListener::activeListener != this->GetObject()->GetId()) {
		if (this->is3D) {
			ENGINE->GetAudioManager()->Enable3DSound();
		}
		else {
			ENGINE->GetAudioManager()->Disable3DSound();
		}

		Transform* trans = this->GetObject()->GetComponent<Transform>();
		ENGINE->GetAudioManager()->SetListenerPosition(trans->GetPositionWorld());
		ENGINE->GetAudioManager()->SetListenerOrientation(trans->GetOrientationWorld());

		ENGINE->GetAudioManager()->SetListenerVelocity(this->velocity);
		ENGINE->GetAudioManager()->SetListenerVolume(this->volume);

		AudioListener::activeListener = this->GetObject()->GetId();
	}
}

void AudioListener::init() {
	this->is3D = false;
	this->velocity = ZERO_VEC3;
	this->volume = 1.0f;

	this->addSubscriptionToMessageType(MESSAGE_TYPE_TRANSFORM_CHANGED_EVENT, this->GetTypeId(), true);
}

void AudioListener::destroy() {

}

void AudioListener::onTransformChanged() {
	if (AudioListener::activeListener == this->GetObject()->GetId()) {
		Transform* trans = this->GetObject()->GetComponent<Transform>();
		ENGINE->GetAudioManager()->SetListenerPosition(trans->GetPositionWorld());
		ENGINE->GetAudioManager()->SetListenerOrientation(trans->GetOrientationWorld());
	}
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2012-2013 Danny Y., Rapptz

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#ifndef GEARS_PROGRAM_OPTIONS_ARG_HPP
#define GEARS_PROGRAM_OPTIONS_ARG_HPP

#include <string>

#ifndef GEARS_NO_IOSTREAM
#include <iosfwd>
#endif // GEARS_NO_IOSTREAM

namespace gears {
struct arg {
private:
    friend class program_options;
    std::string name;
    std::string description;
    std::string parameter;
    std::string value;
    bool active = false;
    char short_name = '\0';
public:
    arg(std::string name, std::string desc = "", char shorter = '\0', std::string p = "") noexcept:
    name(std::move(name)), description(std::move(desc)), parameter(std::move(p)), short_name(shorter) {}

    template<typename Elem, typename Traits>
    friend auto operator<<(std::basic_ostream<Elem, Traits>& out, const arg& a) -> decltype(out) {
        int spaces = 30;
        out << "    ";
        if(!a.name.empty()) {
            if(a.short_name != '\0') {
                out << '-' << a.short_name << ", ";
                spaces -= 4;
            }

            out << "--" << a.name;
            spaces -= 2 + a.name.size();

            if(!a.parameter.empty()) {
                out << "[=<" << a.parameter << ">]";
                spaces -= 5 + a.parameter.size();
            }

            for(int i = 0; i < spaces; ++i) {
                out << ' ';
            }

            out << a.description << '\n';
        }

        return out;
    }

    bool is_value() const noexcept {
        return !parameter.empty();
    }

    arg& shorter(char c) noexcept {
        short_name = c;
        return *this;
    }

    arg& help(std::string str) noexcept {
        description = std::move(str);
        return *this;
    }

    arg& param(std::string str) noexcept {
        parameter = std::move(str);
        return *this;
    }
};
} // gears

#endif // GEARS_PROGRAM_OPTIONS_ARG_HPP<commit_msg>Fix issue with GEARS_NO_IOSTREAM missing<commit_after>// The MIT License (MIT)

// Copyright (c) 2012-2013 Danny Y., Rapptz

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#ifndef GEARS_PROGRAM_OPTIONS_ARG_HPP
#define GEARS_PROGRAM_OPTIONS_ARG_HPP

#include <string>

#ifndef GEARS_NO_IOSTREAM
#include <iosfwd>
#endif // GEARS_NO_IOSTREAM

namespace gears {
struct arg {
private:
    friend class program_options;
    std::string name;
    std::string description;
    std::string parameter;
    std::string value;
    bool active = false;
    char short_name = '\0';
public:
    arg(std::string name, std::string desc = "", char shorter = '\0', std::string p = "") noexcept:
    name(std::move(name)), description(std::move(desc)), parameter(std::move(p)), short_name(shorter) {}

    bool is_value() const noexcept {
        return !parameter.empty();
    }

    arg& shorter(char c) noexcept {
        short_name = c;
        return *this;
    }

    arg& help(std::string str) noexcept {
        description = std::move(str);
        return *this;
    }

    arg& param(std::string str) noexcept {
        parameter = std::move(str);
        return *this;
    }

    #ifndef GEARS_NO_IOSTREAM

    template<typename Elem, typename Traits>
    friend auto operator<<(std::basic_ostream<Elem, Traits>& out, const arg& a) -> decltype(out) {
        int spaces = 30;
        out << "    ";
        if(!a.name.empty()) {
            if(a.short_name != '\0') {
                out << '-' << a.short_name << ", ";
                spaces -= 4;
            }

            out << "--" << a.name;
            spaces -= 2 + a.name.size();

            if(!a.parameter.empty()) {
                out << "[=<" << a.parameter << ">]";
                spaces -= 5 + a.parameter.size();
            }

            for(int i = 0; i < spaces; ++i) {
                out << ' ';
            }

            out << a.description << '\n';
        }

        return out;
    }

    #endif // GEARS_NO_IOSTREAM
};
} // gears

#endif // GEARS_PROGRAM_OPTIONS_ARG_HPP<|endoftext|>


<commit_before>#include <iostream>
#include <External/gl_core_3_3.hpp>
#include <SDL2/SDL.h>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <Base/Filesystem/File.h>
#include "messagebox.h"
using namespace gl;

int keyPressed[SDL_NUM_SCANCODES] = {0};
int WINDOW_WIDTH = 1280, WINDOW_HEIGHT = 720;

class String
{
    private:
        const char* stringData;
    public:
        String(const char* cstr) : stringData(cstr) {}
        const char* str() const
        {
            return stringData;
        }
};

class Shader
{
    private:
        GLuint vertID, fragID, progID;
        GLuint newShader(GLenum shaderType, const char* source)
        {
            // Tutorial code from http://www.arcsynthesis.org/gltut
            // Copyright © 2012 Jason L. McKesson
            // Variable names changed take custom parameters and work with glLoadGen func_cpp functions
            GLuint shader = CreateShader(shaderType);
            ShaderSource(shader, 1, &source, NULL);
            CompileShader(shader);
            GLint status;
            GetShaderiv(shader, COMPILE_STATUS, &status);
            if (status == FALSE_)
            {
               GLint infoLogLength;
               GetShaderiv(shader, INFO_LOG_LENGTH, &infoLogLength);

               GLchar *strInfoLog = new GLchar[infoLogLength + 1];
               GetShaderInfoLog(shader, infoLogLength, NULL, strInfoLog);

               const char *strShaderType = NULL;
               switch(shaderType)
               {
               case VERTEX_SHADER: strShaderType = "vertex"; break;
               case GEOMETRY_SHADER: strShaderType = "geometry"; break;
               case FRAGMENT_SHADER: strShaderType = "fragment"; break;
               }

               fprintf(stderr, "Compile failure in %s shader:\n%s\n", strShaderType, strInfoLog);
               delete[] strInfoLog;
            }
            return shader;
        }
    public:
        Shader(const File& vert, const File& frag)
        {
            // Create both shaders
            vertID = newShader(VERTEX_SHADER, vert.string().c_str());
            fragID = newShader(FRAGMENT_SHADER, frag.string().c_str());

            // Create and link a program
            progID = CreateProgram();
            AttachShader(progID, vertID);
            AttachShader(progID, fragID);
            LinkProgram(progID);
            DetachShader(progID, vertID);
            DetachShader(progID, fragID);
        }
        const int getProgram() {return progID;}
};

const float triangle[] = {
    // Positions
    -1, -1,  0,  1,
     1, -1,  0,  1,
     0,  1,  0,  1,
    // Colors
     1,  0,  0,  1,
     0,  1,  0,  1,
     0,  0,  1,  1
};

#include "cube.cpp"

class TestScene
{
    private:
        Shader shader;
        GLuint VBO_id;
        GLuint VAO_id;
        GLuint mvpLocation, offsetLocation;
        unsigned deltaTime;
        unsigned currentFrame;
        unsigned lastFrame;

        glm::mat4 model, view, projection;
        glm::vec4 offset;

    public:
        TestScene() : shader("Shaders/UnlitGeneric.vert", "Shaders/UnlitGeneric.frag"), lastFrame(SDL_GetTicks())
        {
            ClearColor(0, 0, 0, 1);
            ClearDepth(1);
            Enable(DEPTH_TEST);
            Enable(CULL_FACE);
            CullFace(BACK);
            FrontFace(CW);
            Viewport(0, 0, WINDOW_WIDTH, WINDOW_HEIGHT);

            // Set up uniforms for the shader
            mvpLocation = GetUniformLocation(shader.getProgram(), "MVP");
            offsetLocation = GetUniformLocation(shader.getProgram(), "offset");

            // Generate a VBO and VAO
            GenBuffers(1, &VBO_id);
            GenVertexArrays(1, &VAO_id);
            // Bind them
            BindBuffer(ARRAY_BUFFER, VBO_id);
            BindVertexArray(VAO_id);
            // Set the buffer data
            BufferData(ARRAY_BUFFER, sizeof(vertexData), vertexData, STATIC_DRAW);
            // Enable the first two vertex attributes
            EnableVertexAttribArray(0);
            EnableVertexAttribArray(1);
            // Attribute index, 4 values per position, inform they're floats, unknown, space between values, first value
            VertexAttribPointer(0, 4, gl::FLOAT, false, 0, 0);
            VertexAttribPointer(1, 4, gl::FLOAT, false, 0, (void*) (sizeof(vertexData) / 2));
            // And clean up
            DisableVertexAttribArray(0);
            DisableVertexAttribArray(1);
            BindBuffer(ARRAY_BUFFER, 0);
            BindVertexArray(0);
        }
        void draw()
        {
            static unsigned totalTime = 0;
            currentFrame = SDL_GetTicks();
            deltaTime = currentFrame - lastFrame;
            lastFrame = currentFrame;

            totalTime += deltaTime;

            // Matrices
            model = glm::mat4(1);

            view = glm::lookAt(glm::vec3(0, 0, 2),
                               glm::vec3(0, 0, 0),
                               glm::vec3(0, 1, 0));

            int sizeX = WINDOW_WIDTH, sizeY = WINDOW_HEIGHT;

            projection = glm::perspective(60.0, (double)sizeX/(double)sizeY, 0.01, 10000.0);

            glm::mat4 modelViewProjection = projection * view * model;

            static float modelX = 0, modelY = 0;
            if (keyPressed[SDL_SCANCODE_LEFT])
            {
                modelX -= float(deltaTime * 0.001);
            }
            if (keyPressed[SDL_SCANCODE_RIGHT])
            {
                modelX += float(deltaTime * 0.001);
            }
            if (keyPressed[SDL_SCANCODE_UP])
            {
                modelY += float(deltaTime * 0.001);
            }
            if (keyPressed[SDL_SCANCODE_DOWN])
            {
                modelY -= float(deltaTime * 0.001);
            }

            // Offset
            offset.x = modelX;
            offset.y = modelY;
            offset.z = 0;
            offset.w = 1;

            // Set the shader program
            UseProgram(shader.getProgram());
            UniformMatrix4fv(mvpLocation, 1, FALSE_, &modelViewProjection[0][0]);
            Uniform4fv(offsetLocation, 1, &offset[0]);
            // Bind the vertex array
            BindVertexArray(VAO_id);
            EnableVertexAttribArray(0);
            EnableVertexAttribArray(1);
            // Draw the values
            DrawArrays(TRIANGLES, 0, 36);
            // And unbind the vertex array
            DisableVertexAttribArray(0);
            DisableVertexAttribArray(1);
            BindVertexArray(0);
        }
};

void cbfun_windowResized(int width, int height)
{
    WINDOW_WIDTH = width;
    WINDOW_HEIGHT = height;
    Viewport(0, 0, WINDOW_WIDTH, WINDOW_HEIGHT);
}

int main(int argc, char** argv)
{
    PHYSFS_init(argv[0]);
    setRootPath("../Data");
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_PROFILE_MASK, SDL_GL_CONTEXT_PROFILE_CORE);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_FLAGS, SDL_GL_CONTEXT_FORWARD_COMPATIBLE_FLAG);
    SDL_GL_SetAttribute(SDL_GL_MULTISAMPLEBUFFERS, 1);
    SDL_GL_SetAttribute(SDL_GL_MULTISAMPLESAMPLES, 4);
    SDL_Init(SDL_INIT_VIDEO);
    SDL_Window* window = SDL_CreateWindow("SomeGame (SDL)",
                                          SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED,
                                          WINDOW_WIDTH, WINDOW_HEIGHT,
                                          SDL_WINDOW_OPENGL | SDL_WINDOW_RESIZABLE);
    SDL_GLContext context = SDL_GL_CreateContext(window);
    SDL_GL_SetSwapInterval(1);
    if (!window || !sys::LoadFunctions())
    {
        MessageBoxError("Fatal Error", "Could not initialize an OpenGL context\nMake sure your computer supports OpenGL 3.3 and drivers are updated");
        fprintf(stderr, "SDL_GetError(): %s", SDL_GetError());
        return -1;
    }
    fprintf(stdout, "OpenGL version: %s\nDisplay device: %s\nVendor: %s\n", GetString(VERSION), GetString(RENDERER), GetString(VENDOR));

    TestScene scene;
    bool runGame = true;
    while (runGame)
    {
        SDL_Event event;
        while (SDL_PollEvent(&event))
        {
            if (event.type == SDL_KEYDOWN)
            {
                keyPressed[event.key.keysym.scancode] = 1;
            }
            else if (event.type == SDL_KEYUP)
            {
                keyPressed[event.key.keysym.scancode] = 0;
            }

            if (event.type == SDL_QUIT or event.key.keysym.scancode == SDL_SCANCODE_ESCAPE)
            {
                runGame = false;
            }
            else if (event.type == SDL_WINDOWEVENT_RESIZED)
            {
                cbfun_windowResized(event.window.data1, event.window.data2);
            }
        }
        Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
        scene.draw();
        SDL_GL_SwapWindow(window);
    }
    SDL_GL_DeleteContext(context);
    SDL_DestroyWindow(window);
    SDL_Quit();
    return 0;
}
<commit_msg>Fix glViewport() not being called on window resize<commit_after>#include <iostream>
#include <External/gl_core_3_3.hpp>
#include <SDL2/SDL.h>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <Base/Filesystem/File.h>
#include "messagebox.h"
using namespace gl;

int keyPressed[SDL_NUM_SCANCODES] = {0};
int WINDOW_WIDTH = 1280, WINDOW_HEIGHT = 720;

class String
{
    private:
        const char* stringData;
    public:
        String(const char* cstr) : stringData(cstr) {}
        const char* str() const
        {
            return stringData;
        }
};

class Shader
{
    private:
        GLuint vertID, fragID, progID;
        GLuint newShader(GLenum shaderType, const char* source)
        {
            // Tutorial code from http://www.arcsynthesis.org/gltut
            // Copyright © 2012 Jason L. McKesson
            // Variable names changed take custom parameters and work with glLoadGen func_cpp functions
            GLuint shader = CreateShader(shaderType);
            ShaderSource(shader, 1, &source, NULL);
            CompileShader(shader);
            GLint status;
            GetShaderiv(shader, COMPILE_STATUS, &status);
            if (status == FALSE_)
            {
               GLint infoLogLength;
               GetShaderiv(shader, INFO_LOG_LENGTH, &infoLogLength);

               GLchar *strInfoLog = new GLchar[infoLogLength + 1];
               GetShaderInfoLog(shader, infoLogLength, NULL, strInfoLog);

               const char *strShaderType = NULL;
               switch(shaderType)
               {
               case VERTEX_SHADER: strShaderType = "vertex"; break;
               case GEOMETRY_SHADER: strShaderType = "geometry"; break;
               case FRAGMENT_SHADER: strShaderType = "fragment"; break;
               }

               fprintf(stderr, "Compile failure in %s shader:\n%s\n", strShaderType, strInfoLog);
               delete[] strInfoLog;
            }
            return shader;
        }
    public:
        Shader(const File& vert, const File& frag)
        {
            // Create both shaders
            vertID = newShader(VERTEX_SHADER, vert.string().c_str());
            fragID = newShader(FRAGMENT_SHADER, frag.string().c_str());

            // Create and link a program
            progID = CreateProgram();
            AttachShader(progID, vertID);
            AttachShader(progID, fragID);
            LinkProgram(progID);
            DetachShader(progID, vertID);
            DetachShader(progID, fragID);
        }
        const int getProgram() {return progID;}
};

const float triangle[] = {
    // Positions
    -1, -1,  0,  1,
     1, -1,  0,  1,
     0,  1,  0,  1,
    // Colors
     1,  0,  0,  1,
     0,  1,  0,  1,
     0,  0,  1,  1
};

#include "cube.cpp"

class TestScene
{
    private:
        Shader shader;
        GLuint VBO_id;
        GLuint VAO_id;
        GLuint mvpLocation, offsetLocation;
        unsigned deltaTime;
        unsigned currentFrame;
        unsigned lastFrame;

        glm::mat4 model, view, projection;
        glm::vec4 offset;

    public:
        TestScene() : shader("Shaders/UnlitGeneric.vert", "Shaders/UnlitGeneric.frag"), lastFrame(SDL_GetTicks())
        {
            ClearColor(0, 0, 0, 1);
            ClearDepth(1);
            Enable(DEPTH_TEST);
            Enable(CULL_FACE);
            CullFace(BACK);
            FrontFace(CW);
            Viewport(0, 0, WINDOW_WIDTH, WINDOW_HEIGHT);

            // Set up uniforms for the shader
            mvpLocation = GetUniformLocation(shader.getProgram(), "MVP");
            offsetLocation = GetUniformLocation(shader.getProgram(), "offset");

            // Generate a VBO and VAO
            GenBuffers(1, &VBO_id);
            GenVertexArrays(1, &VAO_id);
            // Bind them
            BindBuffer(ARRAY_BUFFER, VBO_id);
            BindVertexArray(VAO_id);
            // Set the buffer data
            BufferData(ARRAY_BUFFER, sizeof(vertexData), vertexData, STATIC_DRAW);
            // Enable the first two vertex attributes
            EnableVertexAttribArray(0);
            EnableVertexAttribArray(1);
            // Attribute index, 4 values per position, inform they're floats, unknown, space between values, first value
            VertexAttribPointer(0, 4, gl::FLOAT, false, 0, 0);
            VertexAttribPointer(1, 4, gl::FLOAT, false, 0, (void*) (sizeof(vertexData) / 2));
            // And clean up
            DisableVertexAttribArray(0);
            DisableVertexAttribArray(1);
            BindBuffer(ARRAY_BUFFER, 0);
            BindVertexArray(0);
        }
        void draw()
        {
            static unsigned totalTime = 0;
            currentFrame = SDL_GetTicks();
            deltaTime = currentFrame - lastFrame;
            lastFrame = currentFrame;

            totalTime += deltaTime;

            // Matrices
            model = glm::mat4(1);

            view = glm::lookAt(glm::vec3(0, 0, 2),
                               glm::vec3(0, 0, 0),
                               glm::vec3(0, 1, 0));

            int sizeX = WINDOW_WIDTH, sizeY = WINDOW_HEIGHT;

            projection = glm::perspective(60.0, (double)sizeX/(double)sizeY, 0.01, 10000.0);

            glm::mat4 modelViewProjection = projection * view * model;

            static float modelX = 0, modelY = 0;
            if (keyPressed[SDL_SCANCODE_LEFT])
            {
                modelX -= float(deltaTime * 0.001);
            }
            if (keyPressed[SDL_SCANCODE_RIGHT])
            {
                modelX += float(deltaTime * 0.001);
            }
            if (keyPressed[SDL_SCANCODE_UP])
            {
                modelY += float(deltaTime * 0.001);
            }
            if (keyPressed[SDL_SCANCODE_DOWN])
            {
                modelY -= float(deltaTime * 0.001);
            }

            // Offset
            offset.x = modelX;
            offset.y = modelY;
            offset.z = 0;
            offset.w = 1;

            // Set the shader program
            UseProgram(shader.getProgram());
            UniformMatrix4fv(mvpLocation, 1, FALSE_, &modelViewProjection[0][0]);
            Uniform4fv(offsetLocation, 1, &offset[0]);
            // Bind the vertex array
            BindVertexArray(VAO_id);
            EnableVertexAttribArray(0);
            EnableVertexAttribArray(1);
            // Draw the values
            DrawArrays(TRIANGLES, 0, 36);
            // And unbind the vertex array
            DisableVertexAttribArray(0);
            DisableVertexAttribArray(1);
            BindVertexArray(0);
        }
};

int main(int argc, char** argv)
{
    PHYSFS_init(argv[0]);
    setRootPath("../Data");
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_PROFILE_MASK, SDL_GL_CONTEXT_PROFILE_CORE);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_FLAGS, SDL_GL_CONTEXT_FORWARD_COMPATIBLE_FLAG);
    SDL_GL_SetAttribute(SDL_GL_MULTISAMPLEBUFFERS, 1);
    SDL_GL_SetAttribute(SDL_GL_MULTISAMPLESAMPLES, 4);
    SDL_Init(SDL_INIT_VIDEO);
    SDL_Window* window = SDL_CreateWindow("SomeGame (SDL)",
                                          SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED,
                                          WINDOW_WIDTH, WINDOW_HEIGHT,
                                          SDL_WINDOW_OPENGL | SDL_WINDOW_RESIZABLE);
    SDL_GLContext context = SDL_GL_CreateContext(window);
    SDL_GL_SetSwapInterval(1);
    if (!window || !sys::LoadFunctions())
    {
        MessageBoxError("Fatal Error", "Could not initialize an OpenGL context\nMake sure your computer supports OpenGL 3.3 and drivers are updated");
        fprintf(stderr, "SDL_GetError(): %s", SDL_GetError());
        return -1;
    }
    fprintf(stdout, "OpenGL version: %s\nDisplay device: %s\nVendor: %s\n", GetString(VERSION), GetString(RENDERER), GetString(VENDOR));

    TestScene scene;
    bool runGame = true;
    while (runGame)
    {
        SDL_Event event;
        while (SDL_PollEvent(&event))
        {
            if (event.type == SDL_KEYDOWN)
            {
                keyPressed[event.key.keysym.scancode] = 1;
            }
            else if (event.type == SDL_KEYUP)
            {
                keyPressed[event.key.keysym.scancode] = 0;
            }

            if (event.type == SDL_QUIT or event.key.keysym.scancode == SDL_SCANCODE_ESCAPE)
            {
                runGame = false;
            }

            if (event.type == SDL_WINDOWEVENT and event.window.event == SDL_WINDOWEVENT_RESIZED)
            {
                WINDOW_WIDTH = event.window.data1;
                WINDOW_HEIGHT = event.window.data2;
                Viewport(0, 0, WINDOW_WIDTH, WINDOW_HEIGHT);
            }
        }
        Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
        scene.draw();
        SDL_GL_SwapWindow(window);
    }
    SDL_GL_DeleteContext(context);
    SDL_DestroyWindow(window);
    SDL_Quit();
    return 0;
}
<|endoftext|>


<commit_before>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    ImageRegistrationHistogramPlotter.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

// The example shows how to use the \doxygen{HistogramToImageFilter} class.
// The example registers two images using the gradient descent optimizer.
// The transform used here is a simple translation transform. The metric
// is a MutualInformationHistogramImageToImageMetric.
//
// The jointHistogramWriter is invoked after every iteration of the optimizer.
// the writer here writes the joint histogram after every iteration as
// JointHistogramXXX.mhd
//
// Note: This may not work for some VNL optimizers like Amoeba which do not throw
// events.
//

#if defined(_MSC_VER)
#pragma warning ( disable : 4786 )
#endif

#include "itkImageRegistrationMethod.h"
#include "itkTranslationTransform.h"
#include "itkMutualInformationHistogramImageToImageMetric.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkGradientDescentOptimizer.h"
#include "itkImage.h"
#include "itkNormalizeImageFilter.h"
#include "itkDiscreteGaussianImageFilter.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkResampleImageFilter.h"
#include "itkCastImageFilter.h"
#include "itkHistogramToImageFilter.h"
#include "itkCommand.h"

const    unsigned int    Dimension = 2;

class HistogramWriter
{
public:
  typedef float InternalPixelType;
  typedef itk::Image< InternalPixelType, Dimension > InternalImageType;
  typedef itk::MutualInformationHistogramImageToImageMetric<
                                          InternalImageType,
                                          InternalImageType >    MetricType;
  typedef MetricType::Pointer   MetricPointer;
  typedef MetricType::HistogramType   HistogramType;
  typedef itk::HistogramToImageFilter< HistogramType > HistogramToImageFilterType;
  typedef HistogramToImageFilterType::Pointer   HistogramToImageFilterPointer;
  typedef itk::ImageFileWriter<  HistogramToImageFilterType::OutputImageType  > HistogramWriterType;
  typedef HistogramWriterType::Pointer   HistogramWriterPointer;

  HistogramWriter():
    m_Metric(0)
    {
    this->m_Filter = HistogramToImageFilterType::New();
    this->histogramWriter = HistogramWriterType::New();
    this->histogramWriter->SetInput( this->m_Filter->GetOutput() );

    std::string outputFileBase = "JointHistogram"; // Base of series filenames ( of the joint histogram )
    this->outputFile = outputFileBase + "%03d.";
    this->outputFile += "mhd";   // histogram filename extension
    }

  ~HistogramWriter() { };

  void SetMetric( MetricPointer metric )
    {
    this->m_Metric = metric;
    }

  MetricPointer GetMetric(  )
    {
    return this->m_Metric;
    }

  void WriteHistogramFile( unsigned int iterationNumber )
    {
    char outputFilename[1000];
    sprintf (outputFilename, this->outputFile.c_str(), iterationNumber );

    histogramWriter->SetFileName( outputFilename );
    this->m_Filter->SetInput( m_Metric->GetHistogram() );
    try
      {
      m_Filter->Update();
      }
    catch( itk::ExceptionObject & err )
      {
      std::cerr << "ERROR: ExceptionObject caught !" << std::endl;
      std::cerr << err << std::endl;
      }

    try
      {
      histogramWriter->Update();
      }
    catch( itk::ExceptionObject & excp )
      {
      std::cerr << "Exception thrown " << excp << std::endl;
      }

    std::cout << "Joint Histogram file: " << outputFilename <<
        " written" << std::endl;
    }

private:
  MetricPointer m_Metric;
  HistogramToImageFilterPointer m_Filter;
  HistogramWriterPointer histogramWriter;
  std::string   outputFile;
} jointHistogramWriter;

class CommandIterationUpdate : public itk::Command
{
public:
  typedef  CommandIterationUpdate   Self;
  typedef  itk::Command             Superclass;
  typedef  itk::SmartPointer<Self>  Pointer;
  itkNewMacro( Self );
protected:
  CommandIterationUpdate() {};
public:

  typedef   itk::GradientDescentOptimizer     OptimizerType;
  typedef   const OptimizerType   *           OptimizerPointer;

  void Execute(itk::Object *caller, const itk::EventObject & event)
    {
      Execute( (const itk::Object *)caller, event);
    }

  void Execute(const itk::Object * object, const itk::EventObject & event)
    {
      OptimizerPointer optimizer =
        dynamic_cast< OptimizerPointer >( object );
      if( typeid( event ) != typeid( itk::IterationEvent ) )
        {
        return;
        }
      std::cout << optimizer->GetCurrentIteration() << "   ";
      std::cout << optimizer->GetValue() << "   ";
      std::cout << optimizer->GetCurrentPosition() << std::endl;

      // Write the joint histogram as a file JointHistogramXXX.mhd
      // where XXX is the iteration number
      jointHistogramWriter.WriteHistogramFile( optimizer->GetCurrentIteration() );
    }
};

int main( int argc, char *argv[] )
{
  if( argc < 4 )
    {
    std::cerr << "Missing Parameters " << std::endl;
    std::cerr << "Usage: " << argv[0];
    std::cerr << " fixedImageFile  movingImageFile ";
    std::cerr << "outputImagefile ";
    std::cerr << "Number of histogram bins for writing the MutualInformationHistogramMetric";
    std::cerr <<  std::endl;
    std::cerr << "Example: ImageRegistrationHistogramPlotter BrainT1Slice.png BrainProtonDensitySliceShifted13x17y.png Output.png 32" << std::endl;
    return 1;
    }

  typedef  unsigned short  PixelType;

  typedef itk::Image< PixelType, Dimension >  FixedImageType;
  typedef itk::Image< PixelType, Dimension >  MovingImageType;
  typedef   float     InternalPixelType;
  typedef itk::Image< InternalPixelType, Dimension > InternalImageType;

  typedef itk::TranslationTransform< double, Dimension > TransformType;
  typedef itk::GradientDescentOptimizer                  OptimizerType;
  typedef itk::LinearInterpolateImageFunction<
                                    InternalImageType,
                                    double             > InterpolatorType;
  typedef itk::ImageRegistrationMethod<
                                    InternalImageType,
                                    InternalImageType >  RegistrationType;
  typedef itk::MutualInformationHistogramImageToImageMetric<
                                          InternalImageType,
                                          InternalImageType >    MetricType;

  TransformType::Pointer      transform     = TransformType::New();
  OptimizerType::Pointer      optimizer     = OptimizerType::New();
  InterpolatorType::Pointer   interpolator  = InterpolatorType::New();
  RegistrationType::Pointer   registration  = RegistrationType::New();
  MetricType::Pointer         metric        = MetricType::New();

  registration->SetOptimizer(     optimizer     );
  registration->SetTransform(     transform     );
  registration->SetInterpolator(  interpolator  );

  unsigned int numberOfHistogramBins = atoi( argv[4] );
  MetricType::HistogramType::SizeType histogramSize;
  histogramSize[0] = numberOfHistogramBins;
  histogramSize[1] = numberOfHistogramBins;
  metric->SetHistogramSize( histogramSize );
  const unsigned int numberOfParameters = transform->GetNumberOfParameters();
  typedef MetricType::ScalesType ScalesType;
  ScalesType scales( numberOfParameters );
  scales.Fill( 1.0 );
  metric->SetDerivativeStepLengthScales(scales);

  // Set the metric for the joint histogram writer
  jointHistogramWriter.SetMetric( metric );

  registration->SetMetric( metric  );

  typedef itk::ImageFileReader< FixedImageType  > FixedImageReaderType;
  typedef itk::ImageFileReader< MovingImageType > MovingImageReaderType;

  FixedImageReaderType::Pointer  fixedImageReader  = FixedImageReaderType::New();
  MovingImageReaderType::Pointer movingImageReader = MovingImageReaderType::New();

  fixedImageReader->SetFileName(  argv[1] );
  movingImageReader->SetFileName( argv[2] );

  typedef itk::NormalizeImageFilter<
                                FixedImageType,
                                InternalImageType
                                        > FixedNormalizeFilterType;

  typedef itk::NormalizeImageFilter<
                                MovingImageType,
                                InternalImageType
                                              > MovingNormalizeFilterType;

  FixedNormalizeFilterType::Pointer fixedNormalizer =
                                            FixedNormalizeFilterType::New();

  MovingNormalizeFilterType::Pointer movingNormalizer =
                                            MovingNormalizeFilterType::New();
  typedef itk::DiscreteGaussianImageFilter<
                                      InternalImageType,
                                      InternalImageType
                                                    > GaussianFilterType;

  GaussianFilterType::Pointer fixedSmoother  = GaussianFilterType::New();
  GaussianFilterType::Pointer movingSmoother = GaussianFilterType::New();

  fixedSmoother->SetVariance( 2.0 );
  movingSmoother->SetVariance( 2.0 );
  fixedNormalizer->SetInput(  fixedImageReader->GetOutput() );
  movingNormalizer->SetInput( movingImageReader->GetOutput() );

  fixedSmoother->SetInput( fixedNormalizer->GetOutput() );
  movingSmoother->SetInput( movingNormalizer->GetOutput() );

  registration->SetFixedImage(    fixedSmoother->GetOutput()    );
  registration->SetMovingImage(   movingSmoother->GetOutput()   );

  fixedNormalizer->Update();
  registration->SetFixedImageRegion(
       fixedNormalizer->GetOutput()->GetBufferedRegion() );

  typedef RegistrationType::ParametersType ParametersType;
  ParametersType initialParameters( transform->GetNumberOfParameters() );

  initialParameters[0] = 0.0;  // Initial offset in mm along X
  initialParameters[1] = 0.0;  // Initial offset in mm along Y

  registration->SetInitialTransformParameters( initialParameters );

  optimizer->SetLearningRate( 20.0 );
  optimizer->SetNumberOfIterations( 200 );
  optimizer->MaximizeOn();

  CommandIterationUpdate::Pointer observer = CommandIterationUpdate::New();
  optimizer->AddObserver( itk::IterationEvent(), observer );

  try
    {
    registration->StartRegistration();
    }
  catch( itk::ExceptionObject & err )
    {
    std::cout << "ExceptionObject caught !" << std::endl;
    std::cout << err << std::endl;
    return -1;
    }

  ParametersType finalParameters = registration->GetLastTransformParameters();

  double TranslationAlongX = finalParameters[0];
  double TranslationAlongY = finalParameters[1];

  unsigned int numberOfIterations = optimizer->GetCurrentIteration();

  double bestValue = optimizer->GetValue();

  std::cout << "Result = " << std::endl;
  std::cout << " Translation X = " << TranslationAlongX  << std::endl;
  std::cout << " Translation Y = " << TranslationAlongY  << std::endl;
  std::cout << " Iterations    = " << numberOfIterations << std::endl;
  std::cout << " Metric value  = " << bestValue          << std::endl;

  typedef itk::ResampleImageFilter<
                            MovingImageType,
                            FixedImageType >    ResampleFilterType;

  TransformType::Pointer finalTransform = TransformType::New();

  finalTransform->SetParameters( finalParameters );

  ResampleFilterType::Pointer resample = ResampleFilterType::New();

  resample->SetTransform( finalTransform );
  resample->SetInput( movingImageReader->GetOutput() );

  FixedImageType::Pointer fixedImage = fixedImageReader->GetOutput();

  resample->SetSize(    fixedImage->GetLargestPossibleRegion().GetSize() );
  resample->SetOutputOrigin(  fixedImage->GetOrigin() );
  resample->SetOutputSpacing( fixedImage->GetSpacing() );
  resample->SetDefaultPixelValue( 100 );

  typedef  unsigned char  OutputPixelType;

  typedef itk::Image< OutputPixelType, Dimension > OutputImageType;

  typedef itk::CastImageFilter<
                        FixedImageType,
                        OutputImageType > CastFilterType;

  typedef itk::ImageFileWriter< OutputImageType >  WriterType;

  WriterType::Pointer      writer =  WriterType::New();
  CastFilterType::Pointer  caster =  CastFilterType::New();

  writer->SetFileName( argv[3] );

  caster->SetInput( resample->GetOutput() );
  writer->SetInput( caster->GetOutput()   );
  writer->Update();

  return EXIT_SUCCESS;
}

<commit_msg>ENH: Example showing the use of HistogramToEntropyImageFilter<commit_after>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    ImageRegistrationHistogramPlotter.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

// The example shows how to use the \doxygen{HistogramToEntropyImageFilter} class.
// The example registers two images using the gradient descent optimizer.
// The transform used here is a simple translation transform. The metric
// is a MutualInformationHistogramImageToImageMetric.
//
// The jointHistogramWriter is invoked after every iteration of the optimizer.
// the writer here writes the joint histogram after every iteration as
// JointHistogramXXX.mhd. The output image contains the joint entropy histogram
// given by
// \begin{equation}
// f(I) = -p \log_2 p
// \end{equation}
//
// where
// \begin{equation}
// p = \frac{q_I}{\sum_{i \in I} q_I}
// \end{equation}
//  and $q_I$ is the frequency of measurement vector, I.
//
// p is the probability of the occurance of the measurement vector, \f$q_I\f$.
//
// The output image is of type double.
//
// You may similarly instantiate \doxygen{HistogramToIntensityImageFilter}
// or \doxygen{HistogramToProbabilityImageFilter} or
// \doxygen{HistogramToLogProbabilityImageFilter}.
// \TODO: Put SoftwareGuide comments.

#if defined(_MSC_VER)
#pragma warning ( disable : 4786 )
#endif

#include "itkImageRegistrationMethod.h"
#include "itkTranslationTransform.h"
#include "itkMutualInformationHistogramImageToImageMetric.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkGradientDescentOptimizer.h"
#include "itkImage.h"
#include "itkNormalizeImageFilter.h"
#include "itkDiscreteGaussianImageFilter.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkResampleImageFilter.h"
#include "itkCastImageFilter.h"
#include "itkHistogramToEntropyImageFilter.h"
#include "itkCommand.h"

const    unsigned int    Dimension = 2;

class HistogramWriter
{
public:
  typedef float InternalPixelType;
  typedef itk::Image< InternalPixelType, Dimension > InternalImageType;
  typedef itk::MutualInformationHistogramImageToImageMetric<
                                          InternalImageType,
                                          InternalImageType >    MetricType;
  typedef MetricType::Pointer   MetricPointer;
  typedef MetricType::HistogramType   HistogramType;
  typedef itk::HistogramToEntropyImageFilter< HistogramType > HistogramToEntropyImageFilterType;
  typedef HistogramToEntropyImageFilterType::Pointer   HistogramToImageFilterPointer;
  typedef itk::ImageFileWriter<  HistogramToEntropyImageFilterType::OutputImageType  > HistogramWriterType;
  typedef HistogramWriterType::Pointer   HistogramWriterPointer;

  HistogramWriter():
    m_Metric(0)
    {
    this->m_Filter = HistogramToEntropyImageFilterType::New();
    this->histogramWriter = HistogramWriterType::New();
    this->histogramWriter->SetInput( this->m_Filter->GetOutput() );

    std::string outputFileBase = "JointHistogram"; // Base of series filenames ( of the joint histogram )
    this->outputFile = outputFileBase + "%03d.";
    this->outputFile += "mhd";   // histogram filename extension
    }

  ~HistogramWriter() { };

  void SetMetric( MetricPointer metric )
    {
    this->m_Metric = metric;
    }

  MetricPointer GetMetric(  )
    {
    return this->m_Metric;
    }

  void WriteHistogramFile( unsigned int iterationNumber )
    {
    char outputFilename[1000];
    sprintf (outputFilename, this->outputFile.c_str(), iterationNumber );

    histogramWriter->SetFileName( outputFilename );
    this->m_Filter->SetInput( m_Metric->GetHistogram() );
    try
      {
      m_Filter->Update();
      }
    catch( itk::ExceptionObject & err )
      {
      std::cerr << "ERROR: ExceptionObject caught !" << std::endl;
      std::cerr << err << std::endl;
      }

    try
      {
      histogramWriter->Update();
      }
    catch( itk::ExceptionObject & excp )
      {
      std::cerr << "Exception thrown " << excp << std::endl;
      }

    std::cout << "Joint Histogram file: " << outputFilename <<
        " written" << std::endl;
    }

private:
  MetricPointer m_Metric;
  HistogramToImageFilterPointer m_Filter;
  HistogramWriterPointer histogramWriter;
  std::string   outputFile;
} jointHistogramWriter;

class CommandIterationUpdate : public itk::Command
{
public:
  typedef  CommandIterationUpdate   Self;
  typedef  itk::Command             Superclass;
  typedef  itk::SmartPointer<Self>  Pointer;
  itkNewMacro( Self );
protected:
  CommandIterationUpdate() {};
public:

  typedef   itk::GradientDescentOptimizer     OptimizerType;
  typedef   const OptimizerType   *           OptimizerPointer;

  void Execute(itk::Object *caller, const itk::EventObject & event)
    {
      Execute( (const itk::Object *)caller, event);
    }

  void Execute(const itk::Object * object, const itk::EventObject & event)
    {
      OptimizerPointer optimizer =
        dynamic_cast< OptimizerPointer >( object );
      if( typeid( event ) != typeid( itk::IterationEvent ) )
        {
        return;
        }
      std::cout << optimizer->GetCurrentIteration() << "   ";
      std::cout << optimizer->GetValue() << "   ";
      std::cout << optimizer->GetCurrentPosition() << std::endl;

      // Write the joint histogram as a file JointHistogramXXX.mhd
      // where XXX is the iteration number
      jointHistogramWriter.WriteHistogramFile( optimizer->GetCurrentIteration() );
    }
};

int main( int argc, char *argv[] )
{
  if( argc < 3 )
    {
    std::cerr << "Missing Parameters " << std::endl;
    std::cerr << "Usage: " << argv[0];
    std::cerr << " fixedImageFile  movingImageFile ";
    std::cerr << "outputImagefile ";
    std::cerr << "Number of histogram bins for writing the MutualInformationHistogramMetric";
    std::cerr <<  std::endl;

    return 1;
    }

  typedef  unsigned short  PixelType;

  typedef itk::Image< PixelType, Dimension >  FixedImageType;
  typedef itk::Image< PixelType, Dimension >  MovingImageType;
  typedef   float     InternalPixelType;
  typedef itk::Image< InternalPixelType, Dimension > InternalImageType;

  typedef itk::TranslationTransform< double, Dimension > TransformType;
  typedef itk::GradientDescentOptimizer                  OptimizerType;
  typedef itk::LinearInterpolateImageFunction<
                                    InternalImageType,
                                    double             > InterpolatorType;
  typedef itk::ImageRegistrationMethod<
                                    InternalImageType,
                                    InternalImageType >  RegistrationType;
  typedef itk::MutualInformationHistogramImageToImageMetric<
                                          InternalImageType,
                                          InternalImageType >    MetricType;

  TransformType::Pointer      transform     = TransformType::New();
  OptimizerType::Pointer      optimizer     = OptimizerType::New();
  InterpolatorType::Pointer   interpolator  = InterpolatorType::New();
  RegistrationType::Pointer   registration  = RegistrationType::New();
  MetricType::Pointer         metric        = MetricType::New();

  registration->SetOptimizer(     optimizer     );
  registration->SetTransform(     transform     );
  registration->SetInterpolator(  interpolator  );

  unsigned int numberOfHistogramBins = atoi( argv[4] );
  MetricType::HistogramType::SizeType histogramSize;
  histogramSize[0] = numberOfHistogramBins;
  histogramSize[1] = numberOfHistogramBins;
  metric->SetHistogramSize( histogramSize );
  const unsigned int numberOfParameters = transform->GetNumberOfParameters();
  typedef MetricType::ScalesType ScalesType;
  ScalesType scales( numberOfParameters );
  scales.Fill( 1.0 );
  metric->SetDerivativeStepLengthScales(scales);

  // Set the metric for the joint histogram writer
  jointHistogramWriter.SetMetric( metric );

  registration->SetMetric( metric  );

  typedef itk::ImageFileReader< FixedImageType  > FixedImageReaderType;
  typedef itk::ImageFileReader< MovingImageType > MovingImageReaderType;

  FixedImageReaderType::Pointer  fixedImageReader  = FixedImageReaderType::New();
  MovingImageReaderType::Pointer movingImageReader = MovingImageReaderType::New();

  fixedImageReader->SetFileName(  argv[1] );
  movingImageReader->SetFileName( argv[2] );

  typedef itk::NormalizeImageFilter<
                                FixedImageType,
                                InternalImageType
                                        > FixedNormalizeFilterType;

  typedef itk::NormalizeImageFilter<
                                MovingImageType,
                                InternalImageType
                                              > MovingNormalizeFilterType;

  FixedNormalizeFilterType::Pointer fixedNormalizer =
                                            FixedNormalizeFilterType::New();

  MovingNormalizeFilterType::Pointer movingNormalizer =
                                            MovingNormalizeFilterType::New();
  typedef itk::DiscreteGaussianImageFilter<
                                      InternalImageType,
                                      InternalImageType
                                                    > GaussianFilterType;

  GaussianFilterType::Pointer fixedSmoother  = GaussianFilterType::New();
  GaussianFilterType::Pointer movingSmoother = GaussianFilterType::New();

  fixedSmoother->SetVariance( 2.0 );
  movingSmoother->SetVariance( 2.0 );
  fixedNormalizer->SetInput(  fixedImageReader->GetOutput() );
  movingNormalizer->SetInput( movingImageReader->GetOutput() );

  fixedSmoother->SetInput( fixedNormalizer->GetOutput() );
  movingSmoother->SetInput( movingNormalizer->GetOutput() );

  registration->SetFixedImage(    fixedSmoother->GetOutput()    );
  registration->SetMovingImage(   movingSmoother->GetOutput()   );

  fixedNormalizer->Update();
  registration->SetFixedImageRegion(
       fixedNormalizer->GetOutput()->GetBufferedRegion() );

  typedef RegistrationType::ParametersType ParametersType;
  ParametersType initialParameters( transform->GetNumberOfParameters() );

  initialParameters[0] = 0.0;  // Initial offset in mm along X
  initialParameters[1] = 0.0;  // Initial offset in mm along Y

  registration->SetInitialTransformParameters( initialParameters );

  optimizer->SetLearningRate( 20.0 );
  optimizer->SetNumberOfIterations( 200 );
  optimizer->MaximizeOn();

  CommandIterationUpdate::Pointer observer = CommandIterationUpdate::New();
  optimizer->AddObserver( itk::IterationEvent(), observer );

  try
    {
    registration->StartRegistration();
    }
  catch( itk::ExceptionObject & err )
    {
    std::cout << "ExceptionObject caught !" << std::endl;
    std::cout << err << std::endl;
    return -1;
    }

  ParametersType finalParameters = registration->GetLastTransformParameters();

  double TranslationAlongX = finalParameters[0];
  double TranslationAlongY = finalParameters[1];

  unsigned int numberOfIterations = optimizer->GetCurrentIteration();

  double bestValue = optimizer->GetValue();

  std::cout << "Result = " << std::endl;
  std::cout << " Translation X = " << TranslationAlongX  << std::endl;
  std::cout << " Translation Y = " << TranslationAlongY  << std::endl;
  std::cout << " Iterations    = " << numberOfIterations << std::endl;
  std::cout << " Metric value  = " << bestValue          << std::endl;

  typedef itk::ResampleImageFilter<
                            MovingImageType,
                            FixedImageType >    ResampleFilterType;

  TransformType::Pointer finalTransform = TransformType::New();

  finalTransform->SetParameters( finalParameters );

  ResampleFilterType::Pointer resample = ResampleFilterType::New();

  resample->SetTransform( finalTransform );
  resample->SetInput( movingImageReader->GetOutput() );

  FixedImageType::Pointer fixedImage = fixedImageReader->GetOutput();

  resample->SetSize(    fixedImage->GetLargestPossibleRegion().GetSize() );
  resample->SetOutputOrigin(  fixedImage->GetOrigin() );
  resample->SetOutputSpacing( fixedImage->GetSpacing() );
  resample->SetDefaultPixelValue( 100 );

  typedef  unsigned char  OutputPixelType;

  typedef itk::Image< OutputPixelType, Dimension > OutputImageType;

  typedef itk::CastImageFilter<
                        FixedImageType,
                        OutputImageType > CastFilterType;

  typedef itk::ImageFileWriter< OutputImageType >  WriterType;

  WriterType::Pointer      writer =  WriterType::New();
  CastFilterType::Pointer  caster =  CastFilterType::New();

  writer->SetFileName( argv[3] );

  caster->SetInput( resample->GetOutput() );
  writer->SetInput( caster->GetOutput()   );
  writer->Update();

  return EXIT_SUCCESS;
}

<|endoftext|>


<commit_before>#include <GarrysMod/Lua/Interface.h>
#include <GarrysMod/Lua/LuaInterface.h>
#include <SymbolFinder.hpp>
#include <MologieDetours/detours.h>
#include <stdexcept>
#include <stdint.h>
#include <stdio.h>
#include <sha1.h>

#if defined _WIN32

#define snprintf _snprintf

#define FASTCALL __fastcall
#define THISCALL __thiscall

#define SERVER_BINARY "server.dll"

#define ADDORUPDATEFILE_SYM reinterpret_cast<const uint8_t *>( "\x55\x8B\xEC\x83\xEC\x18\x53\x56\x8B\x75\x08\x83\xC6\x04\x83\x7E" )
#define ADDORUPDATEFILE_SYMLEN 16

#elif defined __linux || defined __APPLE__

#define CDECL __attribute__((cdecl))

#if defined __linux

#define SERVER_BINARY "garrysmod/bin/server_srv.so"

#define SYMBOL_PREFIX "@"

#else

#define SERVER_BINARY "garrysmod/bin/server.dylib"

#define SYMBOL_PREFIX "@_"

#endif

#define ADDORUPDATEFILE_SYM reinterpret_cast<const uint8_t *>( SYMBOL_PREFIX "_ZN12GModDataPack15AddOrUpdateFileEP7LuaFileb" )
#define ADDORUPDATEFILE_SYMLEN 0

#endif

GarrysMod::Lua::ILuaInterface *lua = NULL;

class GModDataPack;

struct LuaFile
{
	uint32_t skip;
	const char *path;
};

#if defined _WIN32

typedef void ( THISCALL *AddOrUpdateFile_t ) ( GModDataPack *self, LuaFile *file, bool force );

#elif defined __linux || defined __APPLE__

typedef void ( CDECL *AddOrUpdateFile_t ) ( GModDataPack *self, LuaFile *file, bool force );

#endif

AddOrUpdateFile_t AddOrUpdateFile = NULL;
MologieDetours::Detour<AddOrUpdateFile_t> *AddOrUpdateFile_d = NULL;

#if defined _WIN32

void FASTCALL AddOrUpdateFile_h( GModDataPack *self, void *, LuaFile *file, bool b )

#elif defined __linux || defined __APPLE__

void CDECL AddOrUpdateFile_h( GModDataPack *self, LuaFile *file, bool b )

#endif

{
	lua->PushSpecial( GarrysMod::Lua::SPECIAL_GLOB );
	if( lua->IsType( -1, GarrysMod::Lua::Type::TABLE ) )
	{
		lua->GetField( -1, "hook" );
		if( lua->IsType( -1, GarrysMod::Lua::Type::TABLE ) )
		{
			lua->GetField( -1, "Run" );
			if( lua->IsType( -1, GarrysMod::Lua::Type::FUNCTION ) )
			{
				lua->PushString( "AddOrUpdateCSLuaFile" );

				lua->PushString( file->path );
				lua->PushBool( b );

				if( lua->PCall( 3, 1, 0 ) != 0 )
				{
					lua->Msg( "[luapack_internal] %s\n", lua->GetString( ) );
					lua->Pop( 3 );
					return AddOrUpdateFile( self, file, b );
				}

				if( lua->IsType( -1, GarrysMod::Lua::Type::BOOL ) && lua->GetBool( ) )
				{
					lua->Pop( 3 );
					return;
				}

				lua->Pop( 1 );
			}
			else
			{
				lua->Pop( 1 );
			}
		}

		lua->Pop( 1 );
	}

	lua->Pop( 1 );
	return AddOrUpdateFile( self, file, b );
}

#define HASHER_METATABLE "SHA1"
#define HASHER_TYPE 31

#define THROW_ERROR( error ) ( LUA->ThrowError( error ), 0 )
#define LUA_ERROR( ) THROW_ERROR( LUA->GetString( ) )

#define GET_USERDATA( index ) reinterpret_cast<GarrysMod::Lua::UserData *>( LUA->GetUserdata( index ) )
#define GET_HASHER( index ) reinterpret_cast<sha1_context *>( GET_USERDATA( index )->data )
#define VALIDATE_HASHER( hasher ) if( hasher == 0 ) return THROW_ERROR( HASHER_METATABLE " object is not valid" )

LUA_FUNCTION_STATIC( hasher__new )
{
	try
	{
		sha1_context *context = new sha1_context;
		sha1_starts( context );

		void *luadata = LUA->NewUserdata( sizeof( GarrysMod::Lua::UserData ) );
		GarrysMod::Lua::UserData *userdata = reinterpret_cast<GarrysMod::Lua::UserData *>( luadata );
		userdata->data = context;
		userdata->type = HASHER_TYPE;

		LUA->CreateMetaTableType( HASHER_METATABLE, HASHER_TYPE );
		LUA->SetMetaTable( -2 );

		return 1;
	}
	catch( std::exception &e )
	{
		LUA->PushString( e.what( ) );
	}

	return LUA_ERROR( );
}

LUA_FUNCTION_STATIC( hasher__gc )
{
	LUA->CheckType( 1, HASHER_TYPE );

	GarrysMod::Lua::UserData *userdata = GET_USERDATA( 1 );
	sha1_context *hasher = reinterpret_cast<sha1_context *>( userdata->data );
	VALIDATE_HASHER( hasher );

	userdata->data = 0;

	delete hasher;
	return 0;
}

LUA_FUNCTION_STATIC( hasher_update )
{
	LUA->CheckType( 1, HASHER_TYPE );
	LUA->CheckType( 2, GarrysMod::Lua::Type::STRING );

	sha1_context *hasher = GET_HASHER( 1 );
	VALIDATE_HASHER( hasher );

	uint32_t len = 0;
	const uint8_t *data = reinterpret_cast<const uint8_t *>( LUA->GetString( 2, &len ) );

	sha1_update( hasher, data, len );

	return 0;
}

LUA_FUNCTION_STATIC( hasher_final )
{
	LUA->CheckType( 1, HASHER_TYPE );

	sha1_context *hasher = GET_HASHER( 1 );
	VALIDATE_HASHER( hasher );

	uint8_t digest[20];
	sha1_finish( hasher, digest );

	LUA->PushString( reinterpret_cast<const char *>( digest ), sizeof( digest ) );
	return 1;
}

LUA_FUNCTION_STATIC( luapack_rename )
{
	LUA->CheckType( 1, GarrysMod::Lua::Type::STRING );

	char newto[70];
	snprintf( newto, sizeof( newto ), "garrysmod/data/luapack/%s.dat", LUA->GetString( 1 ) );

	LUA->PushBool( rename( "garrysmod/data/luapack/temp.dat", newto ) == 0 );
	return 1;
}

GMOD_MODULE_OPEN( )
{
	lua = reinterpret_cast<GarrysMod::Lua::ILuaInterface *>( LUA );

	try
	{
		lua->PushSpecial( GarrysMod::Lua::SPECIAL_GLOB );

		lua->GetField( -1, "luapack" );
		if( !lua->IsType( -1, GarrysMod::Lua::Type::TABLE ) )
			throw std::runtime_error( "luapack table not found" );

		lua->PushCFunction( luapack_rename );
		lua->SetField( -2, "Rename" );

		lua->PushCFunction( hasher__new );
		lua->SetField( -2, HASHER_METATABLE );

		LUA->CreateMetaTableType( HASHER_METATABLE, HASHER_TYPE );

		LUA->Push( -1 );
		LUA->SetField( -2, "__index" );

		LUA->PushCFunction( hasher__gc );
		LUA->SetField( -2, "__gc" );

		LUA->PushCFunction( hasher_update );
		LUA->SetField( -2, "Update" );

		LUA->PushCFunction( hasher_final );
		LUA->SetField( -2, "Final" );

		SymbolFinder symfinder;

		AddOrUpdateFile = reinterpret_cast<AddOrUpdateFile_t>( symfinder.ResolveOnBinary( SERVER_BINARY, ADDORUPDATEFILE_SYM, ADDORUPDATEFILE_SYMLEN ) );
		if( AddOrUpdateFile == NULL )
			throw std::runtime_error( "GModDataPack::AddOrUpdateFile detour failed" );

		AddOrUpdateFile_d = new MologieDetours::Detour<AddOrUpdateFile_t>( AddOrUpdateFile, reinterpret_cast<AddOrUpdateFile_t>( AddOrUpdateFile_h ) );

		AddOrUpdateFile = AddOrUpdateFile_d->GetOriginalFunction( );

		lua->Msg( "[luapack_internal] GModDataPack::AddOrUpdateFile detoured.\n" );

		return 0;
	}
	catch( std::exception &e )
	{
		LUA->PushString( e.what( ) );
	}

	return LUA_ERROR( );
}

GMOD_MODULE_CLOSE( )
{
	delete AddOrUpdateFile_d;
	return 0;
}<commit_msg>luapack.Rename upgraded. Now accepts 2 paths relative to garrysmod. Should still be pretty safe since all ../ are removed and the only accepted extensions are .txt, .dat and .lua<commit_after>#include <GarrysMod/Lua/Interface.h>
#include <GarrysMod/Lua/LuaInterface.h>
#include <SymbolFinder.hpp>
#include <MologieDetours/detours.h>
#include <stdexcept>
#include <string>
#include <stdint.h>
#include <stdio.h>
#include <sha1.h>

#if defined _WIN32

#define snprintf _snprintf

#define FASTCALL __fastcall
#define THISCALL __thiscall

#define SERVER_BINARY "server.dll"

#define ADDORUPDATEFILE_SYM reinterpret_cast<const uint8_t *>( "\x55\x8B\xEC\x83\xEC\x18\x53\x56\x8B\x75\x08\x83\xC6\x04\x83\x7E" )
#define ADDORUPDATEFILE_SYMLEN 16

#elif defined __linux || defined __APPLE__

#define CDECL __attribute__((cdecl))

#if defined __linux

#define SERVER_BINARY "garrysmod/bin/server_srv.so"

#define SYMBOL_PREFIX "@"

#else

#define SERVER_BINARY "garrysmod/bin/server.dylib"

#define SYMBOL_PREFIX "@_"

#endif

#define ADDORUPDATEFILE_SYM reinterpret_cast<const uint8_t *>( SYMBOL_PREFIX "_ZN12GModDataPack15AddOrUpdateFileEP7LuaFileb" )
#define ADDORUPDATEFILE_SYMLEN 0

#endif

GarrysMod::Lua::ILuaInterface *lua = NULL;

class GModDataPack;

struct LuaFile
{
	uint32_t skip;
	const char *path;
};

#if defined _WIN32

typedef void ( THISCALL *AddOrUpdateFile_t ) ( GModDataPack *self, LuaFile *file, bool force );

#elif defined __linux || defined __APPLE__

typedef void ( CDECL *AddOrUpdateFile_t ) ( GModDataPack *self, LuaFile *file, bool force );

#endif

AddOrUpdateFile_t AddOrUpdateFile = NULL;
MologieDetours::Detour<AddOrUpdateFile_t> *AddOrUpdateFile_d = NULL;

#if defined _WIN32

void FASTCALL AddOrUpdateFile_h( GModDataPack *self, void *, LuaFile *file, bool b )

#elif defined __linux || defined __APPLE__

void CDECL AddOrUpdateFile_h( GModDataPack *self, LuaFile *file, bool b )

#endif

{
	lua->PushSpecial( GarrysMod::Lua::SPECIAL_GLOB );
	if( lua->IsType( -1, GarrysMod::Lua::Type::TABLE ) )
	{
		lua->GetField( -1, "hook" );
		if( lua->IsType( -1, GarrysMod::Lua::Type::TABLE ) )
		{
			lua->GetField( -1, "Run" );
			if( lua->IsType( -1, GarrysMod::Lua::Type::FUNCTION ) )
			{
				lua->PushString( "AddOrUpdateCSLuaFile" );

				lua->PushString( file->path );
				lua->PushBool( b );

				if( lua->PCall( 3, 1, 0 ) != 0 )
				{
					lua->Msg( "[luapack_internal] %s\n", lua->GetString( ) );
					lua->Pop( 3 );
					return AddOrUpdateFile( self, file, b );
				}

				if( lua->IsType( -1, GarrysMod::Lua::Type::BOOL ) && lua->GetBool( ) )
				{
					lua->Pop( 3 );
					return;
				}

				lua->Pop( 1 );
			}
			else
			{
				lua->Pop( 1 );
			}
		}

		lua->Pop( 1 );
	}

	lua->Pop( 1 );
	return AddOrUpdateFile( self, file, b );
}

#define HASHER_METATABLE "SHA1"
#define HASHER_TYPE 31

#define THROW_ERROR( error ) ( LUA->ThrowError( error ), 0 )
#define LUA_ERROR( ) THROW_ERROR( LUA->GetString( ) )

#define GET_USERDATA( index ) reinterpret_cast<GarrysMod::Lua::UserData *>( LUA->GetUserdata( index ) )
#define GET_HASHER( index ) reinterpret_cast<sha1_context *>( GET_USERDATA( index )->data )
#define VALIDATE_HASHER( hasher ) if( hasher == 0 ) return THROW_ERROR( HASHER_METATABLE " object is not valid" )

LUA_FUNCTION_STATIC( hasher__new )
{
	try
	{
		sha1_context *context = new sha1_context;
		sha1_starts( context );

		void *luadata = LUA->NewUserdata( sizeof( GarrysMod::Lua::UserData ) );
		GarrysMod::Lua::UserData *userdata = reinterpret_cast<GarrysMod::Lua::UserData *>( luadata );
		userdata->data = context;
		userdata->type = HASHER_TYPE;

		LUA->CreateMetaTableType( HASHER_METATABLE, HASHER_TYPE );
		LUA->SetMetaTable( -2 );

		return 1;
	}
	catch( std::exception &e )
	{
		LUA->PushString( e.what( ) );
	}

	return LUA_ERROR( );
}

LUA_FUNCTION_STATIC( hasher__gc )
{
	LUA->CheckType( 1, HASHER_TYPE );

	GarrysMod::Lua::UserData *userdata = GET_USERDATA( 1 );
	sha1_context *hasher = reinterpret_cast<sha1_context *>( userdata->data );
	VALIDATE_HASHER( hasher );

	userdata->data = 0;

	delete hasher;
	return 0;
}

LUA_FUNCTION_STATIC( hasher_update )
{
	LUA->CheckType( 1, HASHER_TYPE );
	LUA->CheckType( 2, GarrysMod::Lua::Type::STRING );

	sha1_context *hasher = GET_HASHER( 1 );
	VALIDATE_HASHER( hasher );

	uint32_t len = 0;
	const uint8_t *data = reinterpret_cast<const uint8_t *>( LUA->GetString( 2, &len ) );

	sha1_update( hasher, data, len );

	return 0;
}

LUA_FUNCTION_STATIC( hasher_final )
{
	LUA->CheckType( 1, HASHER_TYPE );

	sha1_context *hasher = GET_HASHER( 1 );
	VALIDATE_HASHER( hasher );

	uint8_t digest[20];
	sha1_finish( hasher, digest );

	LUA->PushString( reinterpret_cast<const char *>( digest ), sizeof( digest ) );
	return 1;
}

static void SubstituteChar( std::string &path, char part, char sub )
{
	size_t pos = path.find( part );
	while( pos != path.npos )
	{
		path.erase( pos, 1 );
		path.insert( pos, 1, sub );
		pos = path.find( part, pos + 1 );
	}
}

static void RemovePart( std::string &path, const char *part )
{
	size_t len = strlen( part ), pos = path.find( part );
	while( pos != path.npos )
	{
		path.erase( pos, len );
		pos = path.find( part, pos );
	}
}

inline void FixPath( std::string &path )
{
	SubstituteChar( path, '\\', '/' );
	RemovePart( path, "../" );
	RemovePart( path, "./" );
}

static bool HasWhitelistedExtension( const char *cpath )
{
	std::string path = cpath;
	size_t extstart = path.find( '.' );
	if( extstart != path.npos )
	{
		size_t lastslash = path.find( '/' );
		if( lastslash != path.npos && lastslash > extstart )
			return false;

		std::string ext = path.substr( extstart + 1 );
		return ext == "txt" || ext == "dat" || ext == "lua";
	}

	return false;
}

LUA_FUNCTION_STATIC( luapack_rename )
{
	LUA->CheckType( 1, GarrysMod::Lua::Type::STRING );
	LUA->CheckType( 2, GarrysMod::Lua::Type::STRING );

	const char *luafrom = LUA->GetString( 1 );
	if( HasWhitelistedExtension( luafrom ) )
	{
		LUA->PushBool( false );
		return 1;
	}

	const char *luato = LUA->GetString( 2 );
	if( HasWhitelistedExtension( luato ) )
	{
		LUA->PushBool( false );
		return 1;
	}

	bool success = false;

	// Lua (LuaJIT in x86 in particular) doesn't exactly like RAII.
	// Let's make a specific scope just for our logic.
	{
		std::string from = "garrysmod/";
		from += luafrom;
		FixPath( from );

		std::string to = "garrysmod/";
		to += luato;
		FixPath( to );

		success = rename( from.c_str( ), to.c_str( ) ) == 0;
	}

	LUA->PushBool( success );
	return 1;
}

GMOD_MODULE_OPEN( )
{
	lua = reinterpret_cast<GarrysMod::Lua::ILuaInterface *>( LUA );

	try
	{
		lua->PushSpecial( GarrysMod::Lua::SPECIAL_GLOB );

		lua->GetField( -1, "luapack" );
		if( !lua->IsType( -1, GarrysMod::Lua::Type::TABLE ) )
			throw std::runtime_error( "luapack table not found" );

		lua->PushCFunction( luapack_rename );
		lua->SetField( -2, "Rename" );

		lua->PushCFunction( hasher__new );
		lua->SetField( -2, HASHER_METATABLE );

		LUA->CreateMetaTableType( HASHER_METATABLE, HASHER_TYPE );

		LUA->Push( -1 );
		LUA->SetField( -2, "__index" );

		LUA->PushCFunction( hasher__gc );
		LUA->SetField( -2, "__gc" );

		LUA->PushCFunction( hasher_update );
		LUA->SetField( -2, "Update" );

		LUA->PushCFunction( hasher_final );
		LUA->SetField( -2, "Final" );

		SymbolFinder symfinder;

		AddOrUpdateFile = reinterpret_cast<AddOrUpdateFile_t>( symfinder.ResolveOnBinary( SERVER_BINARY, ADDORUPDATEFILE_SYM, ADDORUPDATEFILE_SYMLEN ) );
		if( AddOrUpdateFile == NULL )
			throw std::runtime_error( "GModDataPack::AddOrUpdateFile detour failed" );

		AddOrUpdateFile_d = new MologieDetours::Detour<AddOrUpdateFile_t>( AddOrUpdateFile, reinterpret_cast<AddOrUpdateFile_t>( AddOrUpdateFile_h ) );

		AddOrUpdateFile = AddOrUpdateFile_d->GetOriginalFunction( );

		lua->Msg( "[luapack_internal] GModDataPack::AddOrUpdateFile detoured.\n" );

		return 0;
	}
	catch( std::exception &e )
	{
		LUA->PushString( e.what( ) );
	}

	return LUA_ERROR( );
}

GMOD_MODULE_CLOSE( )
{
	delete AddOrUpdateFile_d;
	return 0;
}<|endoftext|>