Convert MemoryBufferProtocolWrapper to generic (#1906)

* Convert `MemoryBufferProtocolWrapper` to generic * Autodetect type code
IronLanguages · Feb 6, 2025 · 662a97d · 662a97d
1 parent 3f31f45
commit 662a97d
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 50 deletions.
diff --git a/src/core/IronPython/Runtime/Binding/ConversionBinder.cs b/src/core/IronPython/Runtime/Binding/ConversionBinder.cs
@@ -734,7 +734,7 @@ private DynamicMetaObject ConvertFromMemoryToBufferProtocol(DynamicMetaObject se
             return new DynamicMetaObject(
                 AstUtils.Convert(
                     Ast.New(
-                        typeof(MemoryBufferProtocolWrapper).GetConstructor(new Type[] { fromType }),
+                        typeof(MemoryBufferProtocolWrapper<byte>).GetConstructor([fromType]),
                         AstUtils.Convert(self.Expression, fromType)
                     ),
                     typeof(IBufferProtocol)

diff --git a/src/core/IronPython/Runtime/ConversionWrappers.cs b/src/core/IronPython/Runtime/ConversionWrappers.cs
@@ -10,6 +10,7 @@
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Numerics;
+using System.Runtime.InteropServices;
 
 namespace IronPython.Runtime {
 
@@ -376,79 +377,90 @@ IEnumerator IEnumerable.GetEnumerator() {
         #endregion
     }
 
-    public sealed class MemoryBufferWrapper : IPythonBuffer {
-        private readonly ReadOnlyMemory<byte> _rom;
-        private readonly Memory<byte>? _memory;
-        private readonly BufferFlags _flags;
+    public sealed class MemoryBufferProtocolWrapper<T> : IBufferProtocol where T : unmanaged {
+        private readonly ReadOnlyMemory<T> _rom;
+        private readonly Memory<T>? _memory;
+        private readonly char _format;
 
-        public MemoryBufferWrapper(ReadOnlyMemory<byte> memory, BufferFlags flags) {
+        public MemoryBufferProtocolWrapper(ReadOnlyMemory<T> memory) {
             _rom = memory;
             _memory = null;
-            _flags = flags;
+            _format = GetFormatChar();
         }
 
-        public MemoryBufferWrapper(Memory<byte> memory, BufferFlags flags) {
+        public MemoryBufferProtocolWrapper(Memory<T> memory) {
             _rom = memory;
             _memory = memory;
-            _flags = flags;
+            _format = GetFormatChar();
         }
 
-        public void Dispose() { }
-
-        public object Object => _memory ?? _rom;
-
-        public bool IsReadOnly => !_memory.HasValue;
+        public IPythonBuffer? GetBuffer(BufferFlags flags, bool throwOnError) {
+            if (flags.HasFlag(BufferFlags.Writable) && !_memory.HasValue) {
+                if (throwOnError) {
+                    throw Operations.PythonOps.BufferError("ReadOnlyMemory is not writable.");
+                }
+                return null;
+            }
 
-        public ReadOnlySpan<byte> AsReadOnlySpan() => _rom.Span;
+            return new MemoryBufferWrapper(this, flags);
+        }
+
+        private static char GetFormatChar()
+            => Type.GetTypeCode(typeof(T)) switch {
+                TypeCode.SByte   => 'b',
+                TypeCode.Byte    => 'B',
+                TypeCode.Char    => 'u',
+                TypeCode.Int16   => 'h',
+                TypeCode.UInt16  => 'H',
+                TypeCode.Int32   => 'i',
+                TypeCode.UInt32  => 'I',
+                TypeCode.Int64   => 'q',
+                TypeCode.UInt64  => 'Q',
+                TypeCode.Single  => 'f',
+                TypeCode.Double  => 'd',
+                _ => throw new ArgumentException("Unsupported type"),
+            };
+
+
+        private sealed unsafe class MemoryBufferWrapper : IPythonBuffer {
+            private readonly MemoryBufferProtocolWrapper<T> _wrapper;
+            private readonly BufferFlags _flags;
+
+            public MemoryBufferWrapper(MemoryBufferProtocolWrapper<T> wrapper, BufferFlags flags) {
+                _wrapper = wrapper;
+                _flags = flags;
+            }
 
-        public Span<byte> AsSpan() => _memory.HasValue ? _memory.Value.Span : throw new InvalidOperationException("ReadOnlyMemory is not writable");
+            public void Dispose() { }
 
-        public MemoryHandle Pin() => _rom.Pin();
+            public object Object => _wrapper._memory ?? _wrapper._rom;
 
-        public int Offset => 0;
+            public bool IsReadOnly => !_wrapper._memory.HasValue;
 
-        public string? Format => _flags.HasFlag(BufferFlags.Format) ? "B" : null;
+            public ReadOnlySpan<byte> AsReadOnlySpan() => MemoryMarshal.Cast<T, byte>(_wrapper._rom.Span);
 
-        public int ItemCount => _rom.Length;
+            public Span<byte> AsSpan() 
+                => _wrapper._memory.HasValue
+                    ? MemoryMarshal.Cast<T, byte>(_wrapper._memory.Value.Span)
+                    : throw new InvalidOperationException("ReadOnlyMemory is not writable");
 
-        public int ItemSize => 1;
+            public MemoryHandle Pin() => _wrapper._rom.Pin();
 
-        public int NumOfDims => 1;
+            public int Offset => 0;
 
-        public IReadOnlyList<int>? Shape => null;
+            public string? Format => _flags.HasFlag(BufferFlags.Format) ? _wrapper._format.ToString() : null;
 
-        public IReadOnlyList<int>? Strides => null;
+            public int ItemCount => _wrapper._rom.Length;
 
-        public IReadOnlyList<int>? SubOffsets => null;
-    }
+            public int ItemSize => sizeof(T);
 
-    public class MemoryBufferProtocolWrapper : IBufferProtocol {
-        private readonly ReadOnlyMemory<byte> _rom;
-        private readonly Memory<byte>? _memory;
+            public int NumOfDims => 1;
 
-        public MemoryBufferProtocolWrapper(ReadOnlyMemory<byte> memory) {
-            _rom = memory;
-            _memory = null;
-        }
+            public IReadOnlyList<int>? Shape => null;
 
-        public MemoryBufferProtocolWrapper(Memory<byte> memory) {
-            _rom = memory;
-            _memory = memory;
-        }
-
-        public IPythonBuffer? GetBuffer(BufferFlags flags, bool throwOnError) {
-            if (_memory.HasValue) {
-                return new MemoryBufferWrapper(_memory.Value, flags);
-            }
-
-            if (flags.HasFlag(BufferFlags.Writable)) {
-                if (throwOnError) {
-                    throw Operations.PythonOps.BufferError("ReadOnlyMemory is not writable.");
-                }
-                return null;
-            }
+            public IReadOnlyList<int>? Strides => null;
 
-            return new MemoryBufferWrapper(_rom, flags);
+            public IReadOnlyList<int>? SubOffsets => null;
         }
     }
 }