gaemstone/src/gaemstone/Utility/IL/QueryActionGenerator.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Reflection.Emit;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using gaemstone.ECS;
using static flecs_hub.flecs;

namespace gaemstone.Utility.IL;

public unsafe class QueryActionGenerator
{
	private static readonly PropertyInfo _iteratorUniverseProp   = typeof(Iterator).GetProperty(nameof(Iterator.Universe))!;
	private static readonly PropertyInfo _iteratorDeltaTimeProp  = typeof(Iterator).GetProperty(nameof(Iterator.DeltaTime))!;
	private static readonly PropertyInfo _iteratorCountProp      = typeof(Iterator).GetProperty(nameof(Iterator.Count))!;
	private static readonly MethodInfo _iteratorFieldMethod      = typeof(Iterator).GetMethod(nameof(Iterator.Field))!;
	private static readonly MethodInfo _iteratorFieldIsSetMethod = typeof(Iterator).GetMethod(nameof(Iterator.FieldIsSet))!;
	private static readonly MethodInfo _iteratorEntityMethod     = typeof(Iterator).GetMethod(nameof(Iterator.Entity))!;

	private static readonly MethodInfo _handleFromIntPtrMethod = typeof(GCHandle).GetMethod(nameof(GCHandle.FromIntPtr))!;
	private static readonly PropertyInfo _handleTargetProp     = typeof(GCHandle).GetProperty(nameof(GCHandle.Target))!;

	private static readonly ConditionalWeakTable<MethodInfo, QueryActionGenerator> _cache = new();
	private static readonly Dictionary<Type, Action<ILGeneratorWrapper, IArgument<Iterator>, ILocal<int>>> _uniqueParameters = new() {
		[typeof(Iterator)] = (IL, iter, i) => { IL.Load(iter); },
		[typeof(Universe)] = (IL, iter, i) => { IL.Load(iter, _iteratorUniverseProp); },
		[typeof(TimeSpan)] = (IL, iter, i) => { IL.Load(iter, _iteratorDeltaTimeProp); },
		[typeof(Entity)]   = (IL, iter, i) => { IL.Load(iter); IL.Load(i); IL.Call(_iteratorEntityMethod); },
	};

	public Universe Universe { get; }
	public MethodInfo Method { get; }
	public ParamInfo[] Parameters { get; }

	public ecs_filter_desc_t Filter { get; }
	public Action<object?, Iterator> GeneratedAction { get; }
	public string ReadableString { get; }

	public void RunWithTryCatch(object? instance, Iterator iter)
	{
		try { GeneratedAction(instance, iter); } catch {
			Console.WriteLine("Exception occured while running:");
			Console.WriteLine("    " + Method);
			Console.WriteLine();
			Console.WriteLine("Method's IL code:");
			Console.WriteLine(ReadableString);
			Console.WriteLine();
			throw;
		}
	}

