Skip to content

Commit a472936

Browse files
committed
Add IBufferProtocol to mmap
1 parent 2aca667 commit a472936

File tree

2 files changed

+98
-17
lines changed

2 files changed

+98
-17
lines changed

src/core/IronPython.Modules/mmap.cs

Lines changed: 98 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
#if FEATURE_MMAP
88

99
using System;
10+
using System.Buffers;
11+
using System.Collections.Generic;
1012
using System.Diagnostics;
1113
using System.Globalization;
1214
using System.IO;
@@ -278,7 +280,7 @@ private static MemoryMappedFileAccess ToMmapFileAccess(int access) {
278280
}
279281

280282
[PythonHidden]
281-
public class MmapDefault : IWeakReferenceable {
283+
public class MmapDefault : IWeakReferenceable, IBufferProtocol {
282284
private MemoryMappedFile _file;
283285
private MemoryMappedViewAccessor _view;
284286
private long _position;
@@ -606,6 +608,14 @@ private void Release() {
606608

607609

608610
public void close() {
611+
if ((_state & StateBits.RefCount) != 0) {
612+
GC.Collect();
613+
GC.WaitForPendingFinalizers();
614+
if ((_state & StateBits.RefCount) != 0) {
615+
throw PythonOps.BufferError("cannot close exported pointers exist");
616+
}
617+
}
618+
609619
// close is idempotent; it must never block
610620
#if NET5_0_OR_GREATER
611621
if ((Interlocked.Or(ref _state, StateBits.Closed) & StateBits.Closed) != StateBits.Closed) {
@@ -614,12 +624,10 @@ public void close() {
614624
}
615625
#else
616626
int current = _state;
617-
while (true)
618-
{
627+
while (true) {
619628
int newState = current | StateBits.Closed;
620629
int oldState = Interlocked.CompareExchange(ref _state, newState, current);
621-
if (oldState == current)
622-
{
630+
if (oldState == current) {
623631
// didn't change in the meantime, exchange with newState completed
624632
if ((oldState & StateBits.Closed) != StateBits.Closed) {
625633
// freshly closed, release the construction time reference
@@ -862,6 +870,14 @@ public string readline() {
862870
}
863871

864872
public void resize(long newsize) {
873+
if ((_state & StateBits.RefCount) != 0) {
874+
GC.Collect();
875+
GC.WaitForPendingFinalizers();
876+
if ((_state & StateBits.RefCount) != 0) {
877+
throw PythonOps.BufferError("mmap can't resize with extant buffers exported.");
878+
}
879+
}
880+
865881
using (new MmapLocker(this)) {
866882
if (_fileAccess is not MemoryMappedFileAccess.ReadWrite and not MemoryMappedFileAccess.ReadWriteExecute) {
867883
throw PythonOps.TypeError("mmap can't resize a readonly or copy-on-write memory map.");
@@ -894,13 +910,13 @@ public void resize(long newsize) {
894910
int fd = unchecked((int)_handle.DangerousGetHandle());
895911
PythonNT.ftruncateUnix(fd, newsize);
896912

897-
#if NET8_0_OR_GREATER
913+
#if NET8_0_OR_GREATER
898914
_file = MemoryMappedFile.CreateFromFile(_handle, _mapName, newsize, _fileAccess, HandleInheritability.None, leaveOpen: true);
899-
#else
915+
#else
900916
_sourceStream?.Dispose();
901917
_sourceStream = new FileStream(new SafeFileHandle((IntPtr)fd, ownsHandle: false), FileAccess.ReadWrite);
902918
_file = CreateFromFile(_sourceStream, _mapName, newsize, _fileAccess, HandleInheritability.None, leaveOpen: true);
903-
#endif
919+
#endif
904920
_view = _file.CreateViewAccessor(_offset, newsize, _fileAccess);
905921
return;
906922
} catch {
@@ -1130,8 +1146,10 @@ private long Position {
11301146
}
11311147
}
11321148

1149+
private bool IsReadOnly => _fileAccess is MemoryMappedFileAccess.Read or MemoryMappedFileAccess.ReadExecute;
1150+
11331151
private void EnsureWritable() {
1134-
if (_fileAccess is MemoryMappedFileAccess.Read or MemoryMappedFileAccess.ReadExecute) {
1152+
if (IsReadOnly) {
11351153
throw PythonOps.TypeError("mmap can't modify a read-only memory map.");
11361154
}
11371155
}
@@ -1247,6 +1265,77 @@ void IWeakReferenceable.SetFinalizer(WeakRefTracker value) {
12471265
}
12481266

12491267
#endregion
1268+
1269+
public IPythonBuffer GetBuffer(BufferFlags flags = BufferFlags.Simple) {
1270+
if (flags.HasFlag(BufferFlags.Writable) && IsReadOnly)
1271+
throw PythonOps.BufferError("Object is not writable.");
1272+
1273+
return new MmapBuffer(this, flags);
1274+
}
1275+
1276+
private sealed unsafe class MmapBuffer : IPythonBuffer {
1277+
private readonly MmapDefault _mmap;
1278+
private readonly BufferFlags _flags;
1279+
private SafeMemoryMappedViewHandle? _handle;
1280+
private byte* _pointer = null;
1281+
1282+
public MmapBuffer(MmapDefault mmap, BufferFlags flags) {
1283+
_mmap = mmap;
1284+
_flags = flags;
1285+
if (!_mmap.AddRef()) {
1286+
throw PythonOps.ValueError("mmap closed or invalid");
1287+
}
1288+
_handle = _mmap._view.SafeMemoryMappedViewHandle;
1289+
ItemCount = _mmap.__len__() is int i ? i : throw new NotImplementedException();
1290+
}
1291+
1292+
public object Object => _mmap;
1293+
1294+
public bool IsReadOnly => _mmap.IsReadOnly;
1295+
1296+
public int Offset => 0;
1297+
1298+
public string? Format => _flags.HasFlag(BufferFlags.Format) ? "B" : null;
1299+
1300+
public int ItemCount { get; }
1301+
1302+
public int ItemSize => 1;
1303+
1304+
public int NumOfDims => 1;
1305+
1306+
public IReadOnlyList<int>? Shape => null;
1307+
1308+
public IReadOnlyList<int>? Strides => null;
1309+
1310+
public IReadOnlyList<int>? SubOffsets => null;
1311+
1312+
public unsafe ReadOnlySpan<byte> AsReadOnlySpan() {
1313+
if (_handle is null) throw new ObjectDisposedException(nameof(MmapBuffer));
1314+
if (_pointer is null) _handle.AcquirePointer(ref _pointer);
1315+
return new ReadOnlySpan<byte>(_pointer, ItemCount);
1316+
}
1317+
1318+
public unsafe Span<byte> AsSpan() {
1319+
if (_handle is null) throw new ObjectDisposedException(nameof(MmapBuffer));
1320+
if (IsReadOnly) throw new InvalidOperationException("object is not writable");
1321+
if (_pointer is null) _handle.AcquirePointer(ref _pointer);
1322+
return new Span<byte>(_pointer, ItemCount);
1323+
}
1324+
1325+
public unsafe MemoryHandle Pin() {
1326+
if (_handle is null) throw new ObjectDisposedException(nameof(MmapBuffer));
1327+
if (_pointer is null) _handle.AcquirePointer(ref _pointer);
1328+
return new MemoryHandle(_pointer);
1329+
}
1330+
1331+
public void Dispose() {
1332+
var handle = Interlocked.Exchange(ref _handle, null);
1333+
if (handle is null) return;
1334+
if (_pointer is not null) handle.ReleasePointer();
1335+
_mmap.Release();
1336+
}
1337+
1338+
}
12501339
}
12511340

12521341

src/core/IronPython.Modules/re.cs

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -516,11 +516,6 @@ private string ValidateString(object? @string) {
516516
case IList<byte> b:
517517
str = b.MakeString();
518518
break;
519-
#if FEATURE_MMAP
520-
case MmapModule.MmapDefault mmapFile:
521-
str = mmapFile.GetSearchString().MakeString();
522-
break;
523-
#endif
524519
case string _:
525520
case ExtensibleString _:
526521
throw PythonOps.TypeError("cannot use a bytes pattern on a string-like object");
@@ -537,9 +532,6 @@ private string ValidateString(object? @string) {
537532
break;
538533
case IBufferProtocol _:
539534
case IList<byte> _:
540-
#if FEATURE_MMAP
541-
case MmapModule.MmapDefault _:
542-
#endif
543535
throw PythonOps.TypeError("cannot use a string pattern on a bytes-like object");
544536
default:
545537
throw PythonOps.TypeError("expected string or bytes-like object");

0 commit comments

Comments
 (0)