JCIL/JavaClass/Method.cs

using System.Collections;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.InteropServices;

namespace JCIL.Java.Class;

[Flags]
public enum MethodAccessFlags : ushort
{
    // Declared public; may be accessed from outside its package.
    Public = 0x0001,

    // Declared private; accessible only within the defining class and other classes belonging to the same nest (§5.4.4).
    Private = 0x0002,

    // Declared protected; may be accessed within subclasses.
    Protected = 0x0004,

    // Declared static.
    Static = 0x0008,

    // Declared final; must not be overridden (§5.4.5).
    Final = 0x0010,

    // Declared synchronized; invocation is wrapped by a monitor use.
    Synchronized = 0x0020,

    // A bridge method, generated by the compiler.
    Bridge = 0x0040,

    // Declared with variable number of arguments.
    Varargs = 0x0080,

    // Declared native; implemented in a language other than the Java programming language.
    Native = 0x0100,

    // Declared abstract; no implementation is provided.
    Abstract = 0x0400,

    // In a class file whose major version number is at least 46 and at most 60: Declared strictfp.
    Strict = 0x0800,

    // Declared synthetic; not present in the source code.
    Synthetic = 0x1000,
}

public sealed record MethodSignature(Class ReturnType, IReadOnlyList<Class> ParamTypes) : ISpecialClassMetadata;

public sealed class Method
{
    public readonly Class DeclaringClass;
    public readonly MethodAccessFlags AccessFlags;
    public readonly string Name;
    public readonly Class Type;
    public readonly IReadOnlyList<Attribute> Attributes;

    private Method(Class declaringClass, MethodAccessFlags flags, string name, Class type, Attribute[] attributes)
    {
        DeclaringClass = declaringClass;
        AccessFlags = flags;
        Name = name;
        Type = type;
        Attributes = attributes;
    }

    public static Method Read(Class parent, Reader reader)
    {
        var flags = (MethodAccessFlags)reader.ReadUInt16();
        var nameIndex = reader.ReadUInt16();
        var descriptorIndex = reader.ReadUInt16();
        var attributesCount = reader.ReadUInt16();

        var name = (string)parent.GetConstant(nameIndex)!;
        var descriptor = (string)parent.GetConstant(descriptorIndex)!;
        var descriptorClass = parent.Loader.GetClass(descriptor.AsMemory());

        var attributes = new Attribute[attributesCount];
        foreach (ref var attribute in attributes.AsSpan())
            attribute = Attribute.Read(parent, reader);
        return new Method(parent, flags, name, descriptorClass, attributes);
    }

    public bool TryGetCode([MaybeNullWhen(false)] out Code code, out OpCodeReader opcodes)
    {
        if (Attributes.OfType<Code>().FirstOrDefault() is { } attr)
        {
            MemoryMarshal.TryGetArray(attr.Bytecode, out var array);
            code = attr;
            opcodes = new OpCodeReader(new MemoryStream(array.Array!));
            return true;
        }

        code = null!;
        opcodes = default;
        return false;
    }

    public bool IsCompatibleWith(MethodSignature signature)
    {
        var methodSig = (MethodSignature)Type.SpecialClassMetadata!;
        
        if (!signature.ReturnType.IsAssignableTo(methodSig.ReturnType))
            return false;

        if ((AccessFlags & MethodAccessFlags.Varargs) == 0 && methodSig.ParamTypes.Count != signature.ParamTypes.Count)
            return false;
        
        for (var i = 0; i < signature.ParamTypes.Count; i++)
        {
            var current = methodSig.ParamTypes[Math.Clamp(i, 0, methodSig.ParamTypes.Count - 1)];
            if (i >= methodSig.ParamTypes.Count - 1) current = ((ArrayMetadata)current.SpecialClassMetadata!).ElementClass;
            if (!signature.ParamTypes[i].IsAssignableTo(current)) return false;
        }

        return true;
    }

    public override string ToString()
    {
        var signature = (MethodSignature)Type.SpecialClassMetadata!;
        return $"{signature.ReturnType} {Name}({string.Join(", ", signature.ParamTypes)})";
    }
}

