Fix virtual memory allocation and access (#193)

* Fix virtual memory allocation and access

* Update test dependency lock file

src/SharpEmu.Core/Cpu/TrackedCpuMemory.cs

@@ -50,4 +50,9 @@ public sealed class TrackedCpuMemory : ICpuMemory, ITrackedCpuMemory, IGuestMemo
         address = 0;
         return false;
     }
+
+    public bool TryFreeGuestMemory(ulong address)
+    {
+        return _inner is IGuestMemoryAllocator allocator && allocator.TryFreeGuestMemory(address);
+    }
 }

src/SharpEmu.Core/Memory/PhysicalVirtualMemory.cs

@@ -42,7 +42,8 @@ public sealed unsafe class PhysicalVirtualMemory : IVirtualMemory, IGuestMemoryA
 
     private readonly IHostMemory _hostMemory;
     private ulong _guestAllocationArenaBase;
-    private ulong _guestAllocationOffset;
+    private readonly SortedDictionary<ulong, ulong> _guestAllocationFreeRanges = new();
+    private readonly Dictionary<ulong, (ulong Offset, ulong Size)> _guestAllocations = new();
     private static readonly ulong LazyReservePrimeBytes = ResolveLazyReservePrimeBytes();
 
     public PhysicalVirtualMemory(IHostMemory? hostMemory = null)
@@ -301,7 +302,9 @@ public sealed unsafe class PhysicalVirtualMemory : IVirtualMemory, IGuestMemoryA
                         GuestAllocationArenaSize,
                         executable: false,
                         allowAlternative: true);
-                    _guestAllocationOffset = GuestAllocationArenaStartOffset;
+                    _guestAllocationFreeRanges.Add(
+                        GuestAllocationArenaStartOffset,
+                        GuestAllocationArenaSize - GuestAllocationArenaStartOffset);
                 }
                 catch (Exception)
                 {
@@ -309,14 +312,89 @@ public sealed unsafe class PhysicalVirtualMemory : IVirtualMemory, IGuestMemoryA
                 }
             }
 
-            var alignedOffset = AlignUp(_guestAllocationOffset, alignment);
-            if (alignedOffset > GuestAllocationArenaSize || size > GuestAllocationArenaSize - alignedOffset)
+            ulong rangeOffset = 0;
+            ulong rangeSize = 0;
+            ulong alignedOffset = 0;
+            var found = false;
+            foreach (var range in _guestAllocationFreeRanges)
+            {
+                alignedOffset = AlignUp(range.Key, alignment);
+                if (alignedOffset >= range.Key &&
+                    alignedOffset - range.Key <= range.Value &&
+                    size <= range.Value - (alignedOffset - range.Key))
+                {
+                    rangeOffset = range.Key;
+                    rangeSize = range.Value;
+                    found = true;
+                    break;
+                }
+            }
+
+            if (!found)
             {
                 return false;
             }
 
+            _guestAllocationFreeRanges.Remove(rangeOffset);
+            if (alignedOffset > rangeOffset)
+            {
+                _guestAllocationFreeRanges.Add(rangeOffset, alignedOffset - rangeOffset);
+            }
+
+            var allocationEnd = alignedOffset + size;
+            var rangeEnd = rangeOffset + rangeSize;
+            if (allocationEnd < rangeEnd)
+            {
+                _guestAllocationFreeRanges.Add(allocationEnd, rangeEnd - allocationEnd);
+            }
+
             address = _guestAllocationArenaBase + alignedOffset;
-            _guestAllocationOffset = alignedOffset + size;
+            _guestAllocations.Add(address, (alignedOffset, size));
+            return true;
+        }
+    }
+
+    public bool TryFreeGuestMemory(ulong address)
+    {
+        lock (_guestAllocationGate)
+        {
+            if (!_guestAllocations.Remove(address, out var allocation))
+            {
+                return false;
+            }
+
+            var freeOffset = allocation.Offset;
+            var freeSize = allocation.Size;
+            ulong? previousOffset = null;
+            ulong? nextOffset = null;
+
+            foreach (var range in _guestAllocationFreeRanges)
+            {
+                if (range.Key < freeOffset)
+                {
+                    previousOffset = range.Key;
+                    continue;
+                }
+
+                nextOffset = range.Key;
+                break;
+            }
+
+            if (previousOffset is { } previous &&
+                previous + _guestAllocationFreeRanges[previous] == freeOffset)
+            {
+                freeOffset = previous;
+                freeSize += _guestAllocationFreeRanges[previous];
+                _guestAllocationFreeRanges.Remove(previous);
+            }
+
+            if (nextOffset is { } next && freeOffset + freeSize == next)
+            {
+                freeSize += _guestAllocationFreeRanges[next];
+                _guestAllocationFreeRanges.Remove(next);
+            }
+
+            _guestAllocationFreeRanges.Add(freeOffset, freeSize);
             return true;
         }
     }
