Skip to content

Commit

Permalink
StreamOnSqlBytes Read/Write on Spans (dotnet#86674)
Browse files Browse the repository at this point in the history 
* Implement span overloads

* Add StreamOnBytes tests

* small tests update

* Reroute single byte read/write to span version

* Remove _rgbWorkBuf

* Add WriteInternal

* Internal -> NoValidation

* Swap arguments

* Swap methods

* One more method swap

---------

Co-authored-by: Stephen Toub <stoub@microsoft.com>
  • Loading branch information
@hrrrrustic @stephentoub
hrrrrustic and stephentoub committed Jan 29, 2024
1 parent 027a43b commit 8b082e1
Show file tree
Hide file tree
Showing 2 changed files with 241 additions and 73 deletions.
184 changes: 119 additions & 65 deletions src/libraries/System.Data.Common/src/System/Data/SQLTypes/SQLBytes.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,6 @@ public sealed class SqlBytes : INullable, IXmlSerializable, ISerializable
internal Stream? _stream;
private SqlBytesCharsState _state;

private byte[]? _rgbWorkBuf; // A 1-byte work buffer.

// The max data length that we support at this time.
private const long x_lMaxLen = int.MaxValue;

Expand Down Expand Up @@ -80,8 +78,6 @@ public SqlBytes(byte[]? buffer)
_lCurLen = _rgbBuf.Length;
}

_rgbWorkBuf = null;

AssertValid();
}

Expand All @@ -98,8 +94,6 @@ public SqlBytes(Stream? s)
_stream = s;
_state = (s == null) ? SqlBytesCharsState.Null : SqlBytesCharsState.Stream;

_rgbWorkBuf = null;

AssertValid();
}

Expand Down Expand Up @@ -200,17 +194,14 @@ public byte this[long offset]
ArgumentOutOfRangeException.ThrowIfNegative(offset);
ArgumentOutOfRangeException.ThrowIfGreaterThanOrEqual(offset, Length);

_rgbWorkBuf ??= new byte[1];
byte b = 0;

Read(offset, _rgbWorkBuf, 0, 1);
return _rgbWorkBuf[0];
Read(offset, new Span<byte>(ref b));
return b;
}
set
{
_rgbWorkBuf ??= new byte[1];

_rgbWorkBuf[0] = value;
Write(offset, _rgbWorkBuf, 0, 1);
Write(offset, new ReadOnlySpan<byte>(in value));
}
}

Expand Down Expand Up @@ -290,6 +281,17 @@ public void SetLength(long value)
AssertValid();
}

internal long Read(long offset, Span<byte> buffer)
{
if (IsNull)
throw new SqlNullValueException();

ArgumentOutOfRangeException.ThrowIfGreaterThan(offset, Length);
ArgumentOutOfRangeException.ThrowIfNegative(offset);

return ReadNoValidation(offset, buffer);
}

// Read data of specified length from specified offset into a buffer
public long Read(long offset, byte[] buffer, int offsetInBuffer, int count)
{
Expand All @@ -308,26 +310,50 @@ public long Read(long offset, byte[] buffer, int offsetInBuffer, int count)
ArgumentOutOfRangeException.ThrowIfNegative(count);
ArgumentOutOfRangeException.ThrowIfGreaterThan(count, buffer.Length - offsetInBuffer);

return ReadNoValidation(offset, buffer.AsSpan(offsetInBuffer, count));
}

private long ReadNoValidation(long offset, Span<byte> buffer)
{
if (_state == SqlBytesCharsState.Stream)
{
if (_stream!.Position != offset)
_stream.Seek(offset, SeekOrigin.Begin);
return _stream.Read(buffer);
}

// Adjust count based on data length
if (count > Length - offset)
count = (int)(Length - offset);
int count = Math.Min(buffer.Length, (int)(Length - offset));

Span<byte> span = _rgbBuf!.AsSpan((int)offset, count);
span.CopyTo(buffer);

if (count != 0)
return span.Length;
}

internal void Write(long offset, ReadOnlySpan<byte> buffer)
{
if (FStream())
{
switch (_state)
{
case SqlBytesCharsState.Stream:
if (_stream!.Position != offset)
_stream.Seek(offset, SeekOrigin.Begin);
count = _stream.Read(buffer, offsetInBuffer, count);
break;
if (_stream!.Position != offset)
_stream.Seek(offset, SeekOrigin.Begin);
_stream.Write(buffer);
}
else
{
if (_rgbBuf == null)
throw new SqlTypeException(SR.SqlMisc_NoBufferMessage);

default:
Array.Copy(_rgbBuf!, offset, buffer, offsetInBuffer, count);
break;
}
ArgumentOutOfRangeException.ThrowIfNegative(offset);

if (offset > _rgbBuf.Length)
throw new SqlTypeException(SR.SqlMisc_BufferInsufficientMessage);

if (buffer.Length > _rgbBuf.Length - offset)
throw new SqlTypeException(SR.SqlMisc_BufferInsufficientMessage);

WriteNoValidation(offset, buffer);
}
return count;
}

