using System;

using System.Buffers;

using System.Diagnostics;

using System.Diagnostics.CodeAnalysis;

using System.IO;

using System.Runtime.CompilerServices;

using System.Text;

using System.Threading;

using System.Threading.Tasks;

namespace Npgsql.Internal;

public class PgReader

readonly NpgsqlReadBuffer _buffer;

bool _resumable;

byte[]? _pooledArray;

NpgsqlReadBuffer.ColumnStream? _userActiveStream;

PreparedTextReader? _preparedTextReader;

long _fieldStartPos;

Size _fieldBufferRequirement;

DataFormat _fieldFormat;

int _fieldSize;

// This position is relative to _fieldStartPos, which is why it can be an int.

int _currentStartPos;

Size _currentBufferRequirement;

int _currentSize;

// GetChars Internal state

TextReader? _charsReadReader;

int _charsRead;

// GetChars User state

int? _charsReadOffset;

ArraySegment<char>? _charsReadBuffer;

bool _requiresCleanup;

// The field reading process of doing init/commit and startread/endread pairs is very perf sensitive.

// So this is used in Commit as a fast-path alternative to FieldRemaining to detect if the field was consumed succesfully.

bool _fieldConsumed;

internal PgReader(NpgsqlReadBuffer buffer)

}

internal long FieldStartPos => _fieldStartPos;

internal int FieldSize => _fieldSize;

internal bool Initialized => _fieldStartPos is not -1;

internal int FieldOffset => (int)(_buffer.CumulativeReadPosition - _fieldStartPos);

internal int FieldRemaining => FieldSize - FieldOffset;

bool HasCurrent => _currentSize is not -1;

int CurrentSize => HasCurrent ? _currentSize : _fieldSize;

public ValueMetadata Current => new() { Size = CurrentSize, Format = _fieldFormat, BufferRequirement = CurrentBufferRequirement };

public int CurrentRemaining => HasCurrent ? _currentSize - CurrentOffset : FieldRemaining;

Size CurrentBufferRequirement => HasCurrent ? _currentBufferRequirement : _fieldBufferRequirement;

int CurrentOffset => FieldOffset - _currentStartPos;

internal bool IsAtStart => FieldOffset is 0;

internal bool Resumable => _resumable;

public bool IsResumed => Resumable && CurrentSize != CurrentRemaining;

ArrayPool<byte> ArrayPool => ArrayPool<byte>.Shared;

[MemberNotNullWhen(true, nameof(_charsReadReader))]

internal bool IsCharsRead => _charsReadOffset is not null;

// Here for testing purposes

internal void BreakConnection() => throw _buffer.Connector.Break(new Exception("Broken"));

internal void Revert(int size, int startPos, Size bufferRequirement)

}

[Conditional("DEBUG")]

void CheckBounds(int count)

}

public byte ReadByte()

}

public short ReadInt16()

}

public int ReadInt32()

}

public long ReadInt64()

}

public ushort ReadUInt16()

}

public uint ReadUInt32()

}

public ulong ReadUInt64()

}

public float ReadFloat()

}

public double ReadDouble()

}

public void Read(Span<byte> destination)

}

public async ValueTask<string> ReadNullTerminatedStringAsync(Encoding encoding, CancellationToken cancellationToken = default)

}

public string ReadNullTerminatedString(Encoding encoding)

}

public Stream GetStream(int? length = null) => GetColumnStream(false, length);

internal Stream GetStream(bool canSeek, int? length = null) => GetColumnStream(canSeek, length);

NpgsqlReadBuffer.ColumnStream GetColumnStream(bool canSeek = false, int? length = null)

}

public TextReader GetTextReader(Encoding encoding)

=> GetTextReader(async: false, encoding, CancellationToken.None).GetAwaiter().GetResult();

public ValueTask<TextReader> GetTextReaderAsync(Encoding encoding, CancellationToken cancellationToken)

=> GetTextReader(async: true, encoding, cancellationToken);