@@ -380,7 +458,8 @@ public sealed unsafe class PhysicalVirtualMemory : IVirtualMemory, IGuestMemoryA
             }
 
             _guestAllocationArenaBase = 0;
-            _guestAllocationOffset = 0;
+            _guestAllocationFreeRanges.Clear();
+            _guestAllocations.Clear();
         }
     }
 
@@ -439,12 +518,33 @@ public sealed unsafe class PhysicalVirtualMemory : IVirtualMemory, IGuestMemoryA
 
     private void ApplySegmentProtection(ulong mapStart, ulong mapEnd, ProgramHeaderFlags flags)
     {
+        var runStart = mapStart;
+        var runFlags = ProgramHeaderFlags.None;
+        var hasRun = false;
+
         for (var pageAddress = mapStart; pageAddress < mapEnd; pageAddress += PageSize)
         {
             _pageProtections.TryGetValue(pageAddress, out var existingFlags);
             var mergedFlags = existingFlags | flags;
             _pageProtections[pageAddress] = mergedFlags;
-            SetProtection(pageAddress, PageSize, mergedFlags);
+
+            if (!hasRun)
+            {
+                runStart = pageAddress;
+                runFlags = mergedFlags;
+                hasRun = true;
+            }
+            else if (mergedFlags != runFlags)
+            {
+                SetProtection(runStart, pageAddress - runStart, runFlags);
+                runStart = pageAddress;
+                runFlags = mergedFlags;
+            }
+        }
+
+        if (hasRun)
+        {
+            SetProtection(runStart, mapEnd - runStart, runFlags);
         }
     }
 