[SuppressMessage("ReSharper", "InconsistentNaming")]
public enum OpCodeMnemonic : byte
{
    Nop = 0,
    AConstNull = 1,
    IConstM1 = 2,
    IConst0 = 3,
    IConst1 = 4,
    IConst2 = 5,
    IConst3 = 6,
    IConst4 = 7,
    IConst5 = 8,
    LConst0 = 9,
    LConst1 = 10,
    FConst0 = 11,
    FConst1 = 12,
    FConst2 = 13,
    DConst0 = 14,
    DConst1 = 15,
    PushByte = 16,
    PushShort = 17,
    Ldc = 18,
    LdcW = 19,
    Ldc2W = 20,
    ILoad = 21,
    LLoad = 22,
    FLoad = 23,
    DLoad = 24,
    ALoad = 25,
    ILoad0 = 26,
    ILoad1 = 27,
    ILoad2 = 28,
    ILoad3 = 29,
    LLoad0 = 30,
    LLoad1 = 31,
    LLoad2 = 32,
    LLoad3 = 33,
    FLoad0 = 34,
    FLoad1 = 35,
    FLoad2 = 36,
    FLoad3 = 37,
    DLoad0 = 38,
    DLoad1 = 39,
    DLoad2 = 40,
    DLoad3 = 41,
    ALoad0 = 42,
    ALoad1 = 43,
    ALoad2 = 44,
    ALoad3 = 45,
    IALoad = 46,
    LALoad = 47,
    AALoad = 50,
    BALoad = 51,
    CALoad = 52,
    SALoad = 53,
    IStore = 54,
    LStore = 55,
    FStore = 56,
    DStore = 57,
    AStore = 58,
    IStore0 = 59,
    IStore1 = 60,
    IStore2 = 61,
    IStore3 = 62,
    LStore0 = 63,
    LStore1 = 64,
    LStore2 = 65,
    LStore3 = 66,
    AStore0 = 75,
    AStore1 = 76,
    AStore2 = 77,
    IAStore = 79,
    AStore3 = 78,
    LAStore = 80,
    AAStore = 83,
    BAStore = 84,
    CAStore = 85,
    Pop = 87,
    Pop2 = 88,
    Dup = 89,
    DupX1 = 90,
    DupX2 = 91,
    Dup2 = 92,
    IAdd = 96,
    LAdd = 97,
    FAdd = 98,
    DAdd = 99,
    ISub = 100,
    LSub = 101,
    FSub = 102,
    DSub = 103,
    IMul = 104,
    LMul = 105,
    FMul = 106,
    DMul = 107,
    IDiv = 108,
    LDiv = 109,
    FDiv = 110,
    DDiv = 111,
    IRem = 112,
    LRem = 113,
    FRem = 114,
    DRem = 115,
    INeg = 116,
    LNeg = 117,
    IShl = 120,
    LShl = 121,
    IShr = 122,
    LShr = 123,
    IUShr = 124,
    LUShr = 125,
    IAnd = 126,
    LAnd = 127,
    IOr = 128,
    LOr = 129,
    IXor = 130,
    LXor = 131,
    IInc = 132,
    IntToLong = 133,
    IntToFloat = 134,
    IntToDouble = 135,
    LongToInt = 136,
    LongToFloat = 137,
    LongToDouble = 138,
    FloatToInt = 139,
    FloatToLong = 140,
    FloatToDouble = 141,
    DoubleToInt = 142,
    DoubleToLong = 143,
    DoubleToFloat = 144,
    IntToByte = 145,
    IntToChar = 146,
    IntToShort = 147,
    LCmp = 148,
    FCmpL = 149,
    FCmpG = 150,
    DCmpL = 151,
    DCmpG = 152,
    IfEq = 153,
    IfNe = 154,
    IfLt = 155,
    IfGe = 156,
    IfGt = 157,
    IfLe = 158,
    IfICmpEq = 159,
    IfICmpNe = 160,
    IfICmpLt = 161,
    IfICmpGe = 162,
    IfICmpGt = 163,
    IfICmpLe = 164,
    IfACmpEq = 165,
    IfACmpNe = 166,
    Goto = 167,
    TableSwitch = 170,
    LookupSwitch = 171,
    IReturn = 172,
    LReturn = 173,
    FReturn = 174,
    DReturn = 175,
    AReturn = 176,
    Return = 177,
    GetStaticField = 178,
    SetStaticField = 179,
    GetField = 180,
    SetField = 181,
    InvokeVirtual = 182,
    InvokeSpecial = 183,
    InvokeStatic = 184,
    InvokeInterface = 185,
    InvokeDynamic = 186,
    New = 187,
    NewArray = 188,
    ANewArray = 189,
    ArrayLen = 190,
    AThrow = 191,
    CheckCast = 192,
    InstanceOf = 193,
    MonitorEnter = 194,
    MonitorExit = 195,
    IfNull = 198,
    IfNonNull = 199,
}

