JCIL/CLI/AssemblyBuilder.cs

using System.Collections.Concurrent;
using System.Reflection;
using System.Reflection.Emit;
using System.Runtime.CompilerServices;
using JCIL.CLI.Types;
using JCIL.CLI.Types.java.lang;
using JCIL.Java.Class;
using OpCode = JCIL.Java.Class.OpCode;

namespace JCIL.CLI;

public sealed class AssemblyBuilder
{
	public readonly System.Reflection.Emit.AssemblyBuilder Assembly;
	private readonly ModuleBuilder _module;
	private readonly ConcurrentDictionary<Class, TypeSurrogate> _classes = [];
	private readonly ConcurrentDictionary<Field, FieldInfo> _fields = [];
	private readonly ConcurrentDictionary<Method, LazyMethod> _methods = [];

	public int CompiledMethods { get; private set; }
	public int TotalMethods => _methods.Count;

	public AssemblyBuilder(AssemblyName name, bool runnable)
	{
		Assembly = runnable
			? System.Reflection.Emit.AssemblyBuilder.DefineDynamicAssembly(name, AssemblyBuilderAccess.RunAndCollect)
			: new PersistedAssemblyBuilder(name, typeof(int).Assembly);

		_module = Assembly.DefineDynamicModule("JCIL_Output");
	}

