UnrealEngine/Engine/Source/Developer/NaniteBuilder/Private/GraphPartitioner.h

// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "CoreMinimal.h"
#include "Math/Bounds.h"
#include "metis.h"
#include "Async/ParallelFor.h"
#include "DisjointSet.h"

class FGraphPartitioner
{
public:
	struct FGraphData
	{
		int32	Offset;
		int32	Num;

		TArray< idx_t >	Adjacency;
		TArray< idx_t >	AdjacencyCost;
		TArray< idx_t >	AdjacencyOffset;
	};

	// Inclusive
	struct FRange
	{
		uint32	Begin;
		uint32	End;

		bool operator<( const FRange& Other) const { return Begin < Other.Begin; }
	};
	TArray< FRange >	Ranges;
	TArray< uint32 >	Indexes;
	TArray< uint32 >	SortedTo;

public:
				FGraphPartitioner( uint32 InNumElements, int32 InMinPartitionSize, int32 InMaxPartitionSize );

	FGraphData*	NewGraph( uint32 NumAdjacency ) const;

	void		AddAdjacency( FGraphData* Graph, uint32 AdjIndex, idx_t Cost );
	void		AddLocalityLinks( FGraphData* Graph, uint32 Index, idx_t Cost );

	template< typename FGetCenter >
	void		BuildLocalityLinks( FDisjointSet& DisjointSet, const FBounds3f& Bounds, TConstArrayView< const int32 > GroupIndexes, FGetCenter& GetCenter );

	void		Partition( FGraphData* Graph );
	void		PartitionStrict( FGraphData* Graph, bool bThreaded );

private:
	void		BisectGraph( FGraphData* Graph, FGraphData* ChildGraphs[2] );
	void		RecursiveBisectGraph( FGraphData* Graph );

	uint32		NumElements;
	int32		MinPartitionSize;
	int32		MaxPartitionSize;

	TAtomic< uint32 >	NumPartitions;

	TArray< idx_t >		PartitionIDs;
	TArray< int32 >		SwappedWith;

	TMultiMap< uint32, uint32 >	LocalityLinks;
};

FORCEINLINE void FGraphPartitioner::AddAdjacency( FGraphData* Graph, uint32 AdjIndex, idx_t Cost )
{
	Graph->Adjacency.Add( SortedTo[ AdjIndex ] );
	Graph->AdjacencyCost.Add( Cost );
}

FORCEINLINE void FGraphPartitioner::AddLocalityLinks( FGraphData* Graph, uint32 Index, idx_t Cost )
{
	for( auto Iter = LocalityLinks.CreateKeyIterator( Index ); Iter; ++Iter )
	{
		uint32 AdjIndex = Iter.Value();
		Graph->Adjacency.Add( SortedTo[ AdjIndex ] );
		Graph->AdjacencyCost.Add( Cost );
	}
}

