Move page chunking to PageAllocator

BepInEx.IL2CPP/Hook/Allocator/LinuxPageAllocator.cs

@@ -23,21 +23,24 @@ namespace BepInEx.IL2CPP.Allocator
 				throw new NotImplementedException();
 			}
 		}
-		
+
 		static class Unix
 		{
 			public delegate IntPtr mmapDelegate(IntPtr addr, UIntPtr length, int prot, int flags, int fd, int offset);
+
 			[DynDllImport("mmap")]
 			public static mmapDelegate mmap;
 
 			public delegate int munmapDelegate(IntPtr addr, UIntPtr length);
+
 			[DynDllImport("munmap")]
 			public static munmapDelegate munmap;
-			
+
 			public delegate IntPtr fdopenDelegate(int fd, string mode);
+
 			[DynDllImport("libc")]
 			public static fdopenDelegate fdopen;
-			
+
 			static Unix()
 			{
 				typeof(Unix).ResolveDynDllImports(new Dictionary<string, List<DynDllMapping>>

BepInEx.IL2CPP/Hook/Allocator/MacOsPageAllocator.cs

@@ -8,7 +8,7 @@ namespace BepInEx.IL2CPP.Allocator
 		{
 			throw new System.NotImplementedException();
 		}
-		
+
 		protected class MacOsMemoryMapper : IMemoryMapper
 		{
 			public void Dispose()

BepInEx.IL2CPP/Hook/Allocator/PageAllocator.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using MonoMod.Utils;
 
@@ -6,9 +7,9 @@ namespace BepInEx.IL2CPP.Allocator
 {
 	public class PageAllocatorException : Exception
 	{
-		public PageAllocatorException(string message) : base(message) {}
+		public PageAllocatorException(string message) : base(message) { }
 	}
-	
+
 	/// <summary>
 	///     A general purpose page allocator for patching purposes.
 	///     Allows to allocate pages (4k memory chunks) within the 1GB radius of a given address.
@@ -37,18 +38,64 @@ namespace BepInEx.IL2CPP.Allocator
 		public static PageAllocator Instance => instance ??= Init();
 
 		/// <summary>
+		///     Allocates a single 64k chunk of memory near the given address
+		/// </summary>
+		/// <param name="hint">Address near which to attempt allocate the chunk</param>
+		/// <returns>Allocated chunk</returns>
+		/// <exception cref="PageAllocatorException">Allocation failed</exception>
+		protected abstract IntPtr AllocateChunk(IntPtr hint);
+
+		/// <summary>
 		///     Allocates a single page of size <see cref="PAGE_SIZE" /> near the provided address.
 		///     Attempts to allocate the page within the +-1GB region of the hinted address.
 		/// </summary>
 		/// <param name="hint">Address near which to attempt to allocate the page.</param>
 		/// <returns>Address to the allocated page.</returns>
-		public abstract IntPtr Allocate(IntPtr hint);
+		public virtual IntPtr Allocate(IntPtr hint)
+		{
+			foreach (var allocatedChunk in allocatedChunks)
+			{
+				// Small shortcut to speed up page lookup
+				if (allocatedChunk.UsedPages == PAGES_PER_UNIT)
+					continue;
+				for (var i = 0; i < allocatedChunk.Pages.Length; i++)
+				{
+					if (allocatedChunk.Pages[i])
+						continue;
+					var pageAddr = allocatedChunk.GetPage(i);
+					if (!IsInRelJmpRange(hint, pageAddr))
+						continue;
+					allocatedChunk.Pages[i] = true;
+					allocatedChunk.UsedPages++;
+					return pageAddr;
+				}
+			}
+
+			var chunk = new PageChunk
+			{
+				BaseAddress = AllocateChunk(hint)
+			};
+			allocatedChunks.Add(chunk);
+			chunk.Pages[0] = true;
+			chunk.UsedPages++;
+			return chunk.BaseAddress;
+		}
 
 		/// <summary>
 		///     Frees the page allocated with <see cref="Allocate" />
 		/// </summary>
 		/// <param name="page"></param>
-		public abstract void Free(IntPtr page);
+		public void Free(IntPtr page)
+		{
+			foreach (var allocatedChunk in allocatedChunks)
+			{
+				long index = (page.ToInt64() - allocatedChunk.BaseAddress.ToInt64()) / PAGE_SIZE;
+				if (index < 0 || index > PAGES_PER_UNIT)
+					continue;
+				allocatedChunk.Pages[index] = false;
+				return;
+			}
+		}
 
 		[MethodImpl(MethodImplOptions.AggressiveInlining)]
 		protected static long RoundUp(long num, long unit)
@@ -77,10 +124,24 @@ namespace BepInEx.IL2CPP.Allocator
 				var v when v.Is(Platform.MacOS)   => new MacOsPageAllocator(),
 				_                                 => throw new NotImplementedException()
 			};
+
+		private class PageChunk
+		{
+			public readonly bool[] Pages = new bool[PAGES_PER_UNIT];
+			public IntPtr BaseAddress;
+			public int UsedPages;
+
+			public IntPtr GetPage(int index)
+			{
+				return BaseAddress + index * PAGE_SIZE;
+			}
+		}
+
+		private readonly List<PageChunk> allocatedChunks = new List<PageChunk>();
 	}
-	
+
 	internal static class PlatformExt
 	{
 		public static bool Is(this Platform pl, Platform val) => (pl & val) == val;
 	}
-}
+}

BepInEx.IL2CPP/Hook/Allocator/UnixPageAllocator.cs

@@ -21,9 +21,10 @@ namespace BepInEx.IL2CPP.Allocator
 				if (hint.ToInt64() - addr.ToInt64() < int.MaxValue)
 					result[0] = addr;
 			}
