Add IBufferProtocol to mmap

Author: slozier

src/core/IronPython.Modules/mmap.cs

@@ -7,6 +7,8 @@
 #if FEATURE_MMAP
 
 using System;
+using System.Buffers;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.Globalization;
 using System.IO;
@@ -278,7 +280,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;
@@ -559,10 +561,8 @@ public void __exit__(CodeContext/*!*/ context, params object[] excinfo) {
                 close();
             }
 
-
             public bool closed => (_state & StateBits.Closed) == StateBits.Closed;  // Dispose already requested, will self-dispose when ref count drops to 0.
 
-
             private bool AddRef() {
                 int oldState, newState;
                 do {
@@ -606,6 +606,14 @@ private void Release() {
 
 
             public void close() {
+                if ((_state & StateBits.RefCount) != 0) {
+                    GC.Collect();
+                    GC.WaitForPendingFinalizers();
+                    if ((_state & StateBits.RefCount) != 0) {
+                        throw PythonOps.BufferError("cannot close exported pointers exist");
+                    }
+                }
+
                 // close is idempotent; it must never block
 #if NET5_0_OR_GREATER
                 if ((Interlocked.Or(ref _state, StateBits.Closed) & StateBits.Closed) != StateBits.Closed) {
@@ -614,12 +622,10 @@ public void close() {
                 }
 #else
                 int current = _state;
-                while (true)
-                {
+                while (true) {
                     int newState = current | StateBits.Closed;
                     int oldState = Interlocked.CompareExchange(ref _state, newState, current);
-                    if (oldState == current)
-                    {
+                    if (oldState == current) {
                         // didn't change in the meantime, exchange with newState completed
                         if ((oldState & StateBits.Closed) != StateBits.Closed) {
                             // freshly closed, release the construction time reference
@@ -862,6 +868,14 @@ public string readline() {
             }
 
             public void resize(long newsize) {
+                if ((_state & StateBits.RefCount) != 0) {
+                    GC.Collect();
+                    GC.WaitForPendingFinalizers();
+                    if ((_state & StateBits.RefCount) != 0) {
+                        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.");
@@ -894,13 +908,13 @@ public void resize(long newsize) {
                             int fd = unchecked((int)_handle.DangerousGetHandle());
                             PythonNT.ftruncateUnix(fd, newsize);
 
-    #if NET8_0_OR_GREATER
+#if NET8_0_OR_GREATER
                             _file = MemoryMappedFile.CreateFromFile(_handle, _mapName, newsize, _fileAccess, HandleInheritability.None, leaveOpen: true);
-    #else
+#else
                             _sourceStream?.Dispose();
                             _sourceStream = new FileStream(new SafeFileHandle((IntPtr)fd, ownsHandle: false), FileAccess.ReadWrite);
                             _file = CreateFromFile(_sourceStream, _mapName, newsize, _fileAccess, HandleInheritability.None, leaveOpen: true);
-    #endif
+#endif
                             _view = _file.CreateViewAccessor(_offset, newsize, _fileAccess);
                             return;
                         } catch {
@@ -1130,8 +1144,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.");
                 }
             }
@@ -1247,6 +1263,77 @@ void IWeakReferenceable.SetFinalizer(WeakRefTracker value) {
             }
 
             #endregion
+
+            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 unsafe class MmapBuffer : IPythonBuffer {
+                private readonly MmapDefault _mmap;
+                private readonly BufferFlags _flags;
+                private SafeMemoryMappedViewHandle? _handle;
+                private byte* _pointer = null;
+
+                public MmapBuffer(MmapDefault mmap, BufferFlags flags) {
+                    _mmap = mmap;
+                    _flags = flags;
+                    if (!_mmap.AddRef()) {
+                        throw PythonOps.ValueError("mmap closed or invalid");
+                    }
+                    _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));
+                    if (_pointer is 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");
+                    if (_pointer is null) _handle.AcquirePointer(ref _pointer);
+                    return new Span<byte>(_pointer, ItemCount);
+                }
+
+                public unsafe MemoryHandle Pin() {
+                    if (_handle is null) throw new ObjectDisposedException(nameof(MmapBuffer));
+                    if (_pointer is null) _handle.AcquirePointer(ref _pointer);
+                    return new MemoryHandle(_pointer);
+                }
+
+                public void Dispose() {
+                    var handle = Interlocked.Exchange(ref _handle, null);
+                    if (handle is null) return;
+                    if (_pointer is not null) handle.ReleasePointer();
+                    _mmap.Release();
+                }
+
+            }
         }
 
 

src/core/IronPython.Modules/re.cs

@@ -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");
@@ -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");