	public QueryActionGenerator(Universe universe, MethodInfo method)
	{
		Universe = universe;
		Method   = method;

		Parameters = method.GetParameters().Select(ParamInfo.Build).ToArray();
		if (!Parameters.Any(c => c.IsRequired && (c.Kind != ParamKind.Unique)))
			throw new ArgumentException($"At least one parameter in {method} is required");

		var filter = default(ecs_filter_desc_t);
		var name   = "<>Query_" + string.Join("_", Parameters.Select(p => p.UnderlyingType.Name));
		var genMethod = new DynamicMethod(name, null, new[] { typeof(object), typeof(Iterator) });
		var IL = new ILGeneratorWrapper(genMethod);

		var instanceArg = IL.Argument<object?>(0);
		var iteratorArg = IL.Argument<Iterator>(1);

		var counter = 0; // Counter for fields actually part of the filter terms.
		var fieldLocals = new ILocal[Parameters.Length];
		var tempLocals  = new ILocal[Parameters.Length];

		for (var i = 0; i < Parameters.Length; i++) {
			var p = Parameters[i];
			if (p.Kind == ParamKind.Unique) continue;

			// Update the flecs filter to look for this type.
			// Term index is 0-based and field index (used below) is 1-based, so increasing counter here works out.
			ref var term = ref filter.terms[counter++];
			term.id = Universe.Lookup(p.UnderlyingType);
			term.inout = p.Kind switch {
				ParamKind.In => ecs_inout_kind_t.EcsIn,
				ParamKind.Out => ecs_inout_kind_t.EcsOut,
				ParamKind.Has or ParamKind.Not => ecs_inout_kind_t.EcsInOutNone,
				_ => ecs_inout_kind_t.EcsInOut,
			};
			term.oper = p.Kind switch {
				ParamKind.Not => ecs_oper_kind_t.EcsNot,
				_ when !p.IsRequired => ecs_oper_kind_t.EcsOptional,
				_ => ecs_oper_kind_t.EcsAnd,
			};
			if (p.Source != null) term.src = new() { id = Universe.Lookup(p.Source) };

			// Create a Span<T> local and initialize it to iterator.Field<T>(i).
			var spanType   = typeof(Span<>).MakeGenericType(p.FieldType);
			fieldLocals[i] = IL.Local(spanType, $"field_{counter}");
			if (p.Kind is ParamKind.Has or ParamKind.Not) {
				// If a "has" or "not" parameter is a struct, we require a temporary local that
				// we can later load onto the stack when loading the arguments for the action.
				if (p.ParameterType.IsValueType) {
					IL.Comment($"temp_{counter} = default({p.ParameterType});");
					tempLocals[i] = IL.Local(p.ParameterType);
					IL.LoadAddr(tempLocals[i]);
					IL.Init(tempLocals[i].LocalType);
				}
			} else if (p.IsRequired) {
				IL.Comment($"field_{counter} = iterator.Field<{p.FieldType.Name}>({counter})");
				IL.Load(iteratorArg);
				IL.LoadConst(counter);
				IL.Call(_iteratorFieldMethod.MakeGenericMethod(p.FieldType));
				IL.Store(fieldLocals[i]);
			} else {
				IL.Comment($"field_{counter} = iterator.FieldIsSet({counter}) ? iterator.Field<{p.FieldType.Name}>({counter}) : default");
				var elseLabel = IL.DefineLabel();
				var doneLabel = IL.DefineLabel();
				IL.Load(iteratorArg);
				IL.LoadConst(counter);
				IL.Call(_iteratorFieldIsSetMethod);
				IL.GotoIfFalse(elseLabel);
					IL.Load(iteratorArg);
					IL.LoadConst(counter);
					IL.Call(_iteratorFieldMethod.MakeGenericMethod(p.FieldType));
					IL.Store(fieldLocals[i]);
				IL.Goto(doneLabel);
				IL.MarkLabel(elseLabel);
					IL.LoadAddr(fieldLocals[i]);
					IL.Init(spanType);
				IL.MarkLabel(doneLabel);
			}

			if (p.Kind == ParamKind.Nullable) {
				IL.Comment($"temp_{counter} = default({p.ParameterType});");
				tempLocals[i] = IL.Local(p.ParameterType);
				IL.LoadAddr(tempLocals[i]);
				IL.Init(tempLocals[i].LocalType);
			}
		}

		// If there's any reference type parameters, we need to define a GCHandle local.
		var hasReferenceType = Parameters
			.Where(p => p.Kind != ParamKind.Unique)
			.Any(p => !p.UnderlyingType.IsValueType);
		var handleLocal = hasReferenceType ? IL.Local<GCHandle>() : null;

		using (IL.For(() => IL.Load(iteratorArg, _iteratorCountProp), out var currentLocal)) {
			if (!Method.IsStatic)
				IL.Load(instanceArg);
			for (var i = 0; i < Parameters.Length; i++) {
				var p = Parameters[i];
				if (p.Kind == ParamKind.Unique) {
					IL.Comment($"Unique parameter {p.ParameterType}");
					_uniqueParameters[p.ParameterType](IL, iteratorArg, currentLocal);
				} else if (p.Kind is ParamKind.Has or ParamKind.Not) {
					if (p.ParameterType.IsValueType)
						IL.LoadObj(tempLocals[i]!);
					else IL.LoadNull();
				} else {
					var spanType = typeof(Span<>).MakeGenericType(p.FieldType);
					var spanItemMethod   = spanType.GetProperty("Item")!.GetMethod!;
					var spanLengthMethod = spanType.GetProperty("Length")!.GetMethod!;

					IL.Comment($"Parameter {p.ParameterType}");
					if (p.IsByRef) {
						IL.LoadAddr(fieldLocals[i]!);
						IL.Load(currentLocal);
						IL.Call(spanItemMethod);
					} else if (p.IsRequired) {
						IL.LoadAddr(fieldLocals[i]!);
						IL.Load(currentLocal);
						IL.Call(spanItemMethod);
						IL.LoadObj(p.FieldType);
					} else {
						var elseLabel = IL.DefineLabel();
						var doneLabel = IL.DefineLabel();
						IL.LoadAddr(fieldLocals[i]!);
						IL.Call(spanLengthMethod);
						IL.GotoIfFalse(elseLabel);
							IL.LoadAddr(fieldLocals[i]!);
							IL.Load(currentLocal);
							IL.Call(spanItemMethod);
							IL.LoadObj(p.FieldType);
							if (p.Kind == ParamKind.Nullable)
								IL.New(p.ParameterType);
						IL.Goto(doneLabel);
						IL.MarkLabel(elseLabel);
							if (p.Kind == ParamKind.Nullable)
								IL.LoadObj(tempLocals[i]!);
							else IL.LoadNull();
						IL.MarkLabel(doneLabel);
					}

					if (!p.UnderlyingType.IsValueType) {
						IL.Comment($"Convert nint to {p.UnderlyingType}");
						IL.Call(_handleFromIntPtrMethod);
						IL.Store(handleLocal!);
						IL.LoadAddr(handleLocal!);
						IL.Call(_handleTargetProp.GetMethod!);
						IL.Cast(p.UnderlyingType);
					}
				}
			}
			IL.Call(Method);
		}

		IL.Return();

		Filter = filter;
		GeneratedAction = genMethod.CreateDelegate<Action<object?, Iterator>>();
		ReadableString  = IL.ToReadableString();
	}

