Correctly handle reference types

- Add ReferenceHandle struct,
  a wrapper around GCHandle
- Use component hooks to make sure
  handles are freed when not used
- Move SpanToRef into Iterator class

src/gaemstone.ECS/Component.cs

@@ -1,17 +1,18 @@
 using System;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
 using static flecs_hub.flecs;
 
 namespace gaemstone.ECS;
 
-public static class ComponentExtensions
+public static unsafe class ComponentExtensions
 {
-	public unsafe static EntityRef InitComponent(this EntityRef entity, Type type)
+	public static EntityRef InitComponent(this EntityRef entity, Type type)
 		=> (EntityRef)typeof(ComponentExtensions)
 			.GetMethod(nameof(InitComponent), new[] { typeof(EntityRef) })!
 			.MakeGenericMethod(type).Invoke(null, new[]{ entity })!;
 
-	public unsafe static EntityRef InitComponent<T>(this EntityRef entity)
+	public static EntityRef InitComponent<T>(this EntityRef entity)
 	{
 		if (typeof(T).IsPrimitive) throw new ArgumentException(
 			"Must not be primitive");
@@ -19,12 +20,91 @@ public static class ComponentExtensions
 			"Struct component must satisfy the unmanaged constraint. " +
 			"Consider making it a class if you need to store references.");
 
-		var size = Unsafe.SizeOf<T>();
+		var size = typeof(T).IsValueType ? Unsafe.SizeOf<T>() : sizeof(ReferenceHandle);
-		var typeInfo = new ecs_type_info_t
+		var typeInfo = new ecs_type_info_t { size = size, alignment = size };
-			{ size = size, alignment = size };
+		var componentDesc = new ecs_component_desc_t { entity = entity, type = typeInfo };
-		var componentDesc = new ecs_component_desc_t
-			{ entity = entity, type = typeInfo };
 		ecs_component_init(entity.World, &componentDesc);
+
+		if (!typeof(T).IsValueType) {
+			// Set up component hooks for proper freeing of GCHandles.
+			// Without them, managed classes would never be garbage collected.
+			var typeHooks = new ecs_type_hooks_t {
+				ctor = new() { Data = new() { Pointer = &ReferenceHandle.Construct } },
+				dtor = new() { Data = new() { Pointer = &ReferenceHandle.Destruct } },
+				move = new() { Data = new() { Pointer = &ReferenceHandle.Move } },
+				copy = new() { Data = new() { Pointer = &ReferenceHandle.Copy } },
+			};
+			ecs_set_hooks_id(entity.World, entity, &typeHooks);
+		}
+
 		return entity.CreateLookup(typeof(T));
 	}
 }
+
+public unsafe readonly struct ReferenceHandle
+	: IDisposable
+{
+	public static int NumActiveHandles { get; private set; }
+
+
+	private readonly nint _value;
+
+	public object? Target =>
+		(_value != default)
+			? ((GCHandle)_value).Target
+			: null;
+
+	private ReferenceHandle(nint value)
+		=> _value = value;
+
+	public static ReferenceHandle Alloc(object? target)
+	{
+		if (target == null) return default;
+		NumActiveHandles++;
+		return new((nint)GCHandle.Alloc(target));
+	}
+
+	public ReferenceHandle Clone()
+		=> Alloc(Target);
+
+	public void Dispose()
+	{
+		if (_value == default) return;
+		NumActiveHandles--;
+		((GCHandle)_value).Free();
+	}
+
+
+	[UnmanagedCallersOnly]
+	internal static void Construct(void* ptr, int count, ecs_type_info_t* _)
+		=> new Span<ReferenceHandle>(ptr, count).Clear();
+
+	[UnmanagedCallersOnly]
+	internal static void Destruct(void* ptr, int count, ecs_type_info_t* _)
+	{
+		var span = new Span<ReferenceHandle>(ptr, count);
+		foreach (var handle in span) handle.Dispose();
+		span.Clear();
+	}
+
+	[UnmanagedCallersOnly]
+	internal static void Move(void* dstPtr, void* srcPtr, int count, ecs_type_info_t* _)
+	{
+		var dst = new Span<ReferenceHandle>(dstPtr, count);
+		var src = new Span<ReferenceHandle>(srcPtr, count);
+		foreach (var handle in dst) handle.Dispose();
+		src.CopyTo(dst);
+		src.Clear();
+	}
+
+	[UnmanagedCallersOnly]
+	internal static void Copy(void* dstPtr, void* srcPtr, int count, ecs_type_info_t* _)
+	{
+		var dst = new Span<ReferenceHandle>(dstPtr, count);
+		var src = new Span<ReferenceHandle>(srcPtr, count);
+		for (var i = 0; i < count; i++) {
+			dst[i].Dispose();
+			dst[i] = src[i].Clone();
+		}
+	}
+}

