using System;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Mathematics;
using UnityEngine.Assertions;
namespace UnityEngine.Rendering.Universal
{
struct BuddyAllocation
{
public int level;
public int index;
public BuddyAllocation(int level, int index)
{
this.level = level;
this.index = index;
}
public uint2 index2D => SpaceFillingCurves.DecodeMorton2D((uint)index);
}
unsafe struct BuddyAllocator : IDisposable
{
struct Header
{
public int branchingOrder;
public int levelCount;
public int allocationCount;
public int freeAllocationIdsCount;
}
// This data structure uses one big allocation containing a fixed header and all the arrays.
// Some arrays are sub-divided per order, which can be identified by the presence of an X(int order) method,
// which allows for easy access to the slice of data for the specified order.
// The offsets for the arrays are stored together with the pointer to avoid a double look-up when
// accessing data.
void* m_Data;
ref Header header => ref UnsafeUtility.AsRef<Header>(m_Data);
(int, int) m_ActiveFreeMaskCounts;
NativeArray<int> freeMaskCounts => GetNativeArray<int>(m_ActiveFreeMaskCounts.Item1, m_ActiveFreeMaskCounts.Item2);
(int, int) m_FreeMasksStorage;
NativeArray<ulong> freeMasksStorage => GetNativeArray<ulong>(m_FreeMasksStorage.Item1, m_FreeMasksStorage.Item2);
NativeArray<ulong> FreeMasks(int level) => freeMasksStorage.GetSubArray(LevelOffset64(level, header.branchingOrder), LevelLength64(level, header.branchingOrder));
(int, int) m_FreeMaskIndicesStorage;
NativeArray<int> freeMaskIndicesStorage => GetNativeArray<int>(m_FreeMaskIndicesStorage.Item1, m_FreeMaskIndicesStorage.Item2);
NativeArray<int> FreeMaskIndices(int level) => freeMaskIndicesStorage.GetSubArray(LevelOffset64(level, header.branchingOrder), LevelLength64(level, header.branchingOrder));
Allocator m_Allocator;
#if ENABLE_UNITY_COLLECTIONS_CHECKS
AtomicSafetyHandle m_SafetyHandle;
#endif
public int levelCount => header.levelCount;
public BuddyAllocator(int levelCount, int branchingOrder, Allocator allocator = Allocator.Persistent)
{
// Allows us to support 1D, 2D, and 3D cases.
Assert.IsTrue(branchingOrder is >= 1 and <= 3);
// Memory usage explodes unless capped like this.
Assert.IsTrue((levelCount + branchingOrder - 1) / branchingOrder is >= 1 and <= 24);
var dataSize = sizeof(Header);
m_ActiveFreeMaskCounts = AllocateRange<int>(levelCount, ref dataSize);
m_FreeMasksStorage = AllocateRange<ulong>(LevelOffset64(levelCount, branchingOrder), ref dataSize);
m_FreeMaskIndicesStorage = AllocateRange<int>(LevelOffset64(levelCount, branchingOrder), ref dataSize);
m_Data = UnsafeUtility.Malloc(dataSize, 64, allocator);
UnsafeUtility.MemClear(m_Data, dataSize);
m_Allocator = allocator;
#if ENABLE_UNITY_COLLECTIONS_CHECKS
m_SafetyHandle = AtomicSafetyHandle.Create();
#endif
header = new Header
{
branchingOrder = branchingOrder,
levelCount = levelCount
};
// Initialize level-0 to have 1/1 block available.
var freeMasks0 = FreeMasks(0);
freeMasks0[0] = 0b1111;
var maskCounts = freeMaskCounts;
maskCounts[0] = 1;
}
public bool TryAllocate(int requestedLevel, out BuddyAllocation allocation)
{
allocation = default;
// Find the highest level that has a block available.
var level = requestedLevel;
var maskCounts = freeMaskCounts;
while (level >= 0)
{
if (maskCounts[level] > 0) break;
level--;
}
// No blocks available.
if (level < 0) return false;
// Split a block at the level we found.
int dataIndex;
{
var freeMaskIndices = FreeMaskIndices(level);
var maskIndex = freeMaskIndices[--maskCounts[level]];
var freeMasks = FreeMasks(level);
var freeMask = freeMasks[maskIndex];
Assert.AreNotEqual(freeMask, 0);
var bitIndex = math.tzcnt(freeMask);
freeMask ^= 1ul << bitIndex;
freeMasks[maskIndex] = freeMask;
if (freeMask != 0) freeMaskIndices[maskCounts[level]++] = maskIndex;
dataIndex = maskIndex * 64 + bitIndex;
}
// Walk up the levels until we hit the requested level. For each level we want to mark the remaining parts
// of the newly split blocks as free.
while (level < requestedLevel)
{
level++;
dataIndex <<= header.branchingOrder;
var maskIndex = dataIndex >> 6;
var bitIndex = dataIndex & 63;
var freeMasks = FreeMasks(level);
var freeMask = freeMasks[maskIndex];
// We might have hit a mask that already contained free blocks. If not, add the mask index to the free list.
if (freeMask == 0)
{
var freeMaskIndices = FreeMaskIndices(level);
freeMaskIndices[maskCounts[level]++] = maskIndex;
}
// Mark other bits in the block we just broke apart as free.
// In binary form, 2^b will give us a 1 followed by b 0s. So to get b ones, we subtract 1. Since we want
// the least significant bit to be 0, we subtract another 1.
// E.g. for branching order 1 we get 10b, for 2 we get 1110b, for 3 we get 11111110b.
// Finally we shift according to the data index.
Assert.IsTrue(bitIndex + Pow2(header.branchingOrder) - 1 < 64);
freeMask |= ((1ul << Pow2(header.branchingOrder)) - 2ul) << bitIndex;
freeMasks[maskIndex] = freeMask;
}
allocation.level = level;
allocation.index = dataIndex;
return true;
}
public void Free(BuddyAllocation allocation)
{
var level = allocation.level;
var dataIndex = allocation.index;
while (level >= 0)
{
var maskIndex = dataIndex >> 6;
var bitIndex = dataIndex & 63;
var freeMasks = FreeMasks(level);
var freeMask = freeMasks[maskIndex];
var wasZero = freeMask == 0;
freeMask |= 1ul << bitIndex;
var indices = FreeMaskIndices(level);
var counts = freeMaskCounts;
var superBlockMask = ((1ul << Pow2(header.branchingOrder)) - 1) << ((bitIndex >> header.branchingOrder) * Pow2(header.branchingOrder));
// Check if the whole super-block (i.e. making up one block of the next level) is free.
// If it is, we let the loop continue upwards.
if (level == 0 || (~freeMask & superBlockMask) != 0)
{
freeMasks[maskIndex] = freeMask;
if (wasZero)
{
indices[counts[level]++] = maskIndex;
}
break;
}
freeMask &= ~superBlockMask;
freeMasks[maskIndex] = freeMask;
if (!wasZero && freeMask == 0)
{
for (var i = 0; i < indices.Length; i++)
{
if (indices[i] == maskIndex)
{
indices[i] = indices[--counts[level]];
break;
}
}
}
level--;
dataIndex >>= header.branchingOrder;
}
}
public void Dispose()
{
UnsafeUtility.Free(m_Data, m_Allocator);
#if ENABLE_UNITY_COLLECTIONS_CHECKS
AtomicSafetyHandle.Release(m_SafetyHandle);
#endif
m_Data = default;
m_Allocator = default;
}
NativeArray<T> GetNativeArray<T>(int offset, int length) where T : struct
{
var array = NativeArrayUnsafeUtility.ConvertExistingDataToNativeArray<T>(PtrAdd(m_Data, offset), length, m_Allocator);
#if ENABLE_UNITY_COLLECTIONS_CHECKS
NativeArrayUnsafeUtility.SetAtomicSafetyHandle(ref array, m_SafetyHandle);
#endif
return array;
}
// sum x^i for i=0..(n-1) = (x ^ n - 1) / (x - 1) where n is order and n is 2^branchingOrder
static int LevelOffset(int level, int branchingOrder) => (Pow2(branchingOrder) * (Pow2(branchingOrder * (level - 1) + branchingOrder) - 1)) / (Pow2(branchingOrder) - 1);
static int LevelLength(int level, int branchingOrder) => Pow2N(branchingOrder, level + 1);
// These are for when orders of length <= 64 only take up 1 item, e.g. ulong bitmasks.
static int LevelOffset64(int level, int branchingOrder) => math.min(level, 6/branchingOrder) + LevelOffset(math.max(0, level - 6/branchingOrder), branchingOrder);
static int LevelLength64(int level, int branchingOrder) => Pow2N(branchingOrder, math.max(0, level - 6/branchingOrder + 1));
static (int, int) AllocateRange<T>(int length, ref int dataSize) where T : struct
{
dataSize = AlignForward(dataSize, UnsafeUtility.AlignOf<T>());
var range = (dataSize, length);
dataSize += length * UnsafeUtility.SizeOf<T>();
return range;
}
static int AlignForward(int offset, int alignment)
{
var modulo = offset % alignment;
if (modulo != 0) offset += (alignment - modulo);
return offset;
}
static void* PtrAdd(void* ptr, int bytes) => (void*) ((IntPtr) ptr + bytes);
static int Pow2(int n) => 1 << n;
// (2^x)^n = 2^(x*n)
static int Pow2N(int x, int n) => 1 << (x * n);
}
}