async ValueTask<TextReader> GetTextReader(bool async, Encoding encoding, CancellationToken cancellationToken)

}

public ValueTask ReadBytesAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)

}

public void ReadBytes(Span<byte> buffer)

}

public bool TryReadBytes(int count, out ReadOnlySpan<byte> bytes)

}

public bool TryReadBytes(int count, out ReadOnlyMemory<byte> bytes)

}

/// ReadBytes without memory management, the next read invalidates the underlying buffer(s), only use this for intermediate transformations.

public ReadOnlySequence<byte> ReadBytes(int count)

}

/// ReadBytesAsync without memory management, the next read invalidates the underlying buffer(s), only use this for intermediate transformations.

public async ValueTask<ReadOnlySequence<byte>> ReadBytesAsync(int count, CancellationToken cancellationToken = default)

}

public void Rewind(int count)

}

/// <summary>

///

/// </summary>

/// <param name="async"></param>

/// <returns>The stream length, if any</returns>

async ValueTask DisposeUserActiveStream(bool async)

{

if (_userActiveStream is { IsDisposed: false })

{

if (async)

await _userActiveStream.DisposeAsync().ConfigureAwait(false);

else

_userActiveStream.Dispose();

}

_userActiveStream = null;

}

internal bool GetCharsReadInfo(Encoding encoding, out int charsRead, out TextReader reader, out int charsOffset, out ArraySegment<char>? buffer)

{

if (!IsCharsRead)

throw new InvalidOperationException("No active chars read");

if (_charsReadReader is null)

{

charsRead = 0;

reader = _charsReadReader = GetTextReader(encoding);

charsOffset = _charsReadOffset ??= 0;

buffer = _charsReadBuffer;

return true;

}

charsRead = _charsRead;

reader = _charsReadReader;

charsOffset = _charsReadOffset!.Value;

buffer = _charsReadBuffer;

return false;

}

internal void ResetCharsRead(out int charsRead)

{

if (!IsCharsRead)

throw new InvalidOperationException("No active chars read");

switch (_charsReadReader)

{

case PreparedTextReader reader:

reader.Restart();

break;

case StreamReader reader:

reader.BaseStream.Seek(0, SeekOrigin.Begin);

reader.DiscardBufferedData();

break;

}

_charsRead = charsRead = 0;

}

internal void AdvanceCharsRead(int charsRead)

{

_charsRead += charsRead;

_charsReadOffset = null;

_charsReadBuffer = null;

}

internal void InitCharsRead(int dataOffset, ArraySegment<char>? buffer, out int? charsRead)

{

if (!Resumable)

throw new InvalidOperationException("Wasn't initialized as resumed");

charsRead = _charsReadReader is null ? null : _charsRead;

_charsReadOffset = dataOffset;

_charsReadBuffer = buffer;

}

internal PgReader Init(int fieldLength, DataFormat format, bool resumable = false)

{

if (Initialized)

{

if (resumable)

{

if (Resumable)

return this;

_resumable = true;

}

else

{

if (!IsAtStart)

ThrowHelper.ThrowInvalidOperationException("Cannot be initialized to be non-resumable until a commit is issued.");

_resumable = false;

}

}

Debug.Assert(!_requiresCleanup, "Reader wasn't properly committed before next init");

_fieldStartPos = _buffer.CumulativeReadPosition;

_fieldFormat = format;

_fieldSize = fieldLength;

_resumable = resumable;

_fieldConsumed = false;

return this;

}

internal void StartRead(Size bufferRequirement)

{

Debug.Assert(FieldSize >= 0);

_fieldBufferRequirement = bufferRequirement;

if (ShouldBuffer(bufferRequirement))

Buffer(bufferRequirement);

}

internal ValueTask StartReadAsync(Size bufferRequirement, CancellationToken cancellationToken)

{

Debug.Assert(FieldSize >= 0);

_fieldBufferRequirement = bufferRequirement;

return ShouldBuffer(bufferRequirement) ? BufferAsync(bufferRequirement, cancellationToken) : new();

}

