From 8a8e617dfe107a3962b0960cc6d22db8c5689cc9 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Mon, 14 Apr 2025 15:14:32 +0300
Subject: [PATCH] MdeModulePkg: Added PcdMaxMemoryType.

---
 .../Core/Dxe/DxeUserSpace/DxeUserSpace.inf     |  3 +++
 MdeModulePkg/Include/Library/MemoryPoolLib.h   | 18 +-----------------
 .../Library/MemoryPoolLib/MemoryPoolLib.inf    |  1 +
 MdeModulePkg/Library/MemoryPoolLib/Pool.c      |  2 ++
 MdeModulePkg/MdeModulePkg.dec                  |  6 ++++++
 5 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/DxeUserSpace/DxeUserSpace.inf b/MdeModulePkg/Core/Dxe/DxeUserSpace/DxeUserSpace.inf
index f008d37436..e359cdd05e 100644
--- a/MdeModulePkg/Core/Dxe/DxeUserSpace/DxeUserSpace.inf
+++ b/MdeModulePkg/Core/Dxe/DxeUserSpace/DxeUserSpace.inf
@@ -62,5 +62,8 @@
   gEfiDevicePathProtocolGuid           ## SOMETIMES_CONSUMES
   gEfiUnicodeCollationProtocolGuid     ## SOMETIMES_CONSUMES
 
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdMaxMemoryType  ## CONSUMES
+  
 [Depex]
   TRUE
diff --git a/MdeModulePkg/Include/Library/MemoryPoolLib.h b/MdeModulePkg/Include/Library/MemoryPoolLib.h
index bf459cdb13..3a0fb68f64 100644
--- a/MdeModulePkg/Include/Library/MemoryPoolLib.h
+++ b/MdeModulePkg/Include/Library/MemoryPoolLib.h
@@ -11,22 +11,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include <Library/UefiLib.h>
 
-//
-// +---------------------------------------------------+
-// | 0..(EfiMaxMemoryType - 1)    - Normal memory type |
-// +---------------------------------------------------+
-// | EfiMaxMemoryType..0x6FFFFFFF - Invalid            |
-// +---------------------------------------------------+
-// | 0x70000000..0x7FFFFFFF       - OEM reserved       |
-// +---------------------------------------------------+
-// | 0x80000000..0xFFFFFFFF       - OS reserved        |
-// +---------------------------------------------------+
-//
-#define MEMORY_TYPE_OS_RESERVED_MIN   0x80000000
-#define MEMORY_TYPE_OS_RESERVED_MAX   0xFFFFFFFF
-#define MEMORY_TYPE_OEM_RESERVED_MIN  0x70000000
-#define MEMORY_TYPE_OEM_RESERVED_MAX  0x7FFFFFFF
-
 //
 // Memory type to guard (matching the related PCD definition)
 //
@@ -36,7 +20,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #define GUARD_HEAP_TYPE_ALL         \
         (GUARD_HEAP_TYPE_PAGE|GUARD_HEAP_TYPE_POOL|GUARD_HEAP_TYPE_FREED)
 
-#define MAX_MEMORY_TYPE        100
+#define MAX_MEMORY_TYPE        FixedPcdGet32 (PcdMaxMemoryType)
 
 /**
   Called to initialize the pool.
diff --git a/MdeModulePkg/Library/MemoryPoolLib/MemoryPoolLib.inf b/MdeModulePkg/Library/MemoryPoolLib/MemoryPoolLib.inf
index 4943e6242b..b082b3d0d6 100644
--- a/MdeModulePkg/Library/MemoryPoolLib/MemoryPoolLib.inf
+++ b/MdeModulePkg/Library/MemoryPoolLib/MemoryPoolLib.inf
@@ -38,3 +38,4 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPageType                       ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPoolType                       ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPropertyMask                   ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdMaxMemoryType                           ## CONSUMES
diff --git a/MdeModulePkg/Library/MemoryPoolLib/Pool.c b/MdeModulePkg/Library/MemoryPoolLib/Pool.c
index 190ea17f55..45184d81b1 100644
--- a/MdeModulePkg/Library/MemoryPoolLib/Pool.c
+++ b/MdeModulePkg/Library/MemoryPoolLib/Pool.c
@@ -82,6 +82,8 @@ typedef struct {
 //
 POOL  mPoolHead[MAX_MEMORY_TYPE];
 
+STATIC_ASSERT (MAX_MEMORY_TYPE > EfiMaxMemoryType, "PcdMaxMemoryType must be greater than EfiMaxMemoryType");
+
 STATIC BOOLEAN  mIsCore;
 
 //
diff --git a/MdeModulePkg/MdeModulePkg.dec b/MdeModulePkg/MdeModulePkg.dec
index cb53d3ab3c..fcb7127cf9 100644
--- a/MdeModulePkg/MdeModulePkg.dec
+++ b/MdeModulePkg/MdeModulePkg.dec
@@ -1210,6 +1210,12 @@
   # @Prompt Enable User Space.
   gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUserSpace|FALSE|BOOLEAN|0x30001062
 
+  ## Sets the size of mPoolHead array in MemoryPoolLib. It must be greater than EfiMaxMemoryType.
+  #  In DxeMain this array manages Pools of different memory types.
+  #  In DxeUserSpace this array manages Pools in different User address spaces.
+  # @Prompt Max number of memory types (or User address spaces).
+  gEfiMdeModulePkgTokenSpaceGuid.PcdMaxMemoryType|100|UINT32|0x30001063
+
 [PcdsFixedAtBuild, PcdsPatchableInModule]
   ## Dynamic type PCD can be registered callback function for Pcd setting action.
   #  PcdMaxPeiPcdCallBackNumberPerPcdEntry indicates the maximum number of callback function