path: root/LibGit/Extensions/StreamExtensions.cs

using System.IO.Compression;
using System.Runtime.InteropServices;

namespace LibGit.Extensions;

public static class StreamExtensions
{
    private const bool _debug = false;

    public static long Remaining(this Stream stream)
    {
        //if (_debug) Console.WriteLine($"stream pos: {stream.Position}, stream len: {stream.Length}, stream rem: {stream.Length - stream.Position}");
        return stream.Length - stream.Position;
    }

    public static int Peek(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't peek a non-readable stream");
        if (!stream.CanSeek)
            throw new InvalidOperationException("Can't peek a non-seekable stream");

        int peek = stream.ReadByte();
        if (peek != -1)
            stream.Seek(-1, SeekOrigin.Current);

        return peek;
    }

    public static IEnumerable<byte> Peek(this Stream stream, long count)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't peek a non-readable stream");
        if (!stream.CanSeek)
            throw new InvalidOperationException("Can't peek a non-seekable stream");
        long i;
        for (i = 0; i < count; i++)
        {
            int peek = stream.ReadByte();
            if (peek == -1)
            {
                if (_debug) Console.WriteLine($"Can't peek {count} bytes, only {i} bytes remaining");
                stream.Seek(-i, SeekOrigin.Current);
                yield break;
            }

            yield return (byte)peek;
        }

