|
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
namespace Microsoft.CodeAnalysis;
/// <summary>
/// Helpers to create temporary streams backed by pooled memory
/// </summary>
internal static class SerializableBytes
{
private const int ChunkSize = SharedPools.ByteBufferSize;
internal static PooledStream CreateReadableStream(byte[] bytes)
=> CreateReadableStream(bytes, bytes.Length);
internal static PooledStream CreateReadableStream(byte[] bytes, int length)
{
var stream = CreateWritableStream();
stream.Write(bytes, 0, length);
stream.Position = 0;
return stream;
}
internal static async Task<PooledStream> CreateReadableStreamAsync(Stream stream, CancellationToken cancellationToken)
{
var length = stream.Length;
var chunkCount = (length + ChunkSize - 1) / ChunkSize;
var chunks = new byte[chunkCount][];
try
{
for (long i = 0, c = 0; i < length; i += ChunkSize, c++)
{
var count = (int)Math.Min(ChunkSize, length - i);
var chunk = SharedPools.ByteArray.Allocate();
var chunkOffset = 0;
while (count > 0)
{
var bytesRead = await stream.ReadAsync(chunk, chunkOffset, count, cancellationToken).ConfigureAwait(false);
if (bytesRead > 0)
{
count -= bytesRead;
chunkOffset += bytesRead;
}
else
{
break;
}
}
chunks[c] = chunk;
}
var result = new ReadStream(length, chunks);
chunks = null;
return result;
}
finally
{
BlowChunks(chunks);
}
}
// free any chunks remaining
private static void BlowChunks(byte[][]? chunks)
{
if (chunks != null)
{
for (long c = 0; c < chunks.Length; c++)
{
if (chunks[c] != null)
{
SharedPools.ByteArray.Free(chunks[c]);
chunks[c] = null!;
}
}
}
}
internal static ReadWriteStream CreateWritableStream()
=> new();
public abstract class PooledStream : Stream
{
protected List<byte[]> chunks;
protected long position;
protected long length;
protected PooledStream(long length, List<byte[]> chunks)
{
this.position = 0;
this.length = length;
this.chunks = chunks;
}
public override long Length => this.length;
public override bool CanRead => true;
public override bool CanSeek => true;
public override bool CanWrite => false;
public override void Flush()
{
// nothing to do, this is a read-only stream
}
public override long Position
{
get => this.position;
set
{
if (value < 0 || value >= length)
{
throw new ArgumentOutOfRangeException(nameof(value));
}
this.position = value;
}
}
public override long Seek(long offset, SeekOrigin origin)
{
long target;
try
{
target = origin switch
{
SeekOrigin.Begin => offset,
SeekOrigin.Current => checked(offset + position),
SeekOrigin.End => checked(offset + length),
_ => throw new ArgumentOutOfRangeException(nameof(origin)),
};
}
catch (OverflowException)
{
throw new ArgumentOutOfRangeException(nameof(offset));
}
if (target < 0)
{
throw new ArgumentOutOfRangeException(nameof(offset));
}
position = target;
return target;
}
public override int ReadByte()
{
if (position >= length)
return -1;
var result = chunks[CurrentChunkIndex][CurrentChunkOffset];
this.position++;
return result;
}
public override int Read(byte[] buffer, int index, int count)
{
if (count <= 0 || position >= length)
{
return 0;
}
var totalCopyCount = Read(this.chunks, this.position, this.length, buffer, index, count);
this.position += totalCopyCount;
return totalCopyCount;
}
private static int Read(List<byte[]> chunks, long position, long length, byte[] buffer, int index, int count)
{
var oldPosition = position;
while (count > 0 && position < length)
{
var chunk = chunks[GetChunkIndex(position)];
var currentOffset = GetChunkOffset(position);
var copyCount = Math.Min(Math.Min(ChunkSize - currentOffset, count), (int)(length - position));
Array.Copy(chunk, currentOffset, buffer, index, copyCount);
position += copyCount;
index += copyCount;
count -= copyCount;
}
return (int)(position - oldPosition);
}
public byte[] ToArray()
{
if (this.Length == 0)
{
return [];
}
var array = new byte[this.Length];
// read entire array
Read(this.chunks, 0, this.length, array, 0, array.Length);
return array;
}
public ImmutableArray<byte> ToImmutableArray()
{
// ImmutableArray only supports int-sized arrays
var count = checked((int)Length);
var builder = ImmutableArray.CreateBuilder<byte>(count);
var chunkIndex = 0;
while (count > 0)
{
var chunk = chunks[chunkIndex];
var copyCount = Math.Min(chunk.Length, count);
builder.AddRange(chunk, copyCount);
count -= copyCount;
chunkIndex++;
}
Debug.Assert(count == 0);
return builder.MoveToImmutable();
}
protected int CurrentChunkIndex { get { return GetChunkIndex(this.position); } }
protected int CurrentChunkOffset { get { return GetChunkOffset(this.position); } }
protected static int GetChunkIndex(long value)
=> (int)(value / ChunkSize);
protected static int GetChunkOffset(long value)
=> (int)(value % ChunkSize);
protected override void Dispose(bool disposing)
{
base.Dispose(disposing);
if (chunks != null)
{
foreach (var chunk in chunks)
{
SharedPools.ByteArray.Free(chunk);
}
chunks = null!;
}
}
public override void SetLength(long value)
=> throw new NotSupportedException();
public override void Write(byte[] buffer, int offset, int count)
=> throw new NotSupportedException();
}
private sealed class ReadStream(long length, byte[][] chunks) : PooledStream(length, new List<byte[]>(chunks))
{
}
public sealed class ReadWriteStream : PooledStream
{
public ReadWriteStream()
: base(length: 0, chunks: SharedPools.BigDefault<List<byte[]>>().AllocateAndClear())
{
// growing list on EnsureSize shown as perf bottleneck. reuse shared list so that
// we don't re-allocate as much.
}
public override bool CanWrite => true;
public override long Position
{
get
{
return base.Position;
}
set
{
if (value < 0)
{
throw new ArgumentOutOfRangeException(nameof(value));
}
this.position = value;
}
}
private void EnsureCapacity(long value)
{
var nextIndex = GetChunkIndex(value);
for (var i = this.chunks.Count; i <= nextIndex; i++)
{
// allocate memory and initialize it to zero
var chunk = SharedPools.ByteArray.Allocate();
Array.Clear(chunk, 0, chunk.Length);
chunks.Add(chunk);
}
}
public override void SetLength(long value)
=> SetLength(value, truncate: true);
/// <summary>
/// Sets the length of this stream (see <see cref="SetLength(long)"/>. If <paramref name="truncate"/> is <see
/// langword="false"/>, the internal buffers will be left as is, and the data in them will be left as garbage.
/// If it is <see langword="true"/> then any fully unused chunks will be discarded. If there is a final chunk
/// the stream is partway through, the remainder of that chunk will be zeroed out.
/// </summary>
public void SetLength(long value, bool truncate)
{
EnsureCapacity(value);
if (value < length && truncate)
{
var chunkIndex = GetChunkIndex(value);
var chunkOffset = GetChunkOffset(value);
Array.Clear(chunks[chunkIndex], chunkOffset, chunks[chunkIndex].Length - chunkOffset);
var trimIndex = chunkIndex + 1;
for (var i = trimIndex; i < chunks.Count; i++)
SharedPools.ByteArray.Free(chunks[i]);
chunks.RemoveRange(trimIndex, chunks.Count - trimIndex);
}
length = value;
if (position > value)
{
position = value;
}
}
public override void WriteByte(byte value)
{
EnsureCapacity(this.position + 1);
var currentIndex = CurrentChunkIndex;
var currentOffset = CurrentChunkOffset;
chunks[currentIndex][currentOffset] = value;
this.position++;
if (this.position >= length)
{
this.length = this.position;
}
}
public override void Write(byte[] buffer, int index, int count)
{
EnsureCapacity(this.position + count);
var currentIndex = index;
var countLeft = count;
while (countLeft > 0)
{
var chunk = chunks[CurrentChunkIndex];
var currentOffset = CurrentChunkOffset;
var writeCount = Math.Min(ChunkSize - currentOffset, countLeft);
Array.Copy(buffer, currentIndex, chunk, currentOffset, writeCount);
this.position += writeCount;
currentIndex += writeCount;
countLeft -= writeCount;
}
if (this.position >= length)
{
this.length = this.position;
}
}
protected override void Dispose(bool disposing)
{
var temp = this.chunks;
base.Dispose(disposing);
SharedPools.BigDefault<List<byte[]>>().ClearAndFree(temp);
}
}
}
|