	public static QueryActionGenerator GetOrBuild(Universe universe, MethodInfo method)
		=>_cache.GetValue(method, m => new QueryActionGenerator(universe, m));

	public class ParamInfo
	{
		public ParameterInfo Info { get; }
		public int Index { get; }

		public ParamKind Kind { get; }
		public Type ParameterType { get; }
		public Type UnderlyingType { get; }
		public Type FieldType { get; }

		public Type? Source { get; }

		public bool IsRequired => (Kind <  ParamKind.Nullable);
		public bool IsByRef    => (Kind >= ParamKind.In) && (Kind <= ParamKind.Ref);

		private ParamInfo(
			ParameterInfo info, int index, ParamKind kind,
			Type paramType, Type underlyingType)
		{
			Info  = info;
			Index = index;

			Kind = kind;
			ParameterType  = paramType;
			UnderlyingType = underlyingType;
			// Reference types have a backing type of nint - they're pointers.
			FieldType = underlyingType.IsValueType ? underlyingType : typeof(nint);

			// If the underlying type has EntityAttribute, it's a singleton.
			if (UnderlyingType.Has<EntityAttribute>()) Source = underlyingType;
			if (Info.Get<SourceAttribute>() is SourceAttribute attr) Source = attr.Type;
		}

		public static ParamInfo Build(ParameterInfo info, int index)
		{
			if (info.IsOptional)              throw new ArgumentException("Optional parameters are not supported\nParameter: " + info);
			if (info.ParameterType.IsArray)   throw new ArgumentException("Arrays are not supported\nParameter: " + info);
			if (info.ParameterType.IsPointer) throw new ArgumentException("Pointers are not supported\nParameter: " + info);

			if (_uniqueParameters.ContainsKey(info.ParameterType))
				return new(info, index, ParamKind.Unique, info.ParameterType, info.ParameterType);

			var isByRef    = info.ParameterType.IsByRef;
			var isNullable = info.IsNullable();

			if (info.Has<NotAttribute>()) {
				if (isByRef || isNullable) throw new ArgumentException(
					"Parameter with NotAttribute must not be ByRef or nullable\nParameter: " + info);
				return new(info, index, ParamKind.Not, info.ParameterType, info.ParameterType);
			}

			if (info.Has<HasAttribute>() || info.ParameterType.Has<TagAttribute>()) {
				if (isByRef || isNullable) throw new ArgumentException(
					"Parameter with HasAttribute / TagAttribute must not be ByRef or nullable\nParameter: " + info);
				return new(info, index, ParamKind.Has, info.ParameterType, info.ParameterType);
			}

			var kind = ParamKind.Normal;
			var underlyingType = info.ParameterType;

			if (info.IsNullable()) {
				if (info.ParameterType.IsValueType)
					underlyingType = Nullable.GetUnderlyingType(info.ParameterType)!;
				kind = ParamKind.Nullable;
			}

			if (info.ParameterType.IsByRef) {
				if (kind == ParamKind.Nullable) throw new ArgumentException(
					"ByRef and Nullable are not supported together\nParameter: " + info);
				underlyingType = info.ParameterType.GetElementType()!;
				if (!underlyingType.IsValueType) throw new ArgumentException(
					"Reference types can't also be ByRef\nParameter: " + info);
				kind = info.IsIn  ? ParamKind.In
				     : info.IsOut ? ParamKind.Out
				                  : ParamKind.Ref;
			}

			if (underlyingType.IsPrimitive) throw new ArgumentException(
				"Primitives are not supported\nParameter: " + info);

			return new(info, index, kind, info.ParameterType, underlyingType);
		}
	}

	public enum ParamKind
	{
		/// <summary> Parameter is not part of terms, handled uniquely, such as Universe and Entity. </summary>
		Unique,
		/// <summary> Passed by value. </summary>
		Normal,
		/// <summary> Struct passed with the "in" modifier. </summary>
		In,
		/// <summary> Struct passed with the "out" modifier. </summary>
		Out,
		/// <summary> Struct passed with the "ref" modifier. </summary>
		Ref,
		/// <summary>
		/// Only checks for presence.
		/// Manually applied with <see cref="HasAttribute"/>.
		/// Automatically applied for types with <see cref="TagAttribute"/>.
		/// </summary>
		Has,
		/// <summary> Struct passed as <c>Nullable&lt;T&gt;</c>. </summary>
		Nullable,
		/// <summary>
		/// Only checks for absence.
		/// Applied with <see cref="NotAttribute"/>.
		/// </summary>
		Not,
	}
}