Last active
July 27, 2023 14:12
-
-
Save samloeschen/1d0a1cd1d555fa0d4da894401a002eeb to your computer and use it in GitHub Desktop.
C# unmanaged arena allocator and collections for it
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
using System; | |
using System.Runtime.InteropServices; | |
using System.Runtime.CompilerServices; | |
using Range = System.Range; | |
public unsafe struct Arena { | |
UnmanagedArray<Optional<Block>> _blocks; | |
int _currentBlockIndex; | |
const int MIN_POW2 = 6; | |
const int MAX_POW2 = 30; | |
public Arena() { | |
int numBlocks = MAX_POW2 - MIN_POW2; | |
_blocks = new UnmanagedArray<Optional<Block>>(numBlocks); | |
_currentBlockIndex = MIN_POW2; | |
} | |
public static void Reset(ref Arena arena) { | |
for (int i = 0; i < arena._blocks.Length; i++) { | |
if (arena._blocks[i].TryGetValue(out var block)) { | |
block.Reset(); | |
arena._blocks[i] = block; | |
} | |
} | |
} | |
public static void Delete(ref Arena arena) { | |
for (int i = 0; i < arena._blocks.Length; i++) { | |
if (arena._blocks[i].TryGetValue(out var block)) { | |
Marshal.FreeHGlobal((IntPtr)block.backing); | |
arena._blocks[i] = default; | |
} | |
arena._blocks.Delete(); | |
} | |
} | |
Block AllocBlock(int blockIndex) { | |
int blockSize = 1 << blockIndex; | |
return new Block { | |
backing = (byte*)Marshal.AllocHGlobal(blockSize), | |
offset = 0, | |
size = blockSize, | |
}; | |
} | |
int GetMinBlockIndex(int byteLen) { | |
int minBlockSize = NextPowerOf2(byteLen); | |
int exp2 = (int)Math.Floor(Math.Log2(minBlockSize)); | |
return Math.Max(0, exp2 - MIN_POW2); | |
} | |
T* AllocBytes<T>(int length) where T: unmanaged { | |
int elemSize = Unsafe.SizeOf<T>(); | |
int byteLen = length * elemSize; | |
int minBlockIndex = GetMinBlockIndex(byteLen); | |
if (minBlockIndex > _blocks.Length) { | |
throw new OutOfMemoryException($"Tried to allocate {byteLen} bytes -- this is larger than the max block allocation size of the Arena (1073741824B, or about 1GB)"); | |
} | |
if (minBlockIndex > _currentBlockIndex) { | |
_currentBlockIndex = minBlockIndex; | |
} | |
// if this block is not allocated yet, create it | |
if (!_blocks[_currentBlockIndex].TryGetValue(out var block)) { | |
_blocks[_currentBlockIndex] = block = AllocBlock(_currentBlockIndex); | |
} | |
// if we're asking for more memory than the block has left, go to the next block | |
if (block.offset + byteLen > block.size) { | |
if (_currentBlockIndex == _blocks.Length - 1) { | |
throw new OutOfMemoryException($"Tried to allocate {byteLen} bytes, but the Arena doesn't have enough block memory remaining!"); | |
} | |
_currentBlockIndex += 1; | |
if (!_blocks[_currentBlockIndex].TryGetValue(out block)) { | |
_blocks[_currentBlockIndex] = block = AllocBlock(_currentBlockIndex); | |
} | |
} | |
var ptr = (T*)(block.backing + block.offset); | |
block.offset += byteLen; | |
_blocks[_currentBlockIndex] = block; | |
return ptr; | |
} | |
public static Array<T> AllocArray<T>(int length, ref Arena arena) where T: unmanaged { | |
var ptr = arena.AllocBytes<T>(length); | |
return new Array<T>(ptr, length); | |
} | |
public static List<T> AllocList<T>(int initialCapacity, ref Arena arena) where T: unmanaged { | |
return new List<T>(ref arena, initialCapacity); | |
} | |
struct Block { | |
public byte* backing; | |
public int offset; | |
public int size; | |
public void Reset() => offset = 0; | |
} | |
public struct Array<T> where T: unmanaged { | |
T* _ptr; | |
public int Length { get; } | |
public Array(T* ptr, int length) { | |
_ptr = ptr; | |
Length = length; | |
} | |
public Span<T> Slice(int start, int length) { | |
return new Span<T>((void*)(_ptr + start), length); | |
} | |
public T this[int i] { | |
get { | |
if (i < 0 || i >= Length) { | |
throw new IndexOutOfRangeException($"Index {i} is out of the range 0..<{Length}"); | |
} | |
return _ptr[i]; | |
} | |
set { | |
if (i < 0 || i >= Length) { | |
throw new IndexOutOfRangeException($"Index {i} is out of the range 0..<{Length}"); | |
} | |
_ptr[i] = value; | |
} | |
} | |
public int IndexOf(T value) { | |
for (int i = 0; i < Length; i++) { | |
if (this[i].Equals(value)) { | |
return i; | |
} | |
} | |
return -1; | |
} | |
public bool Contains(T value) => IndexOf(value) >= 0; | |
public Chunk<T> Chunk(int start, int length) => new Chunk<T>(start, length, this); | |
public Chunk<T> Chunk(int length) => Chunk(0, length); | |
public static implicit operator Chunk<T>(Array<T> arr) => arr.Chunk(0, arr.Length); | |
} | |
public struct List<T> where T: unmanaged { | |
public int Count { get; set; } | |
public int Capacity => _array.Length; | |
Arena* _arenaRef; | |
public Array<T> _array; | |
public List(ref Arena arena, int initialCapacity) { | |
_array = Arena.AllocArray<T>(initialCapacity, ref arena); | |
_arenaRef = (Arena*)Unsafe.AsPointer<Arena>(ref arena); | |
Count = 0; | |
} | |
public T this[int i] { | |
get => _array[i]; | |
set => _array[i] = value; | |
} | |
public void Add(T elem) { | |
if (Count == Capacity) { | |
ref var arena = ref Unsafe.AsRef<Arena>(_arenaRef); | |
var newArray = Arena.AllocArray<T>(Capacity * 2, ref arena); | |
_array[..].CopyTo(newArray[0..Count]); | |
_array = newArray; | |
} | |
_array[Count] = elem; | |
Count++; | |
} | |
public void Clear() => Count = 0; | |
public void RemoveAt(int index) { | |
if (index < 0 || index > Count) { | |
throw new IndexOutOfRangeException($"Index {index} is out of bounds of the range 0..{Count}"); | |
} | |
for (int i = index; i < Count - 1; i++) { | |
this[index] = this[index + 1]; | |
} | |
Count--; | |
} | |
public void RemoveAtSwapback(int index) { | |
if (index < 0 || index > Count) { | |
throw new IndexOutOfRangeException($"Index {index} is out of bounds of the range 0..{Count}"); | |
} | |
this[index] = this[Count - 1]; | |
Count--; | |
} | |
public void Remove(T elem) { | |
int index = IndexOf(elem); | |
if (index >= 0) { | |
RemoveAt(index); | |
} | |
} | |
public int IndexOf(T elem) => _array.IndexOf(elem); | |
public bool Contains(T elem) => _array.Contains(elem); | |
public Chunk<T> Chunk(int start, int length) => _array.Chunk(0, Count); | |
public Chunk<T> Chunk(int length) => Chunk(0, length); | |
public Chunk<T> Chunk() => Chunk(Count); | |
public static implicit operator Chunk<T>(List<T> list) => list.Chunk(0, list.Count); | |
} | |
// read only! | |
// both arena lists and arena arrays can be cast to these or sliced into these | |
// basically a persistent span (can also be cast to a ReadOnlySpan<T>) | |
public struct Chunk<T> where T: unmanaged { | |
public int Start; | |
public int Length; | |
// TODO prob could just use the pointer... | |
Array<T> _array; | |
public T this[int i] { | |
get => _array[i]; | |
} | |
public Chunk(int start, int length, Array<T> array) { | |
Start = start; | |
Length = length; | |
_array = array; | |
} | |
public static implicit operator ReadOnlySpan<T>(Chunk<T> s) => s._array[s.Start..(s.Start + s.Length)]; | |
} | |
public static int NextPowerOf2(int v) { | |
v--; | |
v |= v >> 1; | |
v |= v >> 2; | |
v |= v >> 4; | |
v |= v >> 8; | |
v |= v >> 16; | |
v++; | |
return v; | |
} | |
} | |
unsafe struct UnmanagedArray<T> where T: unmanaged { | |
public int Length { get; private set; } | |
public T* Ptr => _ptr; | |
T* _ptr; | |
public UnmanagedArray(int length) { | |
int elemSize = Unsafe.SizeOf<T>(); | |
_ptr = (T*)Marshal.AllocHGlobal(length * elemSize); | |
Length = length; | |
} | |
public void Resize(int newLength) { | |
int elemSize = Unsafe.SizeOf<T>(); | |
var newPtr = (T*)Marshal.AllocHGlobal(newLength * elemSize); | |
Span<T> src = new Span<T>((void*)_ptr, Length); | |
Span<T> dst = new Span<T>((void*)newPtr, Math.Min(Length, newLength)); | |
src.CopyTo(dst); | |
Marshal.FreeHGlobal((IntPtr)_ptr); | |
Length = newLength; | |
_ptr = newPtr; | |
} | |
public void Delete() { | |
Marshal.FreeHGlobal((IntPtr)_ptr); | |
_ptr = null; | |
Length = 0; | |
} | |
public T this[int index] { | |
get => _ptr[index]; | |
set => _ptr[index] = value; | |
} | |
public Span<T> this[Range r] => new Span<T>((void*)(_ptr + r.Start.Value), r.End.Value - r.Start.Value); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment