-
Notifications
You must be signed in to change notification settings - Fork 119
/
SpanOwner.cs
198 lines (176 loc) · 7.54 KB
/
SpanOwner.cs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
using System.Buffers;
using System.ComponentModel;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
namespace DotNext.Buffers;
using Intrinsics = Runtime.Intrinsics;
/// <summary>
/// Represents the memory obtained from the pool or allocated
/// on the stack or heap.
/// </summary>
/// <remarks>
/// This type is aimed to be compatible with memory allocated using <c>stackalloc</c> operator.
/// If stack allocation threshold is reached (e.g. <see cref="StackallocThreshold"/>) then it's possible to use pooled memory from
/// arbitrary <see cref="MemoryPool{T}"/> or <see cref="ArrayPool{T}.Shared"/>. Custom
/// <see cref="ArrayPool{T}"/> is not supported because default <see cref="ArrayPool{T}.Shared"/>
/// is optimized for per-CPU core allocation which is perfect when the same
/// thread is responsible for renting and releasing the array.
/// </remarks>
/// <example>
/// <code>
/// const int stackallocThreshold = 20;
/// var memory = size <=stackallocThreshold ? new SpanOwner<byte>(stackalloc byte[stackallocThreshold], size) : new SpanOwner<byte>(size);
/// </code>
/// </example>
/// <typeparam name="T">The type of the elements in the rented memory.</typeparam>
[StructLayout(LayoutKind.Auto)]
public ref struct SpanOwner<T>
{
/// <summary>
/// Global recommended number of elements that can be allocated on the stack.
/// </summary>
/// <remarks>
/// This property is for internal purposes only and should not be referenced
/// directly in your code.
/// </remarks>
[EditorBrowsable(EditorBrowsableState.Never)]
[CLSCompliant(false)]
public static int StackallocThreshold { get; } = 1 + (LibrarySettings.StackallocThreshold / Unsafe.SizeOf<T>());
private readonly object? owner;
private readonly Span<T> memory;
/// <summary>
/// Rents the memory referenced by the span.
/// </summary>
/// <param name="span">The span that references the memory to rent.</param>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public SpanOwner(Span<T> span)
=> memory = span;
/// <summary>
/// Rents the memory referenced by the span.
/// </summary>
/// <param name="span">The span that references the memory to rent.</param>
/// <param name="length">The actual length of the data.</param>
public SpanOwner(Span<T> span, int length)
: this(span.Slice(0, length))
{
}
/// <summary>
/// Rents the memory from the pool.
/// </summary>
/// <param name="pool">The memory pool.</param>
/// <param name="minBufferSize">The minimum size of the memory to rent.</param>
/// <param name="exactSize"><see langword="true"/> to return the buffer of <paramref name="minBufferSize"/> length; otherwise, the returned buffer is at least of <paramref name="minBufferSize"/>.</param>
/// <exception cref="ArgumentNullException"><paramref name="pool"/> is <see langword="null"/>.</exception>
/// <exception cref="ArgumentOutOfRangeException"><paramref name="minBufferSize"/> is less than or equal to zero.</exception>
public SpanOwner(MemoryPool<T> pool, int minBufferSize, bool exactSize = true)
{
ArgumentNullException.ThrowIfNull(pool);
ArgumentOutOfRangeException.ThrowIfNegativeOrZero(minBufferSize);
var owner = pool.Rent(minBufferSize);
memory = owner.Memory.Span;
if (exactSize)
memory = memory.Slice(0, minBufferSize);
this.owner = owner;
}
/// <summary>
/// Rents the memory from the pool.
/// </summary>
/// <param name="pool">The memory pool.</param>
/// <exception cref="ArgumentNullException"><paramref name="pool"/> is <see langword="null"/>.</exception>
public SpanOwner(MemoryPool<T> pool)
{
ArgumentNullException.ThrowIfNull(pool);
var owner = pool.Rent();
memory = owner.Memory.Span;
this.owner = owner;
}
/// <summary>
/// Rents the memory from <see cref="ArrayPool{T}.Shared"/>, if <typeparamref name="T"/>
/// contains at least one field of reference type; or use <see cref="NativeMemory"/>.
/// </summary>
/// <param name="minBufferSize">The minimum size of the memory to rent.</param>
/// <param name="exactSize"><see langword="true"/> to return the buffer of <paramref name="minBufferSize"/> length; otherwise, the returned buffer is at least of <paramref name="minBufferSize"/>.</param>
/// <exception cref="ArgumentOutOfRangeException"><paramref name="minBufferSize"/> is less than or equal to zero.</exception>
public SpanOwner(int minBufferSize, bool exactSize = true)
{
if (UseNativeAllocation)
{
ArgumentOutOfRangeException.ThrowIfNegativeOrZero(minBufferSize);
unsafe
{
var ptr = NativeMemory.Alloc((uint)minBufferSize, (uint)Unsafe.SizeOf<T>());
memory = new(ptr, minBufferSize);
}
owner = Sentinel.Instance;
}
else
{
var owner = ArrayPool<T>.Shared.Rent(minBufferSize);
memory = exactSize ? new(owner, 0, minBufferSize) : new(owner);
this.owner = owner;
}
}
private static bool UseNativeAllocation
=> !LibrarySettings.DisableNativeAllocation && !RuntimeHelpers.IsReferenceOrContainsReferences<T>() && Intrinsics.AlignOf<T>() <= nuint.Size;
/// <summary>
/// Gets the rented memory.
/// </summary>
public readonly Span<T> Span => memory;
/// <summary>
/// Gets a value indicating that this object
/// doesn't reference rented memory.
/// </summary>
public readonly bool IsEmpty => memory.IsEmpty;
/// <summary>
/// Converts the reference to the already allocated memory
/// into the rental object.
/// </summary>
/// <param name="span">The allocated memory to convert.</param>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static implicit operator SpanOwner<T>(Span<T> span)
=> new(span);
/// <summary>
/// Gets length of the rented memory.
/// </summary>
public readonly int Length => memory.Length;
/// <summary>
/// Gets the memory element by its index.
/// </summary>
/// <param name="index">The index of the memory element.</param>
/// <returns>The managed pointer to the memory element.</returns>
public readonly ref T this[int index] => ref memory[index];
/// <summary>
/// Obtains managed pointer to the first element of the rented array.
/// </summary>
/// <returns>The managed pointer to the first element of the rented array.</returns>
[EditorBrowsable(EditorBrowsableState.Never)]
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public readonly ref T GetPinnableReference() => ref memory.GetPinnableReference();
/// <summary>
/// Gets textual representation of the rented memory.
/// </summary>
/// <returns>The textual representation of the rented memory.</returns>
public override readonly string ToString() => memory.ToString();
/// <summary>
/// Returns the memory back to the pool.
/// </summary>
public void Dispose()
{
if (owner is T[] array)
{
ArrayPool<T>.Shared.Return(array, RuntimeHelpers.IsReferenceOrContainsReferences<T>());
}
else if (ReferenceEquals(owner, Sentinel.Instance))
{
unsafe
{
NativeMemory.Free(Unsafe.AsPointer(ref MemoryMarshal.GetReference(memory)));
}
}
else
{
Unsafe.As<IDisposable>(owner)?.Dispose();
}
this = default;
}
}