template< typename FGetCenter >
void FGraphPartitioner::BuildLocalityLinks( FDisjointSet& DisjointSet, const FBounds3f& Bounds, TConstArrayView< const int32 > GroupIndexes, FGetCenter& GetCenter )
{
	TArray< uint32 > SortKeys;
	SortKeys.AddUninitialized( NumElements );
	SortedTo.AddUninitialized( NumElements );

	const bool bElementGroups = !GroupIndexes.IsEmpty();	// Only create locality links between elements with the same group index

	ParallelFor( TEXT("BuildLocalityLinks.PF"), NumElements, 4096,
		[&]( uint32 Index )
		{
			FVector3f Center = GetCenter( Index );
			FVector3f CenterLocal = ( Center - Bounds.Min ) / FVector3f( Bounds.Max - Bounds.Min ).GetMax();

			uint32 Morton;
			Morton  = FMath::MortonCode3( uint32( CenterLocal.X * 1023 ) );
			Morton |= FMath::MortonCode3( uint32( CenterLocal.Y * 1023 ) ) << 1;
			Morton |= FMath::MortonCode3( uint32( CenterLocal.Z * 1023 ) ) << 2;
			SortKeys[ Index ] = Morton;
		} );

	RadixSort32( SortedTo.GetData(), Indexes.GetData(), NumElements,
		[&]( uint32 Index )
		{
			return SortKeys[ Index ];
		} );
	
	SortKeys.Empty();

	Swap( Indexes, SortedTo );
	for( uint32 i = 0; i < NumElements; i++ )
	{
		SortedTo[ Indexes[i] ] = i;
	}

	TArray< FRange > IslandRuns;
	IslandRuns.AddUninitialized( NumElements );
	
	// Run length acceleration
	// Range of identical IslandID denoting that elements are connected.
	// Used for jumping past connected elements to the next nearby disjoint element.
	{
		uint32 RunIslandID = 0;
		uint32 RunFirstElement = 0;

		for( uint32 i = 0; i < NumElements; i++ )
		{
			uint32 IslandID = DisjointSet.Find( Indexes[i] );

			if( RunIslandID != IslandID )
			{
				// We found the end so rewind to the beginning of the run and fill.
				for( uint32 j = RunFirstElement; j < i; j++ )
				{
					IslandRuns[j].End = i - 1;
				}
			
				// Start the next run
				RunIslandID = IslandID;
				RunFirstElement = i;
			}

			IslandRuns[i].Begin = RunFirstElement;
		}
		// Finish the last run
		for( uint32 j = RunFirstElement; j < NumElements; j++ )
		{
			IslandRuns[j].End = NumElements - 1;
		}
	}

	for( uint32 i = 0; i < NumElements; i++ )
	{
		uint32 Index = Indexes[i];

		uint32 RunLength = IslandRuns[i].End - IslandRuns[i].Begin + 1;
		if( RunLength < 128 )
		{
			uint32 IslandID = DisjointSet[ Index ];
			int32 GroupID = bElementGroups ? GroupIndexes[ Index ] : 0;

			FVector3f Center = GetCenter( Index );

			const uint32 MaxLinksPerElement = 5;

			uint32 ClosestIndex[MaxLinksPerElement];
			float  ClosestDist2[MaxLinksPerElement];
			for (int32 k = 0; k < MaxLinksPerElement; k++)
			{
				ClosestIndex[k] = ~0u;
				ClosestDist2[k] = MAX_flt;
			}

			for( int Direction = 0; Direction < 2; Direction++ )
			{
				uint32 Limit = Direction ? NumElements - 1 : 0;
				uint32 Step  = Direction ? 1 : -1;

				uint32 Adj = i;
				for( int32 Iterations = 0; Iterations < 16; Iterations++ )
				{
					if( Adj == Limit )
						break;
					Adj += Step;

					uint32 AdjIndex = Indexes[ Adj ];
					uint32 AdjIslandID = DisjointSet[ AdjIndex ];
					int32 AdjGroupID = bElementGroups ? GroupIndexes[AdjIndex] : 0;
					if( IslandID == AdjIslandID || ( GroupID != AdjGroupID ) )
					{
						// Skip past this run
						if( Direction )
							Adj = IslandRuns[ Adj ].End;
						else
							Adj = IslandRuns[ Adj ].Begin;
					}
					else
					{
						// Add to sorted list
						float AdjDist2 = ( Center - GetCenter( AdjIndex ) ).SizeSquared();
						for( int k = 0; k < MaxLinksPerElement; k++ )
						{
							if( AdjDist2 < ClosestDist2[k] )
							{
								Swap( AdjIndex, ClosestIndex[k] );
								Swap( AdjDist2, ClosestDist2[k] );
							}
						}
					}
				}
			}

			for( int k = 0; k < MaxLinksPerElement; k++ )
			{
				if( ClosestIndex[k] != ~0u )
				{
					// Add both directions
					LocalityLinks.AddUnique( Index, ClosestIndex[k] );
					LocalityLinks.AddUnique( ClosestIndex[k], Index );
				}
			}
		}
	}
}