Lua-CSharp/src/Lua/LuaTable.cs

using System.Runtime.CompilerServices;
using Lua.Internal;

namespace Lua;

public sealed class LuaTable
{
    public LuaTable() : this(8, 8)
    {
    }

    public LuaTable(int arrayCapacity, int dictionaryCapacity)
    {
        array = new LuaValue[arrayCapacity];
        dictionary = new(dictionaryCapacity);
    }

    LuaValue[] array;
    Dictionary<LuaValue, LuaValue> dictionary;
    LuaTable? metatable;

    public LuaValue this[LuaValue key]
    {
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        get
        {
            if (key.Type is LuaValueType.Nil) ThrowIndexIsNil();

            if (TryGetInteger(key, out var index))
            {
                if (index > 0 && index <= array.Length)
                {
                    // Arrays in Lua are 1-origin...
                    return MemoryMarshalEx.UnsafeElementAt(array, index - 1);
                }
            }

            if (dictionary.TryGetValue(key, out var value)) return value;
            return LuaValue.Nil;
        }
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        set
        {
            if (key.Type is LuaValueType.Number && double.IsNaN(key.UnsafeRead<double>()))
            {
                ThrowIndexIsNaN();
            }

            if (TryGetInteger(key, out var index))
            {
                if (0 < index && index <= Math.Max(array.Length * 2, 8))
                {
                    EnsureArrayCapacity(index);
                    MemoryMarshalEx.UnsafeElementAt(array, index - 1) = value;
                    return;
                }
            }

            dictionary[key] = value;
        }
    }

    public int HashMapCount
    {
        get => dictionary.Count(x => x.Value.Type is not LuaValueType.Nil);
    }

    public int ArrayLength
    {
        get
        {
            for (int i = 0; i < array.Length; i++)
            {
                if (array[i].Type is LuaValueType.Nil) return i;
            }
            return array.Length;
        }
    }

    public LuaTable? Metatable
    {
        get => metatable;
        set => metatable = value;
    }

    public bool TryGetValue(LuaValue key, out LuaValue value)
    {
        if (key.Type is LuaValueType.Nil)
        {
            value = default;
            return false;
        }

        if (TryGetInteger(key, out var index))
        {
            if (index > 0 && index <= array.Length)
            {
                value = MemoryMarshalEx.UnsafeElementAt(array, index - 1);
                return value.Type is not LuaValueType.Nil;
            }
        }

        return dictionary.TryGetValue(key, out value) && value.Type is not LuaValueType.Nil;
    }

    public bool ContainsKey(LuaValue key)
    {
        if (key.Type is LuaValueType.Nil)
        {
            return false;
        }

        if (TryGetInteger(key, out var index))
        {
            return index > 0 && index <= array.Length &&
                MemoryMarshalEx.UnsafeElementAt(array, index - 1).Type != LuaValueType.Nil;
        }

        return dictionary.TryGetValue(key, out var value) && value.Type is not LuaValueType.Nil;
    }

    public LuaValue RemoveAt(int index)
    {
        if (index <= 0 || index > array.Length)
        {
            throw new IndexOutOfRangeException();
        }

        var arrayIndex = index - 1;
        var value = MemoryMarshalEx.UnsafeElementAt(array, arrayIndex);

        if (arrayIndex < array.Length - 1)
        {
            array.AsSpan(arrayIndex + 1).CopyTo(array.AsSpan(arrayIndex));
        }
        
        MemoryMarshalEx.UnsafeElementAt(array, array.Length - 1) = default;

        return value;
    }

    public void Insert(int index, LuaValue value)
    {
        if (index <= 0 || index > array.Length + 1)
        {
            throw new IndexOutOfRangeException();
        }

        var arrayIndex = index - 1;
        EnsureArrayCapacity(array.Length + 1);

        if (arrayIndex != array.Length - 1)
        {
            array.AsSpan(arrayIndex, array.Length - arrayIndex - 1).CopyTo(array.AsSpan(arrayIndex + 1));
        }
        
        array[arrayIndex] = value;
    }

    public bool TryGetNext(LuaValue key, out KeyValuePair<LuaValue, LuaValue> pair)
    {
        var index = -1;
        if (key.Type is LuaValueType.Nil)
        {
            index = 0;
        }
        else if (TryGetInteger(key, out var integer) && integer > 0 && integer <= array.Length)
        {
            index = integer;
        }

        if (index != -1)
        {
            var span = array.AsSpan(index);
            for (int i = 0; i < span.Length; i++)
            {
                if (span[i].Type is not LuaValueType.Nil)
                {
                    pair = new(index + i + 1, span[i]);
                    return true;
                }
            }

            foreach (var kv in dictionary)
            {
                if (kv.Value.Type is not LuaValueType.Nil)
                {
                    pair = kv;
                    return true;
                }
            }
        }
        else
        {
            var foundKey = false;
            foreach (var kv in dictionary)
            {
                if (foundKey && kv.Value.Type is not LuaValueType.Nil)
                {
                    pair = kv;
                    return true;
                }

                if (kv.Key.Equals(key))
                {
                    foundKey = true;
                }
            }
        }

        pair = default;
        return false;
    }

    public void Clear()
    {
        dictionary.Clear();
    }

    public Memory<LuaValue> GetArrayMemory()
    {
        return array.AsMemory();
    }

    public Span<LuaValue> GetArraySpan()
    {
        return array.AsSpan();
    }

    internal void EnsureArrayCapacity(int newCapacity)
    {
        if (array.Length >= newCapacity) return;

        var prevLength = array.Length;
        var newLength = array.Length;
        if (newLength == 0) newLength = 8;

        while (newLength < newCapacity)
        {
            newLength *= 2;
        }

        Array.Resize(ref array, newLength);

        using var indexList = new PooledList<(int, LuaValue)>(dictionary.Count);

        // Move some of the elements of the hash part to a newly allocated array
        foreach (var kv in dictionary)
        {
            if (TryGetInteger(kv.Key, out var index))
            {
                if (index > prevLength && index <= newLength)
                {
                    indexList.Add((index, kv.Value));
                }
            }
        }

        foreach ((var index, var value) in indexList.AsSpan())
        {
            dictionary.Remove(index);
            array[index - 1] = value;
        }
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    static bool TryGetInteger(LuaValue value, out int integer)
    {
        if (value.TryReadNumber(out var num) && MathEx.IsInteger(num))
        {
            integer = (int)num;
            return true;
        }

        integer = default;
        return false;
    }

    static void ThrowIndexIsNil()
    {
        throw new ArgumentException("the table index is nil");
    }

    static void ThrowIndexIsNaN()
    {
        throw new ArgumentException("the table index is NaN");
    }
}