src/gaemstone.ECS/EntityRef.cs

@@ -2,7 +2,6 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 using gaemstone.Utility;
 using static flecs_hub.flecs;
 
@@ -88,22 +87,23 @@ public unsafe class EntityRef
 		where T : class
 	{
 		var ptr = ecs_get_id(World, this, id);
-		return (ptr != null) ? (T)((GCHandle)Unsafe.Read<nint>(ptr)).Target! : null;
+		return (T?)Unsafe.Read<ReferenceHandle>(ptr).Target;
 	}
 
 	public override T GetOrThrow<T>(Id id)
 	{
 		var ptr = ecs_get_id(World, this, id);
 		if (ptr == null) throw new Exception($"Component {typeof(T)} not found on {this}");
-		return typeof(T).IsValueType ? Unsafe.Read<T>(ptr)
+		if (typeof(T).IsValueType) return Unsafe.Read<T>(ptr);
-			: (T)((GCHandle)Unsafe.Read<nint>(ptr)).Target!;
+		else return (T)Unsafe.Read<ReferenceHandle>(ptr).Target!;
 	}
 
 	public override ref T GetRefOrNull<T>(Id id)
 	{
 		var @ref = ecs_ref_init_id(World, this, id);
 		var ptr  = ecs_ref_get_id(World, &@ref, id);
-		return ref (ptr != null) ? ref Unsafe.AsRef<T>(ptr) : ref Unsafe.NullRef<T>();
+		return ref (ptr != null) ? ref Unsafe.AsRef<T>(ptr)
+		                         : ref Unsafe.NullRef<T>();
 	}
 
 	public override ref T GetRefOrThrow<T>(Id id)
@@ -135,11 +135,12 @@ public unsafe class EntityRef
 
 	public override EntityRef Set<T>(Id id, T obj) where T : class
 	{
-		var handle = (nint)GCHandle.Alloc(obj);
+		if (obj == null) throw new ArgumentNullException(nameof(obj));
-		// FIXME: Previous handle needs to be freed.
+		var size = (ulong)sizeof(ReferenceHandle);
-		if (ecs_set_id(World, this, id, (ulong)sizeof(nint), &handle).Data == 0)
+		// Dispose this handle afterwards, since Flecs clones it.
+		using var handle = ReferenceHandle.Alloc(obj);
+		if (ecs_set_id(World, this, id, size, &handle).Data == 0)
 			throw new InvalidOperationException();
-		// FIXME: Handle needs to be freed when component is removed!
 		return this;
 	}
 

src/gaemstone.ECS/Iterator.cs

@@ -92,7 +92,7 @@ public unsafe class Iterator
 		where T : unmanaged => FieldIsSet(index) ? Field<T>(index) : default;
 
 	public SpanToRef<T> FieldRef<T>(int index)
-		where T : class => new(Field<nint>(index));
+		where T : class => new(Field<ReferenceHandle>(index));
 
 	public bool FieldIsSet(int index)
 	{
@@ -128,6 +128,15 @@ public unsafe class Iterator
 		public Variable(int index, string name)
 			{ Index = index; Name = name; }
 	}
+
+	public readonly ref struct SpanToRef<T>
+		where T : class
+	{
+		private readonly Span<ReferenceHandle> _span;
+		public int Length => _span.Length;
+		public T? this[int index] => (T?)_span[index].Target;
+		internal SpanToRef(Span<ReferenceHandle> span) => _span = span;
+	}
 }
 
 public enum IteratorType

src/gaemstone.Utility/SpanToRef.cs

@@ -1,19 +0,0 @@
-using System;
-using System.Runtime.InteropServices;
-
-namespace gaemstone.Utility;
-
-public readonly ref struct SpanToRef<T>
-	where T : class
-{
-	private readonly Span<nint> _span;
-
-	public int Length => _span.Length;
-	public T? this[int index] => (_span[index] != 0)
-		? (T)((GCHandle)_span[index]).Target! : null;
-
-	internal SpanToRef(Span<nint> span) => _span = span;
-
-	public void Clear() => _span.Clear();
-	public void CopyTo(SpanToRef<T> dest) => _span.CopyTo(dest._span);
-}