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 6, 2025
1 parent aaf701d commit 1ab991f
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,7 @@
#if FEATURE_MMAP

using System;
using System.Buffers;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
Expand All @@ -24,6 +25,7 @@
using IronPython.Runtime.Types;

using Microsoft.Scripting.Utils;
using Microsoft.Win32.SafeHandles;

[assembly: PythonModule("mmap", typeof(IronPython.Modules.MmapModule))]
namespace IronPython.Modules {
Expand Down Expand Up @@ -127,7 +129,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 @@ -349,12 +351,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 @@ -590,6 +600,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 @@ -813,8 +831,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 @@ -916,6 +936,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 1ab991f

Please sign in to comment.