	public TypeSurrogate MakeType(Class javaClass)
	{
		if (_classes.TryGetValue(javaClass, out var type))
			return type;

		switch (javaClass.SpecialClassMetadata)
		{
			case IntrinsicMetadata intrinsic:
			{
				switch (intrinsic.Type)
				{
					case IntrinsicType.Byte:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Byte(c));
					case IntrinsicType.Char:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Char(c));
					case IntrinsicType.Int:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Integer(c));
					case IntrinsicType.Long:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Long(c));
					case IntrinsicType.Boolean:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Boolean(c));
					case IntrinsicType.Float:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Float(c));
					case IntrinsicType.Double:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Double(c));
					case IntrinsicType.Void:
						return _classes.GetOrAdd(javaClass, c => new JCIL.CLI.Types.Java.Lang.Void(c));
					default: throw new NotImplementedException($"Unimplemented Intrinsic {intrinsic.Type}");
				}
			}
			case ArrayMetadata array:
				return _classes.GetOrAdd(javaClass, c => new ArrayOf(MakeType(array.ElementClass), c));
		}

		var newClass = (NewClass)_classes.GetOrAdd(javaClass, _ =>
		{
			TypeAttributes flags = default;
			if ((javaClass.AccessFlags & ClassAccessFlags.Public) != 0) flags |= TypeAttributes.Public;
			if ((javaClass.AccessFlags & ClassAccessFlags.Abstract) != 0) flags |= TypeAttributes.Abstract;
			if ((javaClass.AccessFlags & ClassAccessFlags.Final) != 0) flags |= TypeAttributes.Sealed;
			if ((javaClass.AccessFlags & ClassAccessFlags.Interface) != 0) flags |= TypeAttributes.Interface;
			var builder = _module.DefineType(javaClass.Name.ToString(), flags);
			return new NewClass(this, builder, javaClass);
		});
		newClass.CreateFields();
		return newClass;
	}

	internal void CompileOpCodes(Class javaClass, MethodBase methodBase, ILGenerator il, Code code,
		OpCodeReader opcodes)
	{
		var thisType = methodBase.DeclaringType!;
		Debug.Writer.WriteLine($"Compiling {thisType.FullName}.{methodBase.Name}");
		Debug.Writer.Indent++;

		var labels = CreateLabels(il, opcodes);
		var position = opcodes.Reader.Stream.Position;
		var stack = new Stack<Value>(code.MaxStack);
		var locals = new Value[code.MaxLocals];
		var localCount = 0;

		{
			var idx = 0;
			if (!methodBase.IsStatic)
				locals[idx] = Value.Parameter(thisType, idx++);

			foreach (var parameter in methodBase.GetParameters())
				locals[idx] = Value.Parameter(parameter.ParameterType, idx++);
		}

		void CompileOpCode(long position, OpCode opcode)
		{
			Debug.Writer.WriteLine($"IL_{position:X04}: {opcode.ToString(javaClass)}");

			Label JmpLabel()
			{
				return labels[position + (short)opcode.P0.UShort];
			}

			void CheckCast(Value value, Type type)
			{
				if (value.Type == typeof(int) && type == typeof(byte))
				{
					il.Emit(OpCodes.Conv_I1);
				}
				else if (!value.Type.IsAssignableTo(type))
				{
					throw new InvalidOperationException(
						$"Value of type {value.Type} cannot be assigned to one of type {type}.");
				}
			}

			switch (opcode.Mnemonic)
			{
				case OpCodeMnemonic.AConstNull:
					stack.Push(new Value(typeof(object)));
					il.Emit(OpCodes.Ldnull);
					break;

				case OpCodeMnemonic.IConst0:
				case OpCodeMnemonic.IConst1:
				case OpCodeMnemonic.IConst2:
				case OpCodeMnemonic.IConst3:
				case OpCodeMnemonic.IConst4:
				case OpCodeMnemonic.IConst5:
					stack.Push(new Value(opcode.Mnemonic - OpCodeMnemonic.IConst0));
					break;

				case OpCodeMnemonic.LConst0:
					stack.Push(new Value(typeof(long)));
					il.Emit(OpCodes.Ldc_I8, 0L);
					break;

				case OpCodeMnemonic.LConst1:
					stack.Push(new Value(typeof(long)));
					il.Emit(OpCodes.Ldc_I8, 1L);
					break;

				case OpCodeMnemonic.PushByte: stack.Push(new Value(opcode.P0.Byte).Load(il)); break;
				case OpCodeMnemonic.PushShort: stack.Push(new Value((short)opcode.P0.UShort).Load(il)); break;

				case OpCodeMnemonic.Ldc:
				case OpCodeMnemonic.LdcW:
				{
					var constant = javaClass.GetConstant<object>(opcode.P0.UShort);
					switch (constant)
					{
						case int val: stack.Push(new Value(val).Load(il)); break;
						case string val: stack.Push(new Value(val).Load(il)); break;
						case Class val: stack.Push(Value.TypeRef(MakeType(val).Type).Load(il)); break;
						default:
						{
							throw new NotImplementedException(
								$"Unknown opcode: {opcode.Mnemonic} {constant.GetType()}");
						}
					}

					break;
				}

				case OpCodeMnemonic.ILoad:
				case OpCodeMnemonic.ILoad0:
				case OpCodeMnemonic.ILoad1:
				case OpCodeMnemonic.ILoad2:
				case OpCodeMnemonic.ILoad3:
				{
					var idx = opcode.Mnemonic switch
					{
						OpCodeMnemonic.ILoad0 => 0,
						OpCodeMnemonic.ILoad1 => 1,
						OpCodeMnemonic.ILoad2 => 2,
						OpCodeMnemonic.ILoad3 => 3,
						OpCodeMnemonic.ILoad => opcode.P0.UShort,
						_ => -1,
					};

					if (locals[idx].Type != typeof(int))
						throw new InvalidOperationException($"Local {idx} is not an integer.");

					stack.Push(locals[idx].Load(il));
					break;
				}

				case OpCodeMnemonic.ALoad:
				case OpCodeMnemonic.ALoad0:
				case OpCodeMnemonic.ALoad1:
				case OpCodeMnemonic.ALoad2:
				case OpCodeMnemonic.ALoad3:
				{
					var idx = opcode.Mnemonic switch
					{
						OpCodeMnemonic.ALoad0 => 0,
						OpCodeMnemonic.ALoad1 => 1,
						OpCodeMnemonic.ALoad2 => 2,
						OpCodeMnemonic.ALoad3 => 3,
						OpCodeMnemonic.ALoad => opcode.P0.UShort,
						_ => -1,
					};

					if (locals[idx] is not { Type.IsClass: true })
						throw new InvalidOperationException($"Local {idx} is not a class.");

					stack.Push(locals[idx].Load(il));
					break;
				}

				case OpCodeMnemonic.AALoad:
				{
					var value = stack.Pop();

					if (!value.Type.IsArray)
						throw new InvalidOperationException("Value is not an array.");

					if (value.Type.GetElementType() is not { IsClass: true } elementType)
						throw new InvalidOperationException("Array element is not an object.");

					il.Emit(OpCodes.Ldelem_Ref, elementType);
					stack.Push(new Value(elementType));
					break;
				}

				case OpCodeMnemonic.LStore:
				case OpCodeMnemonic.LStore0:
				case OpCodeMnemonic.LStore1:
				case OpCodeMnemonic.LStore2:
				case OpCodeMnemonic.LStore3:
				{
					var value = stack.Pop();
					CheckCast(value, typeof(long));
					ref var loc = ref locals[opcode.Mnemonic switch
					{
						OpCodeMnemonic.LStore => opcode.P0.UShort,
						_ => opcode.Mnemonic - OpCodeMnemonic.LStore0,
					}];
					if (loc.Type != typeof(long))
					{
						il.DeclareLocal(typeof(long));
						loc = Value.Local(typeof(long), localCount++);
					}
					loc.Store(il);
					break;
				}

				case OpCodeMnemonic.Pop:
				{
					stack.Pop();
					il.Emit(OpCodes.Pop);
					break;
				}

				case OpCodeMnemonic.Dup:
				{
					var value = stack.Pop();
					stack.Push(value);
					stack.Push(value);
					il.Emit(OpCodes.Dup);
					break;
				}

				case OpCodeMnemonic.IfEq:
				{
					var value = stack.Pop();
					if (value.Type != typeof(int) && value.Type != typeof(bool))
						throw new InvalidOperationException($"Expected integer, found {value.Type}.");

					il.Emit(OpCodes.Brfalse, JmpLabel());
					break;
				}

				case OpCodeMnemonic.IfACmpNe:
				{
					if (!stack.Pop().Type.IsClass || !stack.Pop().Type.IsClass)
						throw new InvalidOperationException("Value is not an object.");

					il.Emit(OpCodes.Ceq);
					il.Emit(OpCodes.Brfalse, JmpLabel());
					break;
				}

				case OpCodeMnemonic.IReturn:
				case OpCodeMnemonic.LReturn:
				case OpCodeMnemonic.FReturn:
				case OpCodeMnemonic.DReturn:
				case OpCodeMnemonic.AReturn:
				{
					var value = stack.Pop();
					il.Emit(OpCodes.Ret);
					break;
				}

				case OpCodeMnemonic.Return:
				{
					il.Emit(OpCodes.Ret);
					break;
				}

				case OpCodeMnemonic.GetStaticField:
				{
					var fieldRef = javaClass.GetConstant<FieldRef>(opcode.P0.UShort);
					var parent = MakeType(fieldRef.Class);
					var field = parent.GetField(fieldRef.Field.Name);
					stack.Push(new Value(field).Load(il));
					break;
				}

				case OpCodeMnemonic.GetField:
				{
					var fieldRef = javaClass.GetConstant<FieldRef>(opcode.P0.UShort);
					var parent = MakeType(fieldRef.Class);
					var field = parent.GetField(fieldRef.Field.Name);

					var value = stack.Pop();
					if (!value.Type.IsAssignableTo(parent.Type))
						throw new InvalidOperationException($"Value is not compatible with type {parent}.");

					stack.Push(new Value(field).Load(il));
					break;
				}

				case OpCodeMnemonic.SetField:
				{
					var fieldRef = javaClass.GetConstant<FieldRef>(opcode.P0.UShort);
					var parent = MakeType(fieldRef.Class);
					var field = parent.GetField(fieldRef.Field.Name);
					var value = stack.Pop();

					CheckCast(value, field.FieldType);
					il.Emit(OpCodes.Stfld, field);
					break;
				}

				case OpCodeMnemonic.InvokeVirtual:
				case OpCodeMnemonic.InvokeSpecial:
				case OpCodeMnemonic.InvokeStatic:
				{
					var methodRef = javaClass.GetConstant<MethodRef>(opcode.P0.UShort);
					var parent = MakeType(methodRef.Class);

					var signature = (MethodSignature)javaClass.Loader.GetClass(methodRef.Method.Descriptor.AsMemory())
						.SpecialClassMetadata!;

					stack.Push(new Value(parent.CallMethod(methodRef.Method.Name, signature, il)));
					break;
				}

				case OpCodeMnemonic.New:
				{
					var type = MakeType(javaClass.GetConstant<Class>(opcode.P0.UShort));
					var posBackup = opcodes.Reader.Stream.Position;
					var stackBackup = new Stack<Value>(stack);
					opcodes.MoveNext();
					if (opcodes.Current.Mnemonic == OpCodeMnemonic.Dup)
					{
						opcodes.MoveNext();
						position = opcodes.Reader.Stream.Position;
						while (opcodes.Current.Mnemonic != OpCodeMnemonic.InvokeSpecial)
						{
							CompileOpCode(position, opcodes.Current);
							opcodes.MoveNext();
							position = opcodes.Reader.Stream.Position;
						}

						var methodRef = javaClass.GetConstant<MethodRef>(opcodes.Current.P0.UShort);
						var parent = MakeType(methodRef.Class);

						var signature = (MethodSignature)javaClass.Loader
							.GetClass(methodRef.Method.Descriptor.AsMemory())
							.SpecialClassMetadata!;

						var ctor = (ConstructorInfo)parent.GetMethod(methodRef.Method.Name, signature);
						while (stack.Count > stackBackup.Count) stack.Pop();
						il.Emit(OpCodes.Newobj, ctor);
						stack.Push(new Value(parent.Type));
						break;
					}

					opcodes.Reader.Stream.Position = posBackup;
					stack = stackBackup;
					var func = RuntimeHelpers.GetUninitializedObject;
					il.Emit(OpCodes.Ldtoken, type.Type);
					il.Emit(OpCodes.Call, func.Method);
					stack.Push(new Value(type.Type));
					break;
				}

				case OpCodeMnemonic.NewArray:
				{
					var elementType = opcode.P0.Byte switch
					{
						4 => typeof(bool),
						5 => typeof(char),
						6 => typeof(float),
						7 => typeof(double),
						8 => typeof(byte),
						9 => typeof(short),
						10 => typeof(int),
						11 => typeof(long),
						_ => typeof(void),
					};
					il.Emit(OpCodes.Newarr, elementType);
					stack.Push(new Value(elementType.MakeArrayType()));
					break;
				}

				case OpCodeMnemonic.IfNull:
				{
					if (!stack.Pop().Type.IsClass)
						throw new InvalidOperationException("Value is not an object.");

					il.Emit(OpCodes.Brfalse, JmpLabel());
					break;
				}

				case OpCodeMnemonic.IfNonNull:
				{
					if (!stack.Pop().Type.IsClass)
						throw new InvalidOperationException("Value is not an object.");

					il.Emit(OpCodes.Brtrue, JmpLabel());
					break;
				}

				default:
				{
					Console.ForegroundColor = ConsoleColor.Red;
					Console.WriteLine($"Unknown opcode: {opcode.Mnemonic}");
					Console.ResetColor();
					il.Emit(OpCodes.Nop);
					break;
				}
			}
		}

		for (; opcodes.MoveNext();)
		{
			il.MarkLabel(labels[position]);
			var opcode = opcodes.Current;
			CompileOpCode(position, opcode);
			position = opcodes.Reader.Stream.Position;
		}

		Debug.Writer.Indent--;
		CompiledMethods += 1;
	}

	private static Dictionary<long, Label> CreateLabels(ILGenerator il, OpCodeReader opcodes)
	{
		var start = opcodes.Reader.Stream.Position;
		var position = start;
		var labels = new Dictionary<long, Label>();
		while (opcodes.MoveNext())
		{
			labels.Add(position, il.DefineLabel());
			position = opcodes.Reader.Stream.Position;
		}

		opcodes.Reader.Stream.Position = start;
		return labels;
	}
}