public struct OpCode
{
    public OpCodeMnemonic Mnemonic;
    public OpCodeParameter P0;
    public OpCodeParameter P1;
    
    public string ToString(Class javaClass)
    {
        switch (Mnemonic)
        {
            case OpCodeMnemonic.GetStaticField:
            {
                var fieldRef = javaClass.GetConstant<FieldRef>(P0.UShort);
                return $"{Mnemonic} {fieldRef.Class}.{fieldRef.Field.Name}";
            }

            case OpCodeMnemonic.InvokeStatic:
            case OpCodeMnemonic.InvokeVirtual:
            case OpCodeMnemonic.InvokeSpecial:
            {
                var methodRef = javaClass.GetConstant<MethodRef>(P0.UShort);
                return $"{Mnemonic} {methodRef.Class}.{methodRef.Method.Name}";
            }
            
            default: return ToString();
        }
    }

    public override string ToString()
    {
        var p0 = P0.Object is null ? P0.ULong.ToString() : P0.Object.ToString();
        var p1 = P1.Object is null ? P1.ULong.ToString() : P1.Object.ToString();
        return $"{Mnemonic} {p0}, {p1}";
    }
};

[StructLayout(LayoutKind.Explicit)]
public struct OpCodeParameter
{
    [FieldOffset(0)] public byte Byte;
    [FieldOffset(0)] public ushort UShort;
    [FieldOffset(0)] public uint UInt;
    [FieldOffset(0)] public ulong ULong;
    [FieldOffset(0)] public float Float;
    [FieldOffset(0)] public double Double;
    [FieldOffset(8)] public object? Object;
}

public struct OpCodeReader(Stream bytecode) : IEnumerator<OpCode>
{
    private OpCode _current;
    public Reader Reader = new(bytecode);