+
 			return false;
 		}
-		
+
 		[MethodImpl(MethodImplOptions.AggressiveInlining)]
 		private static bool CheckFreeRegionAfter(IntPtr end, IntPtr hint, IntPtr[] result)
 		{
@@ -33,15 +34,16 @@ namespace BepInEx.IL2CPP.Allocator
 					result[1] = end;
 				return true;
 			}
+
 			return false;
 		}
-			
+
 		private IntPtr[] GetFreeAddresses(IntPtr hint)
 		{
 			var result = new IntPtr[2];
 			var prevEnd = IntPtr.Zero;
 			using var mapper = OpenMemoryMap();
-			
+
 			while (mapper.FindNextFree(out var start, out var end))
 			{
 				if ((prevEnd + PAGE_SIZE).ToInt64() <= start.ToInt64())
@@ -49,24 +51,20 @@ namespace BepInEx.IL2CPP.Allocator
 						return result;
 				prevEnd = end;
 			}
+
 			if (CheckFreeRegionAfter(prevEnd, hint, result))
 				return result;
 			throw new PageAllocatorException($"Could not find free region near {hint.ToInt64():X8}");
 		}
-		
-		public override IntPtr Allocate(IntPtr hint)
-		{
-			throw new NotImplementedException();
-		}
 
-		public override void Free(IntPtr page)
+		protected override IntPtr AllocateChunk(IntPtr hint)
 		{
 			throw new NotImplementedException();
 		}
 
-		protected interface IMemoryMapper: IDisposable
+		protected interface IMemoryMapper : IDisposable
 		{
 			bool FindNextFree(out IntPtr start, out IntPtr end);
 		}
 	}
-}
+}

BepInEx.IL2CPP/Hook/Allocator/WindowsPageAllocator.cs

@@ -10,15 +10,7 @@ namespace BepInEx.IL2CPP.Allocator
 	/// </summary>
 	internal class WindowsPageAllocator : PageAllocator
 	{
-		private readonly List<PageChunk> allocatedChunks = new List<PageChunk>();
-
-		/// <summary>
-		///     Allocates a single 64k chunk of memory near the given address
-		/// </summary>
-		/// <param name="hint">Address near which to attempt allocate the chunk</param>
-		/// <returns>Allocated chunk</returns>
-		/// <exception cref="Win32Exception">Allocation failed</exception>
-		private PageChunk AllocateChunk(IntPtr hint)
+		protected override IntPtr AllocateChunk(IntPtr hint)
 		{
 			while (true)
 			{
@@ -51,72 +43,17 @@ namespace BepInEx.IL2CPP.Allocator
 				throw new Win32Exception(error);
 			}
 
-			var result = new PageChunk
-			{
-				BaseAddress = chunk
-			};
-			allocatedChunks.Add(result);
-			return result;
-		}
-
-		private IntPtr AllocPage(IntPtr hint)
-		{
-			foreach (var allocatedChunk in allocatedChunks)
-			{
-				// Small shortcut to speed up page lookup
-				if (allocatedChunk.UsedPages == PAGES_PER_UNIT)
-					continue;
-				for (var i = 0; i < allocatedChunk.Pages.Length; i++)
-				{
-					if (allocatedChunk.Pages[i])
-						continue;
-					var pageAddr = allocatedChunk.GetPage(i);
-					if (!IsInRelJmpRange(hint, pageAddr))
-						continue;
-					allocatedChunk.Pages[i] = true;
-					allocatedChunk.UsedPages++;
-					return pageAddr;
-				}
-			}
-
-			var chunk = AllocateChunk(hint);
-			chunk.Pages[0] = true;
-			chunk.UsedPages++;
-			return chunk.BaseAddress;
+			return chunk;
 		}
 
 		public override IntPtr Allocate(IntPtr hint)
 		{
-			var pageAddress = AllocPage(hint);
+			var pageAddress = base.Allocate(hint);
 			if (WinApi.VirtualAlloc(pageAddress, (UIntPtr)PAGE_SIZE, WinApi.AllocationType.MEM_COMMIT, WinApi.ProtectConstant.PAGE_READWRITE) == IntPtr.Zero)
 				throw new Win32Exception(Marshal.GetLastWin32Error());
 			return pageAddress;
 		}
 
-		public override void Free(IntPtr page)
-		{
-			foreach (var allocatedChunk in allocatedChunks)
-			{
-				long index = (page.ToInt64() - allocatedChunk.BaseAddress.ToInt64()) / PAGE_SIZE;
-				if (index < 0 || index > PAGES_PER_UNIT)
-					continue;
-				allocatedChunk.Pages[index] = false;
-				return;
-			}
-		}
-
-		private class PageChunk
-		{
-			public readonly bool[] Pages = new bool[PAGES_PER_UNIT];
-			public IntPtr BaseAddress;
-			public int UsedPages;
-
-			public IntPtr GetPage(int index)
-			{
-				return BaseAddress + index * PAGE_SIZE;
-			}
-		}
-
 		private static class WinApi
 		{
 			[Flags]
@@ -178,4 +115,4 @@ namespace BepInEx.IL2CPP.Allocator
 			}
 		}
 	}
-}
+}