Fix virtual memory allocation and access (#193)

* Fix virtual memory allocation and access * Update test dependency lock file
2026-07-24 11:48:39 +08:00 · 2026-07-14 21:50:54 -06:00
parent 373100a6b0
commit 9d88542efd
9 changed files with 626 additions and 40 deletions
@@ -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
        {