internal void EndRead()

{

if (_resumable)

return;

// If it was upper bound we should consume.

if (_fieldBufferRequirement is { Kind: SizeKind.UpperBound })

{

Consume(FieldRemaining);

return;

}

if (FieldOffset != FieldSize)

ThrowNotConsumedExactly();

_fieldConsumed = true;

}

internal ValueTask EndReadAsync()

{

if (_resumable)

return new();

// If it was upper bound we should consume.

if (_fieldBufferRequirement is { Kind: SizeKind.UpperBound })

return ConsumeAsync(FieldRemaining);

if (FieldOffset != FieldSize)

ThrowNotConsumedExactly();

_fieldConsumed = true;

return new();

}

internal async ValueTask<NestedReadScope> BeginNestedRead(bool async, int size, Size bufferRequirement, CancellationToken cancellationToken = default)

{

if (size > CurrentRemaining)

throw new ArgumentOutOfRangeException(nameof(size), "Cannot begin a read for a larger size than the current remaining size.");

if (size < 0)

throw new ArgumentOutOfRangeException(nameof(size), "Cannot be negative");

var previousSize = CurrentSize;

var previousStartPos = _currentStartPos;

var previousBufferRequirement = CurrentBufferRequirement;

_currentSize = size;

_currentBufferRequirement = bufferRequirement;

_currentStartPos = FieldOffset;

await Buffer(async, bufferRequirement, cancellationToken).ConfigureAwait(false);

return new NestedReadScope(async, this, previousSize, previousStartPos, previousBufferRequirement);

}

public NestedReadScope BeginNestedRead(int size, Size bufferRequirement)

=> BeginNestedRead(async: false, size, bufferRequirement, CancellationToken.None).GetAwaiter().GetResult();

public ValueTask<NestedReadScope> BeginNestedReadAsync(int size, Size bufferRequirement, CancellationToken cancellationToken = default)

=> BeginNestedRead(async: true, size, bufferRequirement, cancellationToken);

internal void Seek(int offset)

{

if (CurrentOffset > offset)

Rewind(CurrentOffset - offset);

else if (CurrentOffset < offset)

Consume(offset - CurrentOffset);

}

internal async ValueTask Consume(bool async, int? count = null, CancellationToken cancellationToken = default)

{

if (count <= 0 || FieldSize < 0 || FieldRemaining == 0)

return;

var remaining = count ?? CurrentRemaining;

CheckBounds(remaining);

var origOffset = FieldOffset;

// A breaking exception unwind from a nested scope should not try to consume its remaining data.

if (!_buffer.Connector.IsBroken)

await _buffer.Skip(remaining, async).ConfigureAwait(false);

Debug.Assert(FieldRemaining == FieldSize - origOffset - remaining);

}

public void Consume(int? count = null) => Consume(async: false, count).GetAwaiter().GetResult();

public ValueTask ConsumeAsync(int? count = null, CancellationToken cancellationToken = default) => Consume(async: true, count, cancellationToken);

internal void ThrowIfStreamActive()

{

if (_userActiveStream is { IsDisposed: false})

ThrowHelper.ThrowInvalidOperationException("A stream is already open for this reader");

}

internal bool CommitHasIO(bool resuming) => Initialized && !resuming && FieldRemaining > 0;

[MethodImpl(MethodImplOptions.AggressiveInlining)]

internal void Commit(bool resuming)

