Skip to content

Commit

Permalink
Add IBufferProtocol to mmap
Browse files Browse the repository at this point in the history
  • Loading branch information
slozier committed Jan 11, 2025
1 parent 9082851 commit e068943
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 16 deletions.
114 changes: 106 additions & 8 deletions Src/IronPython.Modules/mmap.cs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
#if FEATURE_MMAP

using System;
using System.Buffers;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.Globalization;
Expand Down Expand Up @@ -277,7 +279,7 @@ private static MemoryMappedFileAccess ToMmapFileAccess(int access) {
}

[PythonHidden]
public class MmapDefault : IWeakReferenceable {
public class MmapDefault : IWeakReferenceable, IBufferProtocol {
private MemoryMappedFile _file;
private MemoryMappedViewAccessor _view;
private long _position;
Expand Down Expand Up @@ -570,12 +572,20 @@ public void __exit__(CodeContext/*!*/ context, params object[] excinfo) {
public bool closed => _isClosed;

public void close() {
if (!_isClosed) {
lock (this) {
if (!_isClosed) {
_isClosed = true;
CloseWorker();
}
if (_isClosed) return;

if (_refCount > 1) {
GC.Collect();
GC.WaitForPendingFinalizers();
if (_refCount > 1) {
throw PythonOps.BufferError("cannot close exported pointers exist");
}
}

lock (this) {
if (!_isClosed) {
_isClosed = true;
CloseWorker();
}
}
}
Expand Down Expand Up @@ -820,6 +830,14 @@ public string readline() {
}

public void resize(long newsize) {
if (_refCount > 1) {
GC.Collect();
GC.WaitForPendingFinalizers();
if (_refCount > 1) {
throw PythonOps.BufferError("mmap can't resize with extant buffers exported.");
}
}

using (new MmapLocker(this)) {
if (_fileAccess is not MemoryMappedFileAccess.ReadWrite and not MemoryMappedFileAccess.ReadWriteExecute) {
throw PythonOps.TypeError("mmap can't resize a readonly or copy-on-write memory map.");
Expand Down Expand Up @@ -1076,8 +1094,10 @@ private long Position {
}
}

private bool IsReadOnly => _fileAccess is MemoryMappedFileAccess.Read or MemoryMappedFileAccess.ReadExecute;

private void EnsureWritable() {
if (_fileAccess is MemoryMappedFileAccess.Read or MemoryMappedFileAccess.ReadExecute) {
if (IsReadOnly) {
throw PythonOps.TypeError("mmap can't modify a read-only memory map.");
}
}
Expand Down Expand Up @@ -1198,6 +1218,84 @@ void IWeakReferenceable.SetFinalizer(WeakRefTracker value) {
}

#endregion

#nullable enable

public IPythonBuffer GetBuffer(BufferFlags flags = BufferFlags.Simple) {
if (flags.HasFlag(BufferFlags.Writable) && IsReadOnly)
throw PythonOps.BufferError("Object is not writable.");

return new MmapBuffer(this, flags);
}

private sealed class MmapBuffer : IPythonBuffer {
private readonly MmapDefault _mmap;
private readonly BufferFlags _flags;
private SafeMemoryMappedViewHandle? _handle;

public MmapBuffer(MmapDefault mmap, BufferFlags flags) {
mmap.EnsureOpen();

_mmap = mmap;
_flags = flags;
Interlocked.Increment(ref _mmap._refCount);
_handle = _mmap._view.SafeMemoryMappedViewHandle;
ItemCount = _mmap.__len__() is int i ? i : throw new NotImplementedException();
}

public object Object => _mmap;

public bool IsReadOnly => _mmap.IsReadOnly;

public int Offset => 0;

public string? Format => _flags.HasFlag(BufferFlags.Format) ? "B" : null;

public int ItemCount { get; }

public int ItemSize => 1;

public int NumOfDims => 1;

public IReadOnlyList<int>? Shape => null;

public IReadOnlyList<int>? Strides => null;

public IReadOnlyList<int>? SubOffsets => null;

public unsafe ReadOnlySpan<byte> AsReadOnlySpan() {
if (_handle is null) throw new ObjectDisposedException(nameof(MmapBuffer));
byte* pointer = null;
_handle.AcquirePointer(ref pointer);
return new ReadOnlySpan<byte>(pointer, ItemCount);
}

public unsafe Span<byte> AsSpan() {
if (_handle is null) throw new ObjectDisposedException(nameof(MmapBuffer));
if (IsReadOnly) throw new InvalidOperationException("object is not writable");
byte* pointer = null;
_handle.AcquirePointer(ref pointer);
return new Span<byte>(pointer, ItemCount);
}

public unsafe MemoryHandle Pin() {
if (_handle is null) throw new ObjectDisposedException(nameof(MmapBuffer));
byte* pointer = null;
_handle.AcquirePointer(ref pointer);
return new MemoryHandle(pointer);
}

public void Dispose() {
var handle = Interlocked.Exchange(ref _handle, null);
if (handle is null) return;
handle.Dispose();
_mmap.CloseWorker();
}

}

#nullable restore

}

#region P/Invoke for allocation granularity
Expand Down
8 changes: 0 additions & 8 deletions Src/IronPython.Modules/re.cs
Original file line number Diff line number Diff line change
Expand Up @@ -516,11 +516,6 @@ private string ValidateString(object? @string) {
case IList<byte> b:
str = b.MakeString();
break;
#if FEATURE_MMAP
case MmapModule.MmapDefault mmapFile:
str = mmapFile.GetSearchString().MakeString();
break;
#endif
case string _:
case ExtensibleString _:
throw PythonOps.TypeError("cannot use a bytes pattern on a string-like object");
Expand All @@ -537,9 +532,6 @@ private string ValidateString(object? @string) {
break;
case IBufferProtocol _:
case IList<byte> _:
#if FEATURE_MMAP
case MmapModule.MmapDefault _:
#endif
throw PythonOps.TypeError("cannot use a string pattern on a bytes-like object");
default:
throw PythonOps.TypeError("expected string or bytes-like object");
Expand Down

0 comments on commit e068943

Please sign in to comment.