Add IBufferProtocol to mmap

Author: slozier

Src/IronPython.Modules/mmap.cs

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

Src/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");