{

if (!Initialized)

return;

if (resuming)

{

if (!Resumable)

ThrowHelper.ThrowInvalidOperationException("Cannot resume a non-resumable read.");

return;

}

// We don't rely on CurrentRemaining, just to make sure we consume fully in the event of a nested scope not being disposed.

// Also shut down any streaming, pooled arrays etc.

if (_requiresCleanup || (!_fieldConsumed && FieldRemaining > 0))

{

CommitSlow();

return;

}

_fieldStartPos = -1;

Debug.Assert(!Initialized);

// These will always be re-initialized by Init()

// _fieldSize = default;

// _fieldFormat = default;

// _resumable = default;

// _fieldCompleted = default;

if (HasCurrent)

{

_currentStartPos = 0;

_currentBufferRequirement = default;

_currentSize = -1;

Debug.Assert(!HasCurrent);

}

[MethodImpl(MethodImplOptions.NoInlining)]

void CommitSlow()

{

// Shut down any streaming and pooling going on on the column.

if (_requiresCleanup)

{

if (_userActiveStream is { IsDisposed: false })

DisposeUserActiveStream(async: false).GetAwaiter().GetResult();

if (_pooledArray is not null)

{

ArrayPool.Return(_pooledArray);

_pooledArray = null;

}

if (_charsReadReader is not null)

{

_charsReadReader.Dispose();

_charsReadReader = null;

_charsRead = default;

}

_requiresCleanup = false;

}

Consume(async: false, count: FieldRemaining).GetAwaiter().GetResult();

_fieldStartPos = -1;

Debug.Assert(!Initialized);

// These will always be re-initialized by Init()

// _fieldSize = default;

// _fieldFormat = default;

// _resumable = default;

// _fieldCompleted = default;

if (HasCurrent)

{

_currentStartPos = 0;

_currentBufferRequirement = default;

_currentSize = -1;

Debug.Assert(!HasCurrent);

}

}

}

[MethodImpl(MethodImplOptions.AggressiveInlining)]

internal ValueTask CommitAsync(bool resuming)

{

if (!Initialized)

return new();

if (resuming)

{

if (!Resumable)

ThrowHelper.ThrowInvalidOperationException("Cannot resume a non-resumable read.");

return new();

}

// We don't rely on CurrentRemaining, just to make sure we consume fully in the event of a nested scope not being disposed.

// Also shut down any streaming, pooled arrays etc.

if (_requiresCleanup || (!_fieldConsumed && FieldRemaining > 0))

return CommitSlow();

_fieldStartPos = -1;

Debug.Assert(!Initialized);

// These will always be re-initialized by Init()

// _fieldSize = default;

// _fieldFormat = default;

// _resumable = default;

// _fieldCompleted = default;

if (HasCurrent)

{

_currentStartPos = 0;

_currentBufferRequirement = default;

_currentSize = -1;

Debug.Assert(!HasCurrent);

}

return new();

async ValueTask CommitSlow()

{

// Shut down any streaming and pooling going on on the column.

if (_requiresCleanup)

{

if (_userActiveStream is { IsDisposed: false })

await DisposeUserActiveStream(async: true).ConfigureAwait(false);

if (_pooledArray is not null)

{

ArrayPool.Return(_pooledArray);

_pooledArray = null;

}

if (_charsReadReader is not null)

{

_charsReadReader.Dispose();

_charsReadReader = null;

_charsRead = default;

}

_requiresCleanup = false;

}

await Consume(async: true, count: FieldRemaining).ConfigureAwait(false);

_fieldStartPos = -1;

Debug.Assert(!Initialized);

// These will always be re-initialized by Init()

// _fieldSize = default;

// _fieldFormat = default;

// _resumable = default;

// _fieldCompleted = default;

if (HasCurrent)

{

_currentStartPos = 0;

_currentBufferRequirement = default;

_currentSize = -1;

Debug.Assert(!HasCurrent);

}

}

}

byte[] RentArray(int count)

{

_requiresCleanup = true;

var pooledArray = _pooledArray;

if (pooledArray is not null)

{

if (pooledArray.Length >= count)

return pooledArray;

ArrayPool.Return(pooledArray);

}

var array = _pooledArray = ArrayPool.Rent(count);

return array;

}

[MethodImpl(MethodImplOptions.AggressiveInlining)]

int GetBufferRequirementByteCount(Size bufferRequirement)

=> bufferRequirement is { Kind: SizeKind.UpperBound }

