gaemstone/src/gaemstone/Utility/IL/IterActionGenerator.cs

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Reflection;
using System.Reflection.Emit;
using System.Runtime.CompilerServices;
using gaemstone.ECS;
using GCHandle = System.Runtime.InteropServices.GCHandle;

namespace gaemstone.Utility.IL;

// TODO: Support tuple syntax to match relationship pairs.
public unsafe class IterActionGenerator
{
	private static readonly ConstructorInfo _entityRefCtor = typeof(EntityRef).GetConstructors().Single();

	private static readonly PropertyInfo _iteratorWorldProp      = typeof(Iterator).GetProperty(nameof(Iterator.World))!;
	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 _iteratorEntityMethod     = typeof(Iterator).GetMethod(nameof(Iterator.Entity))!;
	private static readonly MethodInfo _iteratorFieldMethod      = typeof(Iterator).GetMethod(nameof(Iterator.Field))!;
	private static readonly MethodInfo _iteratorFieldOrEmptyMethod = typeof(Iterator).GetMethod(nameof(Iterator.FieldOrEmpty))!;

	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, IterActionGenerator> _cache = new();
	private static readonly Dictionary<Type, Action<ILGeneratorWrapper, IArgument<Iterator>>> _globalUniqueParameters = new() {
		[typeof(World)]    = (IL, iter) => { IL.Load(iter, _iteratorWorldProp); },
		[typeof(Universe)] = (IL, iter) => { IL.Load(iter, _iteratorWorldProp); IL.Cast(typeof(Universe)); },
		[typeof(TimeSpan)] = (IL, iter) => { IL.Load(iter, _iteratorDeltaTimeProp); },
	};
	private static readonly Dictionary<Type, Action<ILGeneratorWrapper, IArgument<Iterator>, ILocal<int>>> _uniqueParameters = new() {
		[typeof(Iterator)]  = (IL, iter, i) => { IL.Load(iter); },
		[typeof(EntityRef)] = (IL, iter, i) => { IL.Load(iter); IL.Load(i); IL.Call(_iteratorEntityMethod); },
	};

	public World World { get; }
	public MethodInfo Method { get; }
	public IReadOnlyList<ParamInfo> Parameters { get; }

	public IReadOnlyList<Term> Terms { get; }
	public Action<object?, Iterator> GeneratedAction { get; }
	public string ReadableString { get; }

	public void RunWithTryCatch(object? instance, Iterator iter)
	{
		try { GeneratedAction(instance, iter); }
		catch { Console.Error.WriteLine(ReadableString); throw; }
	}