@@ -789,9 +889,14 @@ public sealed unsafe class PhysicalVirtualMemory : IVirtualMemory, IGuestMemoryA
         _gate.EnterReadLock();
         try
         {
-            return FindRegion(virtualAddress, 1) is not null
-                ? (void*)virtualAddress
-                : null;
+            var region = FindRegion(virtualAddress, 1);
+            if (region is null ||
+                (region.IsReservedOnly && !EnsureRangeCommitted(virtualAddress, 1, region)))
+            {
+                return null;
+            }
+
+            return (void*)virtualAddress;
         }
         finally
         {

src/SharpEmu.Core/Memory/VirtualMemory.cs

@@ -41,15 +41,14 @@ public sealed class VirtualMemory : IVirtualMemory
 
         lock (_gate)
         {
-            foreach (var existing in _regions)
+            var insertionIndex = FindInsertionIndex(virtualAddress);
+            if ((insertionIndex > 0 && virtualAddress < _regions[insertionIndex - 1].EndAddress) ||
+                (insertionIndex < _regions.Count && endAddress > _regions[insertionIndex].Region.VirtualAddress))
             {
-                if (virtualAddress < existing.EndAddress && endAddress > existing.Region.VirtualAddress)
-                {
-                    throw new InvalidOperationException("Attempted to map an overlapping virtual memory region.");
-                }
+                throw new InvalidOperationException("Attempted to map an overlapping virtual memory region.");
             }
 
-            _regions.Add(new MappedRegion(
+            _regions.Insert(insertionIndex, new MappedRegion(
                 new VirtualMemoryRegion(virtualAddress, memorySize, fileOffset, (ulong)fileData.Length, protection),
                 endAddress,
                 backingMemory));
@@ -74,12 +73,12 @@ public sealed class VirtualMemory : IVirtualMemory
     {
         lock (_gate)
         {
-            if (!TryResolveRegion(virtualAddress, destination.Length, out var region, out var offset))
+            if (!TryValidateRange(virtualAddress, destination.Length, ProgramHeaderFlags.Read, out var regionIndex))
             {
                 return false;
             }
 
-            region.BackingMemory.AsSpan(offset, destination.Length).CopyTo(destination);
+            CopyFromRegions(virtualAddress, destination, regionIndex);
             return true;
         }
     }
@@ -88,39 +87,127 @@ public sealed class VirtualMemory : IVirtualMemory
     {
         lock (_gate)
         {
-            if (!TryResolveRegion(virtualAddress, source.Length, out var region, out var offset))
+            if (!TryValidateRange(virtualAddress, source.Length, ProgramHeaderFlags.Write, out var regionIndex))
             {
                 return false;
             }
 
-            source.CopyTo(region.BackingMemory.AsSpan(offset, source.Length));
+            CopyToRegions(virtualAddress, source, regionIndex);
             return true;
         }
     }
 
-    private bool TryResolveRegion(ulong virtualAddress, int length, out MappedRegion region, out int offset)
+    private bool TryValidateRange(
+        ulong virtualAddress,
+        int length,
+        ProgramHeaderFlags requiredProtection,
+        out int regionIndex)
     {
-        foreach (var candidate in _regions)
+        regionIndex = FindContainingRegionIndex(virtualAddress);
+        if (regionIndex < 0)
         {
-            if (virtualAddress < candidate.Region.VirtualAddress || virtualAddress >= candidate.EndAddress)
-            {
-                continue;
-            }
-
-            var candidateOffset = checked((int)(virtualAddress - candidate.Region.VirtualAddress));
-            if (candidateOffset + length > candidate.BackingMemory.Length)
-            {
-                break;
-            }
-
-            region = candidate;
-            offset = candidateOffset;
-            return true;
+            return false;
         }
 
-        region = default;
-        offset = 0;
-        return false;
+        var currentAddress = virtualAddress;
+        var remaining = length;
+        var currentIndex = regionIndex;
+        while (true)
+        {
+            if (currentIndex >= _regions.Count)
+            {
+                return false;
+            }
+
+            var region = _regions[currentIndex];
+            if (currentAddress < region.Region.VirtualAddress ||
+                currentAddress >= region.EndAddress ||
+                (region.Region.Protection & requiredProtection) == 0)
+            {
+                return false;
+            }
+
+            if (remaining == 0)
+            {
+                return true;
+            }
+
+            var available = region.EndAddress - currentAddress;
+            var chunkLength = (int)Math.Min((ulong)remaining, available);
+            remaining -= chunkLength;
+            if (remaining == 0)
+            {
+                return true;
+            }
+
+            currentAddress += (ulong)chunkLength;
+            currentIndex++;
+        }
+    }
+
+    private int FindContainingRegionIndex(ulong virtualAddress)
+    {
+        var insertionIndex = FindInsertionIndex(virtualAddress);
+        if (insertionIndex < _regions.Count &&
+            _regions[insertionIndex].Region.VirtualAddress == virtualAddress)
+        {
+            return insertionIndex;
+        }
+
+        var candidateIndex = insertionIndex - 1;
+        return candidateIndex >= 0 && virtualAddress < _regions[candidateIndex].EndAddress
+            ? candidateIndex
+            : -1;
+    }
+
+    private void CopyFromRegions(ulong virtualAddress, Span<byte> destination, int regionIndex)
+    {
+        var copied = 0;
+        var currentAddress = virtualAddress;
+        while (copied < destination.Length)
+        {
+            var region = _regions[regionIndex++];
+            var regionOffset = checked((int)(currentAddress - region.Region.VirtualAddress));
+            var chunkLength = Math.Min(destination.Length - copied, region.BackingMemory.Length - regionOffset);
+            region.BackingMemory.AsSpan(regionOffset, chunkLength).CopyTo(destination[copied..]);
+            copied += chunkLength;
+            currentAddress += (ulong)chunkLength;
+        }
+    }
+
+    private void CopyToRegions(ulong virtualAddress, ReadOnlySpan<byte> source, int regionIndex)
+    {
+        var copied = 0;
+        var currentAddress = virtualAddress;
+        while (copied < source.Length)
+        {
+            var region = _regions[regionIndex++];
+            var regionOffset = checked((int)(currentAddress - region.Region.VirtualAddress));
+            var chunkLength = Math.Min(source.Length - copied, region.BackingMemory.Length - regionOffset);
+            source.Slice(copied, chunkLength).CopyTo(region.BackingMemory.AsSpan(regionOffset, chunkLength));
+            copied += chunkLength;
+            currentAddress += (ulong)chunkLength;
+        }
+    }
+
+    private int FindInsertionIndex(ulong virtualAddress)
+    {
+        var lower = 0;
+        var upper = _regions.Count;
+        while (lower < upper)
+        {
+            var middle = lower + ((upper - lower) / 2);
+            if (_regions[middle].Region.VirtualAddress < virtualAddress)
+            {
+                lower = middle + 1;
+            }
+            else
+            {
+                upper = middle;
+            }
+        }
+
+        return lower;
     }
 
     private readonly record struct MappedRegion(VirtualMemoryRegion Region, ulong EndAddress, byte[] BackingMemory);