    public bool MoveNext()
    {
        var b = Reader.Stream.ReadByte();
        if (b == -1)
        {
            _current = default;
            return false;
        }
        _current = new OpCode { Mnemonic = (OpCodeMnemonic)(byte)b };
        switch (_current.Mnemonic)
        {
            case OpCodeMnemonic.PushByte:
            case OpCodeMnemonic.ILoad:
            case OpCodeMnemonic.LLoad:
            case OpCodeMnemonic.FLoad:
            case OpCodeMnemonic.DLoad:
            case OpCodeMnemonic.ALoad:
            case OpCodeMnemonic.IStore:
            case OpCodeMnemonic.AStore:
            case OpCodeMnemonic.LStore:
            case OpCodeMnemonic.FStore:
            case OpCodeMnemonic.DStore:
            case OpCodeMnemonic.NewArray:
            case OpCodeMnemonic.Ldc:
            {
                _current.P0.UShort = Reader.ReadUInt8();
                return true;
            }

            case OpCodeMnemonic.IInc:
            {
                _current.P0.UShort = Reader.ReadUInt8();
                _current.P1.Byte = Reader.ReadUInt8();
                return true;
            }

            case OpCodeMnemonic.PushShort:
            case OpCodeMnemonic.LdcW:
            case OpCodeMnemonic.Ldc2W:
            case OpCodeMnemonic.InvokeVirtual:
            case OpCodeMnemonic.InvokeSpecial:
            case OpCodeMnemonic.InvokeStatic:
            case OpCodeMnemonic.IfNull:
            case OpCodeMnemonic.IfNonNull:
            case OpCodeMnemonic.IfEq:
            case OpCodeMnemonic.IfNe:
            case OpCodeMnemonic.IfLt:
            case OpCodeMnemonic.IfLe:
            case OpCodeMnemonic.IfGe:
            case OpCodeMnemonic.IfGt:
            case OpCodeMnemonic.IfICmpEq:
            case OpCodeMnemonic.IfICmpNe:
            case OpCodeMnemonic.IfICmpLt:
            case OpCodeMnemonic.IfICmpLe:
            case OpCodeMnemonic.IfICmpGe:
            case OpCodeMnemonic.IfICmpGt:
            case OpCodeMnemonic.IfACmpEq:
            case OpCodeMnemonic.IfACmpNe:
            case OpCodeMnemonic.Goto:
            case OpCodeMnemonic.New:
            case OpCodeMnemonic.GetField:
            case OpCodeMnemonic.SetField:
            case OpCodeMnemonic.GetStaticField:
            case OpCodeMnemonic.SetStaticField:
            case OpCodeMnemonic.CheckCast:
            case OpCodeMnemonic.InstanceOf:
            case OpCodeMnemonic.ANewArray:
            {
                _current.P0.UShort = Reader.ReadUInt16();
                return true;
            }

            case OpCodeMnemonic.InvokeInterface:
            case OpCodeMnemonic.InvokeDynamic:
            {
                _current.P0.UShort = Reader.ReadUInt16();
                _current.P1.UShort = Reader.ReadUInt16();
                return true;
            }

            case OpCodeMnemonic.TableSwitch:
            {
                while (Reader.Stream.Position % 4 != 0)
                    if (Reader.Stream.ReadByte() < 0)
                        return false;

                var table = new TableSwitch();
                _current.P0.Object = table;
                table.Default = Reader.ReadInt32();
                table.Low = Reader.ReadInt32();
                table.High = Reader.ReadInt32();
                table.Offsets = new int[table.High - table.Low + 1];
                foreach (ref var offset in table.Offsets.AsSpan())
                    offset = Reader.ReadInt32();

                return true;
            }

            case OpCodeMnemonic.LookupSwitch:
            {
                while (Reader.Stream.Position % 4 != 0)
                    if (Reader.Stream.ReadByte() < 0)
                        return false;

                var table = new LookupSwitch();
                _current.P0.Object = table;
                table.Default = Reader.ReadInt32();
                table.Entries = new KeyValuePair<int, int>[Reader.ReadUInt32()];
                foreach (ref var kv in table.Entries.AsSpan())
                    kv = new KeyValuePair<int, int>(Reader.ReadInt32(), Reader.ReadInt32());

                return true;
            }
            case OpCodeMnemonic.Nop:
            case OpCodeMnemonic.AConstNull:
            case OpCodeMnemonic.IConstM1:
            case OpCodeMnemonic.IConst0:
            case OpCodeMnemonic.IConst1:
            case OpCodeMnemonic.IConst2:
            case OpCodeMnemonic.IConst3:
            case OpCodeMnemonic.IConst4:
            case OpCodeMnemonic.IConst5:
            case OpCodeMnemonic.LConst0:
            case OpCodeMnemonic.LConst1:
            case OpCodeMnemonic.FConst0:
            case OpCodeMnemonic.FConst1:
            case OpCodeMnemonic.FConst2:
            case OpCodeMnemonic.DConst0:
            case OpCodeMnemonic.DConst1:
            case OpCodeMnemonic.ILoad0:
            case OpCodeMnemonic.ILoad1:
            case OpCodeMnemonic.ILoad2:
            case OpCodeMnemonic.ILoad3:
            case OpCodeMnemonic.LLoad0:
            case OpCodeMnemonic.LLoad1:
            case OpCodeMnemonic.LLoad2:
            case OpCodeMnemonic.LLoad3:
            case OpCodeMnemonic.FLoad0:
            case OpCodeMnemonic.FLoad1:
            case OpCodeMnemonic.FLoad2:
            case OpCodeMnemonic.FLoad3:
            case OpCodeMnemonic.DLoad0:
            case OpCodeMnemonic.DLoad1:
            case OpCodeMnemonic.DLoad2:
            case OpCodeMnemonic.DLoad3:
            case OpCodeMnemonic.ALoad0:
            case OpCodeMnemonic.ALoad1:
            case OpCodeMnemonic.ALoad2:
            case OpCodeMnemonic.ALoad3:
            case OpCodeMnemonic.IALoad:
            case OpCodeMnemonic.LALoad:
            case OpCodeMnemonic.AALoad:
            case OpCodeMnemonic.BALoad:
            case OpCodeMnemonic.CALoad:
            case OpCodeMnemonic.SALoad:
            case OpCodeMnemonic.IStore0:
            case OpCodeMnemonic.IStore1:
            case OpCodeMnemonic.IStore2:
            case OpCodeMnemonic.IStore3:
            case OpCodeMnemonic.LStore0:
            case OpCodeMnemonic.LStore1:
            case OpCodeMnemonic.LStore2:
            case OpCodeMnemonic.LStore3:
            case OpCodeMnemonic.AStore0:
            case OpCodeMnemonic.AStore1:
            case OpCodeMnemonic.AStore2:
            case OpCodeMnemonic.AStore3:
            case OpCodeMnemonic.LAStore:
            case OpCodeMnemonic.IAStore:
            case OpCodeMnemonic.AAStore:
            case OpCodeMnemonic.BAStore:
            case OpCodeMnemonic.CAStore:
            case OpCodeMnemonic.Pop:
            case OpCodeMnemonic.Pop2:
            case OpCodeMnemonic.Dup:
            case OpCodeMnemonic.DupX1:
            case OpCodeMnemonic.DupX2:
            case OpCodeMnemonic.Dup2:
            case OpCodeMnemonic.IAdd:
            case OpCodeMnemonic.LAdd:
            case OpCodeMnemonic.FAdd:
            case OpCodeMnemonic.DAdd:
            case OpCodeMnemonic.ISub:
            case OpCodeMnemonic.LSub:
            case OpCodeMnemonic.FSub:
            case OpCodeMnemonic.DSub:
            case OpCodeMnemonic.IMul:
            case OpCodeMnemonic.LMul:
            case OpCodeMnemonic.FMul:
            case OpCodeMnemonic.DMul:
            case OpCodeMnemonic.IDiv:
            case OpCodeMnemonic.LDiv:
            case OpCodeMnemonic.FDiv:
            case OpCodeMnemonic.DDiv:
            case OpCodeMnemonic.IRem:
            case OpCodeMnemonic.LRem:
            case OpCodeMnemonic.FRem:
            case OpCodeMnemonic.DRem:
            case OpCodeMnemonic.INeg:
            case OpCodeMnemonic.LNeg:
            case OpCodeMnemonic.IShl:
            case OpCodeMnemonic.LShl:
            case OpCodeMnemonic.IShr:
            case OpCodeMnemonic.LShr:
            case OpCodeMnemonic.IUShr:
            case OpCodeMnemonic.LUShr:
            case OpCodeMnemonic.IAnd:
            case OpCodeMnemonic.LAnd:
            case OpCodeMnemonic.IOr:
            case OpCodeMnemonic.LOr:
            case OpCodeMnemonic.IXor:
            case OpCodeMnemonic.LXor:
            case OpCodeMnemonic.IntToLong:
            case OpCodeMnemonic.IntToFloat:
            case OpCodeMnemonic.IntToDouble:
            case OpCodeMnemonic.IntToByte:
            case OpCodeMnemonic.IntToChar:
            case OpCodeMnemonic.IntToShort:
            case OpCodeMnemonic.LongToInt:
            case OpCodeMnemonic.LongToFloat:
            case OpCodeMnemonic.LongToDouble:
            case OpCodeMnemonic.FloatToInt:
            case OpCodeMnemonic.FloatToLong:
            case OpCodeMnemonic.FloatToDouble:
            case OpCodeMnemonic.DoubleToInt:
            case OpCodeMnemonic.DoubleToLong:
            case OpCodeMnemonic.DoubleToFloat:
            case OpCodeMnemonic.LCmp:
            case OpCodeMnemonic.FCmpG:
            case OpCodeMnemonic.FCmpL:
            case OpCodeMnemonic.DCmpG:
            case OpCodeMnemonic.DCmpL:
            case OpCodeMnemonic.IReturn:
            case OpCodeMnemonic.AReturn:
            case OpCodeMnemonic.LReturn:
            case OpCodeMnemonic.FReturn:
            case OpCodeMnemonic.DReturn:
            case OpCodeMnemonic.Return:
            case OpCodeMnemonic.ArrayLen:
            case OpCodeMnemonic.AThrow:
            case OpCodeMnemonic.MonitorEnter:
            case OpCodeMnemonic.MonitorExit:
                return true;

            default:
                throw new InvalidDataException($"{_current.Mnemonic} is not a valid opcode");
        }
    }

    public void Reset()
    {
        Reader.Stream.Position = 0;
        _current = default;
    }

    public OpCode Current => _current;
    object IEnumerator.Current => _current;

    public void Dispose()
    {
        Reader.Stream.Dispose();
    }
}

public sealed class TableSwitch
{
    public int Default;
    public int Low;
    public int High;
    public int[] Offsets = [];
}

public sealed class LookupSwitch
{
    public int Default;
    public KeyValuePair<int, int>[] Entries = [];
}