internal class LazyMethod(MethodBase method, Action compile)
{
	public bool Compiled { get; private set; }

	public MethodBase Method
	{
		get
		{
			if (!Compiled)
			{
				Compiled = true;
				compile();
			}

			return method;
		}
	}
}

internal readonly struct Value
{
	public delegate Value LoadDelegate(ILGenerator il, OpCodeParameter data);

	public delegate void StoreDelegate(ILGenerator il, OpCodeParameter data);

	public readonly Type Type;
	private LoadDelegate? _load { get; init; }
	private StoreDelegate? _store { get; init; }
	private OpCodeParameter _data { get; init; }
	public static implicit operator Type(Value v) => v.Type;

	public Value Load(ILGenerator il)
	{
		return _load is { } load ? load(il, _data) : throw new InvalidOperationException("Value cannot be loaded.");
	}

	public void Store(ILGenerator il)
	{
		if (_store is not { } store) throw new InvalidOperationException("Value cannot be stored.");
		store(il, _data);
	}

	public Value(Type type)
	{
		Type = type;
	}

	public Value(int value)
	{
		Type = typeof(int);
		_data = new OpCodeParameter { UInt = (uint)value };
		_load = (il, data) =>
		{
			var val = (int)data.UInt;
			switch (val)
			{
				case 0: il.Emit(OpCodes.Ldc_I4_0); break;
				case 1: il.Emit(OpCodes.Ldc_I4_1); break;
				case 2: il.Emit(OpCodes.Ldc_I4_2); break;
				case 3: il.Emit(OpCodes.Ldc_I4_3); break;
				case 4: il.Emit(OpCodes.Ldc_I4_4); break;
				case 5: il.Emit(OpCodes.Ldc_I4_5); break;
				case 6: il.Emit(OpCodes.Ldc_I4_6); break;
				case 7: il.Emit(OpCodes.Ldc_I4_7); break;
				case 8: il.Emit(OpCodes.Ldc_I4_8); break;
				case -1: il.Emit(OpCodes.Ldc_I4_M1); break;
				case <= 255: il.Emit(OpCodes.Ldc_I4_S, (sbyte)val); break;
				default: il.Emit(OpCodes.Ldc_I4, val); break;
			}

			return new Value(val);
		};
	}

	public Value(string value)
	{
		Type = typeof(string);
		_data = new OpCodeParameter { Object = value };
		_load = (il, data) =>
		{
			var val = Unsafe.As<string>(data.Object)!;
			il.Emit(OpCodes.Ldstr, val);
			return new Value(val);
		};
	}

	public Value(FieldInfo field)
	{
		Type = field.FieldType;
		_data = new OpCodeParameter { Object = field };
		_load = (il, data) =>
		{
			var val = Unsafe.As<FieldInfo>(data.Object)!;
			il.Emit(val.IsStatic ? OpCodes.Ldsfld : OpCodes.Ldfld, val);
			return new Value(val);
		};
	}

	public static Value TypeRef(Type type)
	{
		return new Value(typeof(Type))
		{
			_data = new OpCodeParameter { Object = type },
			_load = (il, data) =>
			{
				var val = Unsafe.As<Type>(data.Object)!;
				il.Emit(OpCodes.Ldtoken, val);
				return new Value(val);
			},
		};
	}

	public static Value Parameter(Type type, int idx)
	{
		return new Value(type)
		{
			_data = new OpCodeParameter { UInt = (uint)idx, Object = type },
			_load = (il, data) =>
			{
				var idx = (int)data.UInt;
				var type = Unsafe.As<Type>(data.Object);
				switch (idx)
				{
					case 0: il.Emit(OpCodes.Ldarg_0); break;
					case 1: il.Emit(OpCodes.Ldarg_1); break;
					case 2: il.Emit(OpCodes.Ldarg_2); break;
					case 3: il.Emit(OpCodes.Ldarg_3); break;
					case <= 255: il.Emit(OpCodes.Ldarg_S, (byte)idx); break;
					default: il.Emit(OpCodes.Ldarg, (short)idx); break;
				}

				return Value.Parameter(type, idx);
			},
			_store = (il, data) =>
			{
				var idx = (int)data.UInt;
				switch (idx)
				{
					case <= 255: il.Emit(OpCodes.Starg_S, (byte)idx); break;
					default: il.Emit(OpCodes.Starg, (short)idx); break;
				}
			},
		};
	}

	public static Value Local(Type type, int idx)
	{
		return new Value(type)
		{
			_data = new OpCodeParameter { UInt = (uint)idx, Object = type },
			_load = (il, data) =>
			{
				var idx = (int)data.UInt;
				var type = Unsafe.As<Type>(data.Object);
				switch (idx)
				{
					case 0: il.Emit(OpCodes.Ldloc_0); break;
					case 1: il.Emit(OpCodes.Ldloc_1); break;
					case 2: il.Emit(OpCodes.Ldloc_2); break;
					case 3: il.Emit(OpCodes.Ldloc_3); break;
					case <= 255: il.Emit(OpCodes.Ldloc_S, (byte)idx); break;
					default: il.Emit(OpCodes.Ldloc, (short)idx); break;
				}

				return Local(type, idx);
			},
			_store = (il, data) =>
			{
				var idx = (int)data.UInt;
				switch (idx)
				{
					case 0: il.Emit(OpCodes.Stloc_0); break;
					case 1: il.Emit(OpCodes.Stloc_1); break;
					case 2: il.Emit(OpCodes.Stloc_2); break;
					case 3: il.Emit(OpCodes.Stloc_3); break;
					case <= 255: il.Emit(OpCodes.Stloc_S, (byte)idx); break;
					default: il.Emit(OpCodes.Stloc, (short)idx); break;
				}
			},
		};
	}
}