From 7fb834ffacbe46680dc7cf508ad4c327d573272c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Ha=CC=88user?=
 <8659494+mhaeuser@users.noreply.github.com>
Date: Sun, 26 Mar 2023 23:06:23 +0200
Subject: [PATCH] EmbeddedPkg/NonCoherentDmaLib: Utilize AllocateAlignedPagesEx

---
 .../NonCoherentDmaLib/NonCoherentDmaLib.c     | 83 +++----------------
 1 file changed, 11 insertions(+), 72 deletions(-)

diff --git a/EmbeddedPkg/Library/NonCoherentDmaLib/NonCoherentDmaLib.c b/EmbeddedPkg/Library/NonCoherentDmaLib/NonCoherentDmaLib.c
index 403cc8e2fd..bff4c99a9f 100644
--- a/EmbeddedPkg/Library/NonCoherentDmaLib/NonCoherentDmaLib.c
+++ b/EmbeddedPkg/Library/NonCoherentDmaLib/NonCoherentDmaLib.c
@@ -15,6 +15,7 @@
 #include <Library/DmaLib.h>
 #include <Library/DxeServicesTableLib.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Library/MemoryAllocationLibEx.h>
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/IoLib.h>
 #include <Library/BaseMemoryLib.h>
@@ -81,82 +82,20 @@ InternalAllocateAlignedPages (
 {
   EFI_STATUS            Status;
   EFI_PHYSICAL_ADDRESS  Memory;
-  UINTN                 AlignedMemory;
-  UINTN                 AlignmentMask;
-  UINTN                 UnalignedPages;
-  UINTN                 RealPages;
 
-  //
-  // Alignment must be a power of two or zero.
-  //
-  ASSERT ((Alignment & (Alignment - 1)) == 0);
-
-  if (Pages == 0) {
+  Memory = mDmaHostAddressLimit;
+  Status = AllocateAlignedPagesEx (
+             AllocateMaxAddress,
+             MemoryType,
+             Pages,
+             Alignment,
+             &Memory
+             );
+  if (EFI_ERROR (Status)) {
     return NULL;
   }
 
-  if (Alignment > EFI_PAGE_SIZE) {
-    //
-    // Calculate the total number of pages since alignment is larger than page
-    // size.
-    //
-    AlignmentMask = Alignment - 1;
-    RealPages     = Pages + EFI_SIZE_TO_PAGES (Alignment);
-    //
-    // Make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not
-    // overflow.
-    //
-    ASSERT (RealPages > Pages);
-
-    Memory = mDmaHostAddressLimit;
-    Status = gBS->AllocatePages (
-                    AllocateMaxAddress,
-                    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 = gBS->FreePages (Memory, UnalignedPages);
-      ASSERT_EFI_ERROR (Status);
-    }
-
-    Memory         = AlignedMemory + EFI_PAGES_TO_SIZE (Pages);
-    UnalignedPages = RealPages - Pages - UnalignedPages;
-    if (UnalignedPages > 0) {
-      //
-      // Free last unaligned page(s).
-      //
-      Status = gBS->FreePages (Memory, UnalignedPages);
-      ASSERT_EFI_ERROR (Status);
-    }
-  } else {
-    //
-    // Do not over-allocate pages in this case.
-    //
-    Memory = mDmaHostAddressLimit;
-    Status = gBS->AllocatePages (
-                    AllocateMaxAddress,
-                    MemoryType,
-                    Pages,
-                    &Memory
-                    );
-    if (EFI_ERROR (Status)) {
-      return NULL;
-    }
-
-    AlignedMemory = (UINTN)Memory;
-  }
-
-  return (VOID *)AlignedMemory;
+  return (VOID *)(UINTN)Memory;
 }
 
 /**