Skip to content

Instantly share code, notes, and snippets.

@samloeschen
Last active July 27, 2023 14:12
Show Gist options
  • Save samloeschen/1d0a1cd1d555fa0d4da894401a002eeb to your computer and use it in GitHub Desktop.
Save samloeschen/1d0a1cd1d555fa0d4da894401a002eeb to your computer and use it in GitHub Desktop.
C# unmanaged arena allocator and collections for it
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