        stream.Seek(-i, SeekOrigin.Current);
    }

    public static IEnumerable<byte> ReadBytes(this Stream stream, long count)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        var buffer = new byte[count];
        var readCount = stream.Read(buffer, 0, (int)count);
        for (int i = 0; i < readCount; i++)
        {
            yield return buffer[i];
        }
    }

    public static byte[] ReadBlock(this Stream stream, int count)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<byte> buffer = stackalloc byte[count];
        stream.ReadExactly(buffer);
        return buffer.ToArray();
    }

    public static bool StartsWith(this Stream stream, IEnumerable<byte> sequence)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");
        if (!stream.CanSeek)
            throw new InvalidOperationException("Can't read a non-seekable stream");

        if (_debug)
        {
            Console.WriteLine($"Expected: {sequence.AsHexString()} ({sequence.AsString()})");
            Console.WriteLine($"Actual:   {stream.Peek(sequence.Count()).AsHexString()} ({stream.Peek(sequence.Count()).AsString()})");
        }

        int readCount = 0;
        foreach (int b in sequence)
        {
            int read = stream.ReadByte();
            readCount++;
            if (read == -1)
            {
                stream.Seek(-readCount, SeekOrigin.Current);
                return false;
            }

            if (read != b)
            {
                if (_debug)
                {
                    Console.ForegroundColor = ConsoleColor.Red;
                    Console.WriteLine("^^".PadLeft(readCount * 3 + 9));
                    Console.ResetColor();
                }

                stream.Seek(-readCount, SeekOrigin.Current);
                return false;
            }
        }

        stream.Seek(-readCount, SeekOrigin.Current);

        return true;
    }

    public static bool StartsWith(this Stream stream, string ascii_seq)
    {
        return stream.StartsWith(ascii_seq.Select(x => (byte)x));
    }

    public static Stream Skip(this Stream stream, long count = 1)
    {
        if (!stream.CanSeek)
            throw new InvalidOperationException("Can't skip a non-seekable stream");
        stream.Seek(count, SeekOrigin.Current);
        return stream;
    }

    public static IEnumerable<byte> ReadNullTerminatedField(this Stream stream, IEnumerable<byte>? binaryPrefix = null, string? asciiPrefix = null) =>
        ReadTerminatedField(stream: stream, terminator: 0x00, binaryPrefix: binaryPrefix, asciiPrefix: asciiPrefix);

    public static IEnumerable<byte> ReadSpaceTerminatedField(this Stream stream, IEnumerable<byte>? binaryPrefix = null, string? asciiPrefix = null) =>
        ReadTerminatedField(stream: stream, terminator: 0x20, binaryPrefix: binaryPrefix, asciiPrefix: asciiPrefix);

    public static IEnumerable<byte> ReadTerminatedField(this Stream stream, byte terminator, IEnumerable<byte>? binaryPrefix = null, string? asciiPrefix = null)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");
        if (binaryPrefix != null)
            if (!stream.StartsWith(binaryPrefix))
                throw new InvalidDataException($"Binary prefix {stream.Peek(binaryPrefix.Count()).AsHexString()} does not match expected value of {binaryPrefix.AsHexString()}!");
            else stream.Skip(binaryPrefix.Count());
        else if (asciiPrefix != null)
            if (!stream.StartsWith(asciiPrefix))
                throw new InvalidDataException(
                    $"Text prefix {stream.Peek(asciiPrefix.Length).AsHexString()} ({stream.Peek(asciiPrefix.Length).AsString()}) does not match expected value of {asciiPrefix.AsBytes().AsHexString()} ({asciiPrefix})!");
            else stream.Skip(asciiPrefix.Length);

        var read = 0;
        while (stream.Peek() != terminator)
        {
            if (_debug) Console.WriteLine($"ReadTerminatedField -- pos: {stream.Position}/+{stream.Remaining()}/{stream.Length} | next: {(char)stream.Peek()} | Length: {read}");
            if (stream.Peek() == -1)
            {
                Console.WriteLine($"Warning: Reached end of stream while reading null-terminated field");
                yield break;
            }

            read++;
            yield return (byte)stream.ReadByte();
        }

        if (stream.Peek() == terminator) stream.Skip();
    }

    public static IEnumerable<byte> ReadToEnd(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        while (stream.Peek() != -1)
            yield return (byte)stream.ReadByte();
    }

    public static int ReadInt32BE(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<byte> b = stackalloc byte[4];
        stream.ReadExactly(b);
        return ((b[0] << 24) | (b[1] << 16) | (b[2] << 8) | b[3]);
    }

    public static int[] MultiReadInt32BE(this Stream stream, int count)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<byte> b = stackalloc byte[4 * count];
        stream.ReadExactly(b);
        int[] results = new int[count];
        for (int i = 0; i < count; i++)
        {
            results[i] = (b[i * 4 + 0] << 24) | (b[i * 4 + 1] << 16) | (b[i * 4 + 2] << 8) | b[i * 4 + 3];
        }

        return results;
    }

    public static uint ReadUInt32BE(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<byte> b = stackalloc byte[4];
        stream.ReadExactly(b);
        return (uint)((b[0] << 24) | (b[1] << 16) | (b[2] << 8) | b[3]);
    }

    public static uint[] MultiReadUInt32BE(this Stream stream, int count)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<byte> b = count >= 32767 ? new byte[4 * count] : stackalloc byte[4 * count];
        stream.ReadExactly(b);
        uint[] results = new uint[count];
        for (int i = 0; i < count; i++)
        {
            results[i] = (uint)((b[i * 4 + 0] << 24) | (b[i * 4 + 1] << 16) | (b[i * 4 + 2] << 8) | b[i * 4 + 3]);
        }

        return results;
    }

    public static ulong ReadUInt64BE(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<byte> b = stackalloc byte[8];
        stream.ReadExactly(b);
        return ((ulong)b[0] << 56) | ((ulong)b[1] << 48) | ((ulong)b[2] << 40) | ((ulong)b[3] << 32) |
               ((ulong)b[4] << 24) | ((ulong)b[5] << 16) | ((ulong)b[6] << 8) | b[7];
    }

    public static ulong[] MultiReadUInt64BE(this Stream stream, int count)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<byte> b = count >= 16384 ? new byte[8 * count] : stackalloc byte[8 * count];
        stream.ReadExactly(b);
        ulong[] results = new ulong[count];
        for (int i = 0; i < count; i++)
        {
            results[i] = ((ulong)b[i * 8 + 0] << 56) | ((ulong)b[i * 8 + 1] << 48) | ((ulong)b[i * 8 + 2] << 40) | ((ulong)b[i * 8 + 3] << 32) |
                         ((ulong)b[i * 8 + 4] << 24) | ((ulong)b[i * 8 + 5] << 16) | ((ulong)b[i * 8 + 6] << 8) | b[i * 8 + 7];
        }

        return results;
    }

    public static T[] MultiReadInlineArray<T>(this Stream stream, int count) where T : unmanaged
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        Span<T> buffer = count >= 8192 ? new T[count] : stackalloc T[count];
        Span<byte> byteBuffer = MemoryMarshal.AsBytes(buffer);
        stream.ReadExactly(byteBuffer);
        return buffer.ToArray();
    }

    //read variable length number
    public static int ReadVLQ(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        int result = 0;
        int shift = 0;
        byte b;
        do
        {
            b = (byte)stream.ReadByte();
            result |= (b & 0b0111_1111) << shift;
            shift += 7;
        } while ((b & 0b1000_0000) != 0);

        return result;
    }

    public static int ReadVLQBigEndian(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        int result = 0;
        int shift = 0;
        byte b;
        do
        {
            b = (byte)stream.ReadByte();
            result = (result << 7) | (b & 0b0111_1111);
        } while ((b & 0b1000_0000) != 0);

        return result;
    }

    // for some reason this is special...
    public static int ReadGitPackOffsetModifiedVLQ(this Stream stream)
    {
        if (!stream.CanRead)
            throw new InvalidOperationException("Can't read a non-readable stream");

        int b = stream.ReadByte();
        int result = b & 0x7F;

        while ((b & 0x80) != 0)
        {
            b = stream.ReadByte();
            result = ((result + 1) << 7) | (b & 0x7F);
        }

        return result;
    }
}