// Write data of specified length into the SqlBytes from specified offset
Expand Down Expand Up @@ -360,41 +386,47 @@ public void Write(long offset, byte[] buffer, int offsetInBuffer, int count)
if (count > _rgbBuf.Length - offset)
throw new SqlTypeException(SR.SqlMisc_BufferInsufficientMessage);

if (IsNull)
{
// If NULL and there is buffer inside, we only allow writing from
// offset zero.
//
if (offset != 0)
throw new SqlTypeException(SR.SqlMisc_WriteNonZeroOffsetOnNullMessage);

// treat as if our current length is zero.
// Note this has to be done after all inputs are validated, so that
// we won't throw exception after this point.
//
_lCurLen = 0;
_state = SqlBytesCharsState.Buffer;
}
else if (offset > _lCurLen)
{
// Don't allow writing from an offset that this larger than current length.
// It would leave uninitialized data in the buffer.
//
throw new SqlTypeException(SR.SqlMisc_WriteOffsetLargerThanLenMessage);
}
WriteNoValidation(offset, buffer.AsSpan(offsetInBuffer, count));
}

if (count != 0)
{
Array.Copy(buffer, offsetInBuffer, _rgbBuf, offset, count);
AssertValid();
}

// If the last position that has been written is after
// the current data length, reset the length
if (_lCurLen < offset + count)
_lCurLen = offset + count;
}
private void WriteNoValidation(long offset, ReadOnlySpan<byte> buffer)
{
if (IsNull)
{
// If NULL and there is buffer inside, we only allow writing from
// offset zero.
//
if (offset != 0)
throw new SqlTypeException(SR.SqlMisc_WriteNonZeroOffsetOnNullMessage);

// treat as if our current length is zero.
// Note this has to be done after all inputs are validated, so that
// we won't throw exception after this point.
//
_lCurLen = 0;
_state = SqlBytesCharsState.Buffer;
}
else if (offset > _lCurLen)
{
// Don't allow writing from an offset that this larger than current length.
// It would leave uninitialized data in the buffer.
//
throw new SqlTypeException(SR.SqlMisc_WriteOffsetLargerThanLenMessage);
}

AssertValid();
if (buffer.Length != 0)
{
Span<byte> span = _rgbBuf.AsSpan((int)offset, buffer.Length);
buffer.CopyTo(span);

// If the last position that has been written is after
// the current data length, reset the length
if (_lCurLen < offset + buffer.Length)
_lCurLen = offset + buffer.Length;
}
}

public SqlBinary ToSqlBinary()
Expand Down Expand Up @@ -436,7 +468,6 @@ private void AssertValid()
Debug.Assert(FStream() || (_rgbBuf != null && _lCurLen <= _rgbBuf.Length));
Debug.Assert(!FStream() || (_lCurLen == x_lNull));
}
Debug.Assert(_rgbWorkBuf == null || _rgbWorkBuf.Length == 1);
}

// Copy the data from the Stream to the array buffer.
Expand Down Expand Up @@ -686,17 +717,35 @@ public override long Seek(long offset, SeekOrigin origin)
return _lPosition;
}

public override int Read(Span<byte> buffer)
{
CheckIfStreamClosed();

return ReadNoValidation(buffer);
}
// The Read/Write/ReadByte/WriteByte simply delegates to SqlBytes
public override int Read(byte[] buffer, int offset, int count)
{
CheckIfStreamClosed();

ValidateBufferArguments(buffer, offset, count);

int iBytesRead = (int)_sb.Read(_lPosition, buffer, offset, count);
_lPosition += iBytesRead;
return ReadNoValidation(buffer.AsSpan(offset, count));
}

private int ReadNoValidation(Span<byte> buffer)
{
int bytesRead = (int)_sb.Read(_lPosition, buffer);
_lPosition += bytesRead;

return iBytesRead;
return bytesRead;
}

public override void Write(ReadOnlySpan<byte> buffer)
{
CheckIfStreamClosed();

WriteNoValidation(buffer);
}

public override void Write(byte[] buffer, int offset, int count)
Expand All @@ -705,8 +754,13 @@ public override void Write(byte[] buffer, int offset, int count)

ValidateBufferArguments(buffer, offset, count);

_sb.Write(_lPosition, buffer, offset, count);
_lPosition += count;
WriteNoValidation(buffer);
}

private void WriteNoValidation(ReadOnlySpan<byte> buffer)
{
_sb.Write(_lPosition, buffer);
_lPosition += buffer.Length;
}

public override int ReadByte()
Expand Down
Loading

0 comments on commit 8b082e1

Please sign in to comment.