[PATCH 4/5] MdePkg PeiMemoryAllocationLib: Update InternalAllocateAlignedPages

[Star Zeng <star.zeng@intel.com>]

Update InternalAllocateAlignedPages to use PeiServicesFreePages.

Let the InternalAllocateAlignedPages in PeiMemoryAllocationLib use
same algorithm with InternalAllocateAlignedPages in
UefiMemoryAllocationLib.

Cc: Liming Gao <liming.gao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Star Zeng <star.zeng@intel.com>
---
 .../PeiMemoryAllocationLib/MemoryAllocationLib.c   | 172 +++++----------------
 1 file changed, 40 insertions(+), 132 deletions(-)

diff --git a/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c b/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c
index 37273dde561a..755f89c72ea2 100644
--- a/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c
+++ b/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c
@@ -181,15 +181,12 @@ InternalAllocateAlignedPages (
   IN UINTN            Alignment
   )
 {
-  EFI_PHYSICAL_ADDRESS   Memory;
-  EFI_PHYSICAL_ADDRESS   AlignedMemory;
-  EFI_PEI_HOB_POINTERS   Hob;
-  BOOLEAN                SkipBeforeMemHob;
-  BOOLEAN                SkipAfterMemHob;
-  EFI_PHYSICAL_ADDRESS   HobBaseAddress;
-  UINT64                 HobLength;
-  EFI_MEMORY_TYPE        HobMemoryType;
-  UINTN                  TotalPages;
+  EFI_STATUS            Status;
+  EFI_PHYSICAL_ADDRESS  Memory;
+  UINTN                 AlignedMemory;
+  UINTN                 AlignmentMask;
+  UINTN                 UnalignedPages;
+  UINTN                 RealPages;
 
   //
   // Alignment must be a power of two or zero.
@@ -199,139 +196,50 @@ InternalAllocateAlignedPages (
   if (Pages == 0) {
     return NULL;
   }
-  //
-  // Make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not overflow.
-  //
-  ASSERT (Pages <= (MAX_ADDRESS - EFI_SIZE_TO_PAGES (Alignment))); 
-
-  //
-  // We would rather waste some memory to save PEI code size.
-  // meaning in addition to the requested size for the aligned mem,
-  // we simply reserve an overhead memory equal to Alignmemt(page-aligned), no matter what.
-  // The overhead mem size could be reduced later with more involved malloc mechanisms
-  // (e.g., somthing that can detect the alignment boundary before allocating memory or 
-  //  can request that memory be allocated at a certain address that is aleady aligned).
-  //
-  TotalPages = Pages + (Alignment <= EFI_PAGE_SIZE ? 0 : EFI_SIZE_TO_PAGES(Alignment));
-  Memory = (EFI_PHYSICAL_ADDRESS) (UINTN) InternalAllocatePages (MemoryType, TotalPages);
-  if (Memory == 0) {
-    DEBUG((DEBUG_INFO, "Out of memory resource! \n"));
-    return NULL;
-  }
-  DEBUG ((DEBUG_INFO, "Allocated Memory unaligned: Address = 0x%LX, Pages = 0x%X, Type = %d \n", Memory, TotalPages, (UINTN) MemoryType));
-
-  //
-  // Alignment calculation
-  //
-  AlignedMemory = Memory;
   if (Alignment > EFI_PAGE_SIZE) {
-    AlignedMemory = ALIGN_VALUE (Memory, Alignment);
-  }
-  DEBUG ((DEBUG_INFO, "After aligning to 0x%X bytes: Address = 0x%LX, Pages = 0x%X \n", Alignment, AlignedMemory, Pages));
-
-  //
-  // In general three HOBs cover the total allocated space.
-  // The aligned portion is covered by the aligned mem HOB and
-  // the unaligned(to be freed) portions before and after the aligned portion are covered by newly created HOBs.
-  //
-  // Before mem HOB covers the region between "Memory" and "AlignedMemory"
-  // Aligned mem HOB covers the region between "AlignedMemory" and "AlignedMemory + EFI_PAGES_TO_SIZE(Pages)"
-  // After mem HOB covers the region between "AlignedMemory + EFI_PAGES_TO_SIZE(Pages)" and "Memory + EFI_PAGES_TO_SIZE(TotalPages)"
-  //
-  // The before or after mem HOBs need to be skipped under special cases where the aligned portion
-  // touches either the top or bottom of the original allocated space.
-  //
-  SkipBeforeMemHob = FALSE;
-  SkipAfterMemHob  = FALSE;
-  if (Memory == AlignedMemory) {
-    SkipBeforeMemHob = TRUE;
-  }
-  if ((Memory + EFI_PAGES_TO_SIZE(TotalPages)) == (AlignedMemory + EFI_PAGES_TO_SIZE(Pages))) {
     //
-    // This condition is never met in the current implementation.
-    // There is always some after-mem since the overhead mem(used in TotalPages)
-    // is no less than Alignment.
+    // Calculate the total number of pages since alignment is larger than page size.
     //
-    SkipAfterMemHob = TRUE;
-  }
-
-  //  
-  // Search for the mem HOB referring to the original(unaligned) allocation 
-  // and update the size and type if needed.
-  //
-  Hob.Raw = GetFirstHob (EFI_HOB_TYPE_MEMORY_ALLOCATION);
-  while (Hob.Raw != NULL) {
-    if (Hob.MemoryAllocation->AllocDescriptor.MemoryBaseAddress == Memory) {
-      break;
-    }
-    Hob.Raw = GET_NEXT_HOB (Hob);
-    Hob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, Hob.Raw);
-  }
-  ASSERT (Hob.Raw != NULL);
-  if (SkipBeforeMemHob) {
+    AlignmentMask  = Alignment - 1;
+    RealPages      = Pages + EFI_SIZE_TO_PAGES (Alignment);
     //
-    // Use this HOB as aligned mem HOB as there is no portion before it.
+    // Make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not overflow.
     //
-    HobLength = EFI_PAGES_TO_SIZE(Pages);
-    Hob.MemoryAllocation->AllocDescriptor.MemoryLength = HobLength;
+    ASSERT (RealPages > Pages);
+
+    Status         = PeiServicesAllocatePages (MemoryType, RealPages, &Memory);
+    if (EFI_ERROR (Status)) {
+      return NULL;
+    }
+    AlignedMemory  = ((UINTN) Memory + AlignmentMask) & ~AlignmentMask;
+    UnalignedPages = EFI_SIZE_TO_PAGES (AlignedMemory - (UINTN) Memory);
+    if (UnalignedPages > 0) {
+      //
+      // Free first unaligned page(s).
+      //
+      Status = PeiServicesFreePages (Memory, UnalignedPages);
+      ASSERT_EFI_ERROR (Status);
+    }
+    Memory         = (EFI_PHYSICAL_ADDRESS) (AlignedMemory + EFI_PAGES_TO_SIZE (Pages));
+    UnalignedPages = RealPages - Pages - UnalignedPages;
+    if (UnalignedPages > 0) {
+      //
+      // Free last unaligned page(s).
+      //
+      Status = PeiServicesFreePages (Memory, UnalignedPages);
+      ASSERT_EFI_ERROR (Status);
+    }
   } else {
     //
-    // Use this HOB as before mem HOB and create a new HOB for the aligned portion 
+    // Do not over-allocate pages in this case.
     //
-    HobLength = (AlignedMemory - Memory); 
-    Hob.MemoryAllocation->AllocDescriptor.MemoryLength = HobLength;
-    Hob.MemoryAllocation->AllocDescriptor.MemoryType = EfiConventionalMemory;
-  }
-
-  HobBaseAddress = Hob.MemoryAllocation->AllocDescriptor.MemoryBaseAddress;
-  HobMemoryType = Hob.MemoryAllocation->AllocDescriptor.MemoryType;
-
-  //
-  // Build the aligned mem HOB if needed
-  //
-  if (!SkipBeforeMemHob) {
-    DEBUG((DEBUG_INFO, "Updated before-mem HOB with BaseAddress = %LX, Length = %LX, MemoryType = %d \n",
-      HobBaseAddress, HobLength, (UINTN) HobMemoryType));
-
-    HobBaseAddress = AlignedMemory;
-    HobLength = EFI_PAGES_TO_SIZE(Pages);
-    HobMemoryType = MemoryType;
-
-    BuildMemoryAllocationHob (
-      HobBaseAddress,
-      HobLength,
-      HobMemoryType
-      );
-
-    DEBUG((DEBUG_INFO, "Created aligned-mem HOB with BaseAddress = %LX, Length = %LX, MemoryType = %d \n",
-      HobBaseAddress, HobLength, (UINTN) HobMemoryType));
-  } else {
-    if (HobBaseAddress != 0) {
-      DEBUG((DEBUG_INFO, "Updated aligned-mem HOB with BaseAddress = %LX, Length = %LX, MemoryType = %d \n",
-        HobBaseAddress, HobLength, (UINTN) HobMemoryType));
+    Status = PeiServicesAllocatePages (MemoryType, Pages, &Memory);
+    if (EFI_ERROR (Status)) {
+      return NULL;
     }
+    AlignedMemory  = (UINTN) Memory;
   }
-
-
-  //
-  // Build the after mem HOB if needed
-  //
-  if (!SkipAfterMemHob) {
-    HobBaseAddress = AlignedMemory + EFI_PAGES_TO_SIZE(Pages);
-    HobLength = (Memory + EFI_PAGES_TO_SIZE(TotalPages)) - (AlignedMemory + EFI_PAGES_TO_SIZE(Pages));
-    HobMemoryType = EfiConventionalMemory;
-
-    BuildMemoryAllocationHob (
-      HobBaseAddress,
-      HobLength,
-      HobMemoryType
-      );
-
-    DEBUG((DEBUG_INFO, "Created after-mem HOB with BaseAddress = %LX, Length = %LX, MemoryType = %d \n",
-      HobBaseAddress, HobLength, (UINTN) HobMemoryType));
-  }
-
-  return (VOID *) (UINTN) AlignedMemory;
+  return (VOID *) AlignedMemory;
 }
 
 /**
-- 
2.7.0.windows.1