	public IterActionGenerator(World world, MethodInfo method)
	{
		World      = world;
		Method     = method;
		Parameters = method.GetParameters().Select(ParamInfo.Build).ToImmutableArray();

		var name = "<>Query_" + string.Join("_", method.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 fieldIndex = 0;
		var paramData  = new List<(ParamInfo Info, Term? Term, ILocal? FieldLocal, ILocal? TempLocal)>();
		foreach (var p in Parameters) {
			// If the parameter is unique, we don't create a term for it.
			if (p.Kind <= ParamKind.Unique)
				{ paramData.Add((p, null, null, null)); continue; }

			// Create a term to add to the query.
			var term = new Term(world.LookupByTypeOrThrow(p.UnderlyingType)) {
				Source = (p.Source != null) ? (TermId)World.LookupByTypeOrThrow(p.Source) : null,
				InOut = p.Kind switch {
					ParamKind.In  => TermInOutKind.In,
					ParamKind.Out => TermInOutKind.Out,
					ParamKind.Not or ParamKind.Not => TermInOutKind.None,
					_ => default,
				},
				Oper = p.Kind switch {
					ParamKind.Not => TermOperKind.Not,
					ParamKind.Or  => TermOperKind.Or,
					_ when !p.IsRequired => TermOperKind.Optional,
					_ => default,
				},
			};

			// If this and the previous parameter are marked with [Or], do not advance the field index.
			if ((fieldIndex == 0) || (p.Kind != ParamKind.Or) || (paramData[^1].Info.Kind != ParamKind.Or))
				fieldIndex++;

			var spanType   = typeof(Span<>).MakeGenericType(p.FieldType);
			var fieldLocal = (ILocal?)null;
			var tempLocal  = (ILocal?)null;

			switch (p.Kind) {
				// FIXME: Currently would not work with [Or]'d components.
				case ParamKind.Has or ParamKind.Not or ParamKind.Or:
					if (!p.ParameterType.IsValueType) break;
					// If parameter is a struct, we require a temporary local that we can
					// later load onto the stack when loading the arguments for the action.
					IL.Comment($"{p.Info.Name}Temp = default({p.ParameterType});");
					tempLocal = IL.Local(p.ParameterType);
					IL.LoadAddr(tempLocal);
					IL.Init(tempLocal.LocalType);
					break;

				case ParamKind.Nullable or ParamKind.Or:
					IL.Comment($"{p.Info.Name}Field = iterator.FieldOrEmpty<{p.FieldType.Name}>({fieldIndex})");
					fieldLocal = IL.Local(spanType, $"{p.Info.Name}Field");
					IL.Load(iteratorArg);
					IL.LoadConst(fieldIndex);
					IL.Call(_iteratorFieldOrEmptyMethod.MakeGenericMethod(p.FieldType));
					IL.Store(fieldLocal);

					if (p.UnderlyingType.IsValueType) {
						IL.Comment($"{p.Info.Name}Temp = default({p.ParameterType});");
						tempLocal = IL.Local(p.ParameterType);
						IL.LoadAddr(tempLocal);
						IL.Init(tempLocal.LocalType);
					}
					break;

				default:
					IL.Comment($"{p.Info.Name}Field = iterator.Field<{p.FieldType.Name}>({fieldIndex})");
					fieldLocal = IL.Local(spanType, $"{p.Info.Name}Field");
					IL.Load(iteratorArg);
					IL.LoadConst(fieldIndex);
					IL.Call(_iteratorFieldMethod.MakeGenericMethod(p.FieldType));
					IL.Store(fieldLocal);
					break;
			}

			paramData.Add((p, term, fieldLocal, tempLocal));
		}

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

		var indexLocal = IL.Local<int>("iter_index");
		var countLocal = IL.Local<int>("iter_count");

		IL.Set(indexLocal, 0);
		IL.Load(iteratorArg, _iteratorCountProp);
		IL.Store(countLocal);

		// If all parameters are fixed, iterator count will be 0, but since
		// the query matched, we want to run the callback at least once.
		IL.Comment("if (iter_count == 0) iter_count = 1;");
		var dontIncrementLabel = IL.DefineLabel();
		IL.Load(countLocal);
		IL.GotoIfTrue(dontIncrementLabel);
		IL.LoadConst(1);
		IL.Store(countLocal);
		IL.MarkLabel(dontIncrementLabel);

		IL.While("IteratorLoop", (@continue) => {
			IL.GotoIf(@continue, indexLocal, Comparison.LessThan, countLocal);
		}, (_, _) => {
			if (!Method.IsStatic)
				IL.Load(instanceArg);

			foreach (var (info, term, fieldLocal, tempLocal) in paramData) {
				var isValueType = info.UnderlyingType.IsValueType;
				var paramName   = info.ParameterType.GetFriendlyName();
				switch (info.Kind) {

					case ParamKind.GlobalUnique:
						IL.Comment($"Global unique parameter {paramName}");
						_globalUniqueParameters[info.ParameterType](IL, iteratorArg);
						break;

					case ParamKind.Unique:
						IL.Comment($"Unique parameter {paramName}");
						_uniqueParameters[info.ParameterType](IL, iteratorArg, indexLocal!);
						break;

					// FIXME: Currently would not work with [Or]'d components.
					case ParamKind.Has or ParamKind.Not or ParamKind.Or:
						IL.Comment($"Has parameter {paramName}");
						if (isValueType) IL.LoadObj(tempLocal!);
						else             IL.LoadNull();
						break;

					default:
						var spanType = typeof(Span<>).MakeGenericType(info.FieldType);
						var spanItemMethod   = spanType.GetProperty("Item")!.GetMethod!;
						var spanLengthMethod = spanType.GetProperty("Length")!.GetMethod!;

						IL.Comment($"Parameter {paramName}");
						if (info.IsByRef) {
							IL.LoadAddr(fieldLocal!);
							if (info.IsFixed) IL.LoadConst(0);
							else IL.Load(indexLocal!);
							IL.Call(spanItemMethod);
						} else if (info.IsRequired) {
							IL.LoadAddr(fieldLocal!);
							if (info.IsFixed) IL.LoadConst(0);
							else IL.Load(indexLocal!);
							IL.Call(spanItemMethod);
							IL.LoadObj(info.FieldType);

							if (!isValueType) {
								IL.Comment($"Convert nint to {paramName}");
								IL.Call(_handleFromIntPtrMethod);
								IL.Store(handleLocal!);
								IL.LoadAddr(handleLocal!);
								IL.Call(_handleTargetProp.GetMethod!);
								IL.Cast(info.UnderlyingType);
							}
						} else {
							var elseLabel = IL.DefineLabel();
							var doneLabel = IL.DefineLabel();
							IL.LoadAddr(fieldLocal!);
							IL.Call(spanLengthMethod);
							IL.GotoIfFalse(elseLabel);
								IL.LoadAddr(fieldLocal!);
								if (info.IsFixed) IL.LoadConst(0);
								else IL.Load(indexLocal!);
								IL.Call(spanItemMethod);
								IL.LoadObj(info.FieldType);
								if (!isValueType) {
									IL.Comment($"Convert nint to {paramName}");
									IL.Call(_handleFromIntPtrMethod);
									IL.Store(handleLocal!);
									IL.LoadAddr(handleLocal!);
									IL.Call(_handleTargetProp.GetMethod!);
									IL.Cast(info.UnderlyingType);
								} else IL.New(info.ParameterType);
							IL.Goto(doneLabel);
							IL.MarkLabel(elseLabel);
								if (!isValueType) IL.LoadNull();
								else IL.LoadObj(tempLocal!);
							IL.MarkLabel(doneLabel);
						}
						break;
				}
			}
			IL.Call(Method);

			IL.Increment(indexLocal);
		});

		IL.Return();

		Terms = paramData.Where(p => p.Term != null).Select(p => p.Term!).ToImmutableList();
		GeneratedAction = genMethod.CreateDelegate<Action<object?, Iterator>>();
		ReadableString  = IL.ToReadableString();
	}

	public static IterActionGenerator GetOrBuild(World world, MethodInfo method)
		=>_cache.GetValue(method, m => new IterActionGenerator(world, m));

	public class ParamInfo
	{
		public ParameterInfo Info { 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 is ParamKind.In or ParamKind.Ref);
		public bool IsFixed    => (Kind == ParamKind.GlobalUnique) || (Source != null);

		private ParamInfo(ParameterInfo info, ParamKind kind,
		                  Type paramType, Type underlyingType)
		{
			Info = info;
			Kind = kind;
			ParameterType  = paramType;
			UnderlyingType = underlyingType;
			// Reference types have a backing type of nint - they're pointers.
			FieldType = underlyingType.IsValueType ? underlyingType : typeof(nint);

			if (UnderlyingType.Has<SingletonAttribute>()) Source = UnderlyingType;
			if (Info.Get<SourceAttribute>()?.Type is Type type) Source = type;
		}

		public static ParamInfo Build(ParameterInfo info)
		{
			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 (info.ParameterType.IsPrimitive) throw new ArgumentException($"Primitives are not supported\nParameter: {info}");

			// Find out initial parameter kind from provided attribute.
			var fromAttributes = new List<ParamKind>();
			if (info.Has< InAttribute>()) fromAttributes.Add(ParamKind.In);
			if (info.Has<OutAttribute>()) fromAttributes.Add(ParamKind.Out);
			if (info.Has<HasAttribute>()) fromAttributes.Add(ParamKind.Has);
			if (info.Has< OrAttribute>()) fromAttributes.Add(ParamKind.Or);
			if (info.Has<NotAttribute>()) fromAttributes.Add(ParamKind.Not);
			// Throw an error if multiple incompatible attributes were found.
			if (fromAttributes.Count > 1) throw new ArgumentException(
				"Parameter must not be marked with multiple attributes: "
					+ string.Join(", ", fromAttributes.Select(a => $"[{a}]"))
					+ $"\nParameter: {info}");
			var kind = fromAttributes.FirstOrNull() ?? ParamKind.Normal;

			// Handle unique parameters such as Universe, EntityRef, ...
			var isGlobalUnique = _globalUniqueParameters.ContainsKey(info.ParameterType);
			var isUnique       = _uniqueParameters.ContainsKey(info.ParameterType);
			if (isGlobalUnique || isUnique) {
				if (kind != ParamKind.Normal) throw new ArgumentException(
					$"Unique parameter {info.ParameterType.Name} does not support [{kind}]\nParameter: {info}");
				kind = isGlobalUnique ? ParamKind.GlobalUnique : ParamKind.Unique;
				return new(info, kind, info.ParameterType, info.ParameterType);
			}

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

			if (info.ParameterType.Has<TagAttribute>() && (kind is not (ParamKind.Has or ParamKind.Not or ParamKind.Or))) {
				if (kind is not ParamKind.Normal) throw new ArgumentException($"Parameter does not support [{kind}]\nParameter: {info}");
				kind = ParamKind.Has;
			}

			if (kind is ParamKind.Not or ParamKind.Has) {
				if (isNullable) throw new ArgumentException($"Parameter does not support Nullable\nParameter: {info}");
				if (isByRef) throw new ArgumentException($"Parameter does not support ByRef\nParameter: {info}");
				return new(info, kind, info.ParameterType, info.ParameterType);
			}

			var underlyingType = info.ParameterType;

			if (isNullable) {
				if (isByRef) throw new ArgumentException($"Parameter does not support ByRef\nParameter: {info}");
				if (info.ParameterType.IsValueType)
					underlyingType = Nullable.GetUnderlyingType(info.ParameterType)!;
				kind = ParamKind.Nullable;
			}

			if (info.ParameterType.IsByRef) {
				if (kind != ParamKind.Normal) throw new ArgumentException(
					$"Parameter does not support [{kind}]\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, kind, info.ParameterType, underlyingType);
		}
	}

	public enum ParamKind
	{
		/// <summary>
		/// Not part of the resulting query's terms.
		/// Same value across a single invocation of a callback.
		/// For example <see cref="ECS.World"/> or <see cref="TimeSpan"/>.
		/// </summary>
		GlobalUnique,

		/// <summary>
		/// Not part of the resulting query's terms.
		/// Unique value for each iterated entity.
		/// For example <see cref="EntityRef"/>.
		/// </summary>
		Unique,

		/// <summary> Passed by value. </summary>
		Normal,

		/// <summary>
		/// Struct passed with the "in" modifier, allowing direct pointer access.
		/// Manually applied with <see cref="InAttribute"/>.
		/// Marks a component as being read from.
		/// </summary>
		In,

		/// <summary>
		/// Struct passed with the "out" modifier, allowing direct pointer access.
		/// Manually applied with <see cref="HasAttribute"/>.
		/// Marks a component as being written to.
		/// </summary>
		Out,

		/// <summary>
		/// Struct passed with the "ref" modifier, allowing direct pointer access.
		/// Marks a component as being read from and written to.
		/// </summary>
		Ref,

		/// <summary>
		/// Only checks for presence.
		/// Manually applied with <see cref="HasAttribute"/>.
		/// Automatically applied for types with <see cref="TagAttribute"/>.
		/// Marks a component as not being accessed.
		/// </summary>
		Has,

		/// <summary>
		/// Struct or class passed as <see cref="T?"/>.
		/// </summary>
		Nullable,

		/// <summary>
		/// Only checks for absence.
		/// Applied with <see cref="NotAttribute"/>.
		/// </summary>
		Not,

		/// <summary>
		/// Matches any terms in a chain of "or" terms.
		/// Applied with <see cref="OrAttribute"/>.
		/// Implies <see cref="Nullable"/>.
		/// </summary>
		Or,
	}
}