? Math.Min(CurrentRemaining, bufferRequirement.Value)

: bufferRequirement.GetValueOrDefault();

internal bool ShouldBufferCurrent() => ShouldBuffer(CurrentBufferRequirement);

public bool ShouldBuffer(Size bufferRequirement)

=> ShouldBuffer(GetBufferRequirementByteCount(bufferRequirement));

public bool ShouldBuffer(int byteCount)

{

return _buffer.ReadBytesLeft < byteCount && ShouldBufferSlow();

[MethodImpl(MethodImplOptions.NoInlining)]

bool ShouldBufferSlow()

{

if (byteCount > _buffer.Size)

ThrowArgumentOutOfRange();

if (byteCount > CurrentRemaining)

ThrowArgumentOutOfRangeOfValue();

return true;

}

static void ThrowArgumentOutOfRange()

=> throw new ArgumentOutOfRangeException(nameof(byteCount),

"Buffer requirement is larger than the buffer size, this can never succeed by buffering data but requires a larger buffer size instead.");

static void ThrowArgumentOutOfRangeOfValue()

=> throw new ArgumentOutOfRangeException(nameof(byteCount),

"Buffer requirement is larger than the remaining length of the value, make sure the value is always at least this size or use an upper bound requirement instead.");

}

public void Buffer(Size bufferRequirement)

=> Buffer(GetBufferRequirementByteCount(bufferRequirement));

public void Buffer(int byteCount) => _buffer.Ensure(byteCount, async: false).GetAwaiter().GetResult();

public ValueTask BufferAsync(Size bufferRequirement, CancellationToken cancellationToken)

=> BufferAsync(GetBufferRequirementByteCount(bufferRequirement), cancellationToken);

public ValueTask BufferAsync(int byteCount, CancellationToken cancellationToken) => _buffer.EnsureAsync(byteCount);

internal ValueTask Buffer(bool async, Size bufferRequirement, CancellationToken cancellationToken)

=> Buffer(async, GetBufferRequirementByteCount(bufferRequirement), cancellationToken);

internal ValueTask Buffer(bool async, int byteCount, CancellationToken cancellationToken)

{

if (async)

return BufferAsync(byteCount, cancellationToken);

Buffer(byteCount);

return new();

}

void ThrowNotConsumedExactly() =>

throw _buffer.Connector.Break(

new InvalidOperationException(

FieldOffset < FieldSize

? $"The read on this field has not consumed all of its bytes (pos: {FieldOffset}, len: {FieldSize})"

: $"The read on this field has consumed all of its bytes and read into the subsequent bytes (pos: {FieldOffset}, len: {FieldSize})"));

public readonly struct NestedReadScope : IDisposable, IAsyncDisposable

readonly PgReader _reader;

readonly int _previousSize;

readonly int _previousStartPos;

readonly Size _previousBufferRequirement;

readonly bool _async;

internal NestedReadScope(bool async, PgReader reader, int previousSize, int previousStartPos, Size previousBufferRequirement)

{

_async = async;

_reader = reader;

_previousSize = previousSize;

_previousStartPos = previousStartPos;

_previousBufferRequirement = previousBufferRequirement;

}

public void Dispose()

{

if (_async)

throw new InvalidOperationException("Cannot synchronously dispose async scopes, call DisposeAsync instead.");

DisposeAsync().GetAwaiter().GetResult();

}

public ValueTask DisposeAsync()

{

if (_reader.CurrentRemaining > 0)

{

if (_async)

return AsyncCore(_reader, _previousSize, _previousStartPos, _previousBufferRequirement);

_reader.Consume();

}

_reader.Revert(_previousSize, _previousStartPos, _previousBufferRequirement);

return new();

static async ValueTask AsyncCore(PgReader reader, int previousSize, int previousStartPos, Size previousBufferRequirement)

{

await reader.ConsumeAsync().ConfigureAwait(false);

reader.Revert(previousSize, previousStartPos, previousBufferRequirement);

}

}