From 5854050bb99d592d2e46bf0bfe0a89950896d704 Mon Sep 17 00:00:00 2001
From: Aaron Robinson <aaronrobin1234@gmail.com>
Date: Thu, 11 Apr 2024 00:58:25 -0500
Subject: [PATCH] Switch MappedMappedDataSource to UnmanagedDataSource

---
 src/AsmResolver/IO/MemoryMappedDataSource.cs  | 65 -------------------
 src/AsmResolver/IO/MemoryMappedInputFile.cs   | 11 ++--
 src/AsmResolver/Shims/MemoryMappedFileShim.cs | 11 +---
 .../IO/MemoryMappedIOTest.cs                  |  1 +
 4 files changed, 7 insertions(+), 81 deletions(-)
 delete mode 100644 src/AsmResolver/IO/MemoryMappedDataSource.cs

diff --git a/src/AsmResolver/IO/MemoryMappedDataSource.cs b/src/AsmResolver/IO/MemoryMappedDataSource.cs
deleted file mode 100644
index f47202bba..000000000
--- a/src/AsmResolver/IO/MemoryMappedDataSource.cs
+++ /dev/null
@@ -1,65 +0,0 @@
-using System;
-using AsmResolver.Shims;
-
-namespace AsmResolver.IO
-{
-    /// <summary>
-    /// Represents a data source that obtains its data from a memory mapped file.
-    /// </summary>
-    public sealed class MemoryMappedDataSource : IDataSource, IDisposable
-#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP2_1_OR_GREATER
-        , ISpanDataSource
-#endif
-    {
-        private readonly MemoryMappedFileShim _file;
-
-        /// <summary>
-        /// Creates a new instance of the <see cref="MemoryMappedDataSource"/> class.
-        /// </summary>
-        /// <param name="file">The memory mapped file to use.</param>
-        internal MemoryMappedDataSource(MemoryMappedFileShim file)
-        {
-            _file = file ?? throw new ArgumentNullException(nameof(file));
-        }
-
-        /// <inheritdoc />
-        public ulong BaseAddress => 0;
-
-        /// <inheritdoc />
-        public byte this[ulong address] => _file.ReadByte((long) address);
-
-        /// <inheritdoc />
-        public ulong Length => (ulong)_file.Size;
-
-        /// <inheritdoc />
-        public bool IsValidAddress(ulong address) => address < Length;
-
-        /// <inheritdoc />
-        public int ReadBytes(ulong address, byte[] buffer, int index, int count)
-        {
-            if (address >= Length)
-                throw new ArgumentOutOfRangeException(nameof(address));
-            int actualLength = (int)Math.Min(Length - address, (ulong)count);
-            _file.GetSpan((long)address, actualLength).CopyTo(buffer.AsSpan(index));
-            return actualLength;
-        }
-
-#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP2_1_OR_GREATER
-        /// <inheritdoc />
-        public unsafe int ReadBytes(ulong address, Span<byte> buffer)
-        {
-            if (!IsValidAddress(address))
-                return 0;
-
-            int actualLength = (int) Math.Min(Length - address, (ulong)buffer.Length);
-
-            _file.GetSpan((long)address, actualLength).CopyTo(buffer);
-
-            return actualLength;
-        }
-#endif
-
-        /// <inheritdoc />
-        public void Dispose() => _file.Dispose();
-    }
-}
diff --git a/src/AsmResolver/IO/MemoryMappedInputFile.cs b/src/AsmResolver/IO/MemoryMappedInputFile.cs
index 5f5a82258..05bb2071a 100644
--- a/src/AsmResolver/IO/MemoryMappedInputFile.cs
+++ b/src/AsmResolver/IO/MemoryMappedInputFile.cs
@@ -10,17 +10,17 @@ namespace AsmResolver.IO
     public sealed class MemoryMappedInputFile : IInputFile
     {
         private readonly MemoryMappedFileShim _file;
-        private readonly MemoryMappedDataSource _dataSource;
+        private readonly UnmanagedDataSource _dataSource;
 
         /// <summary>
         /// Creates a new reader factory for the provided file.
         /// </summary>
         /// <param name="filePath">The path to the file to read.</param>
-        public MemoryMappedInputFile(string filePath)
+        public unsafe MemoryMappedInputFile(string filePath)
         {
             FilePath = filePath ?? throw new ArgumentNullException(nameof(filePath));
             _file = new MemoryMappedFileShim(filePath);
-            _dataSource = new MemoryMappedDataSource(_file);
+            _dataSource = new UnmanagedDataSource(_file.BasePointer, (ulong)_file.Size);
         }
 
         /// <inheritdoc />
@@ -33,14 +33,13 @@ public string FilePath
         public uint Length => (uint) _dataSource.Length;
 
         /// <inheritdoc />
-        public BinaryStreamReader CreateReader(ulong address, uint rva, uint length) =>
-            new(_dataSource, address, rva, length);
+        public unsafe BinaryStreamReader CreateReader(ulong address, uint rva, uint length) =>
+            new(_dataSource, address != 0 ? address : (ulong)_file.BasePointer, rva, length);
 
         /// <inheritdoc />
         public void Dispose()
         {
             _file.Dispose();
-            _dataSource.Dispose();
         }
     }
 }
diff --git a/src/AsmResolver/Shims/MemoryMappedFileShim.cs b/src/AsmResolver/Shims/MemoryMappedFileShim.cs
index 272a363cf..954cb5ee2 100644
--- a/src/AsmResolver/Shims/MemoryMappedFileShim.cs
+++ b/src/AsmResolver/Shims/MemoryMappedFileShim.cs
@@ -17,6 +17,7 @@ public MemoryMappedFileShim(string path)
     }
 
     public long Size => _size;
+    public byte* BasePointer => _file;
 
     public byte ReadByte(long address)
     {
@@ -27,16 +28,6 @@ public byte ReadByte(long address)
         return _file[address];
     }
 
-    public ReadOnlySpan<byte> GetSpan(long offset, int length)
-    {
-        if (_file == null)
-            throw new ObjectDisposedException("disposed");
-        long newSize = _size - offset;
-        if (offset < 0 || newSize < 0 || length > newSize)
-            throw new ArgumentOutOfRangeException();
-        return new(_file + offset, length);
-    }
-
     private void DisposeCore()
     {
         void* filePointer;
diff --git a/test/AsmResolver.Tests/IO/MemoryMappedIOTest.cs b/test/AsmResolver.Tests/IO/MemoryMappedIOTest.cs
index 636f8d3d7..311efd1d8 100644
--- a/test/AsmResolver.Tests/IO/MemoryMappedIOTest.cs
+++ b/test/AsmResolver.Tests/IO/MemoryMappedIOTest.cs
@@ -1,3 +1,4 @@
+using System;
 using System.IO;
 using AsmResolver.IO;
 using AsmResolver.Tests.Runners;