With the new ImagePropertiesRecordLib, UEFI and SMM no longer need
to carry duplicate image properties record logic in their respective
MAT implementations.

Signed-off-by: Taylor Beebe <t...@taylorbeebe.com>
Cc: Jian J Wang <jian.j.w...@intel.com>
Cc: Liming Gao <gaolim...@byosoft.com.cn>
Cc: Dandan Bi <dandan...@intel.com>
---
 MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c  | 784 +------------------
 MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c       |  24 +-
 MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c | 785 +-------------------
 MdeModulePkg/Core/Dxe/DxeMain.h                     |  20 -
 MdeModulePkg/Core/Dxe/DxeMain.inf                   |   1 +
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf           |   1 +
 6 files changed, 45 insertions(+), 1570 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c 
b/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c
index 6e15a76cccae..53aedf7d671e 100644
--- a/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c
+++ b/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c
@@ -14,6 +14,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/DxeServicesTableLib.h>
 #include <Library/DebugLib.h>
 #include <Library/UefiLib.h>
+#include <Library/ImagePropertiesRecordLib.h>
 
 #include <Guid/EventGroup.h>
 
@@ -333,45 +334,6 @@ CoreInitializeMemoryAttributesTable (
 // Below functions are for MemoryMap
 //
 
-/**
-  Converts a number of EFI_PAGEs to a size in bytes.
-
-  NOTE: Do not use EFI_PAGES_TO_SIZE because it handles UINTN only.
-
-  @param  Pages     The number of EFI_PAGES.
-
-  @return  The number of bytes associated with the number of EFI_PAGEs 
specified
-           by Pages.
-**/
-STATIC
-UINT64
-EfiPagesToSize (
-  IN UINT64  Pages
-  )
-{
-  return LShiftU64 (Pages, EFI_PAGE_SHIFT);
-}
-
-/**
-  Converts a size, in bytes, to a number of EFI_PAGESs.
-
-  NOTE: Do not use EFI_SIZE_TO_PAGES because it handles UINTN only.
-
-  @param  Size      A size in bytes.
-
-  @return  The number of EFI_PAGESs associated with the number of bytes 
specified
-           by Size.
-
-**/
-STATIC
-UINT64
-EfiSizeToPages (
-  IN UINT64  Size
-  )
-{
-  return RShiftU64 (Size, EFI_PAGE_SHIFT) + ((((UINTN)Size) & EFI_PAGE_MASK) ? 
1 : 0);
-}
-
 /**
   Acquire memory lock on mMemoryAttributesTableLock.
 **/
@@ -396,48 +358,6 @@ CoreReleasemMemoryAttributesTableLock (
   CoreReleaseLock (&mMemoryAttributesTableLock);
 }
 
-/**
-  Sort memory map entries based upon PhysicalStart, from low to high.
-
-  @param  MemoryMap              A pointer to the buffer in which firmware 
places
-                                 the current memory map.
-  @param  MemoryMapSize          Size, in bytes, of the MemoryMap buffer.
-  @param  DescriptorSize         Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-**/
-STATIC
-VOID
-SortMemoryMap (
-  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
-  IN UINTN                      MemoryMapSize,
-  IN UINTN                      DescriptorSize
-  )
-{
-  EFI_MEMORY_DESCRIPTOR  *MemoryMapEntry;
-  EFI_MEMORY_DESCRIPTOR  *NextMemoryMapEntry;
-  EFI_MEMORY_DESCRIPTOR  *MemoryMapEnd;
-  EFI_MEMORY_DESCRIPTOR  TempMemoryMap;
-
-  MemoryMapEntry     = MemoryMap;
-  NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
-  MemoryMapEnd       = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + 
MemoryMapSize);
-  while (MemoryMapEntry < MemoryMapEnd) {
-    while (NextMemoryMapEntry < MemoryMapEnd) {
-      if (MemoryMapEntry->PhysicalStart > NextMemoryMapEntry->PhysicalStart) {
-        CopyMem (&TempMemoryMap, MemoryMapEntry, sizeof 
(EFI_MEMORY_DESCRIPTOR));
-        CopyMem (MemoryMapEntry, NextMemoryMapEntry, sizeof 
(EFI_MEMORY_DESCRIPTOR));
-        CopyMem (NextMemoryMapEntry, &TempMemoryMap, sizeof 
(EFI_MEMORY_DESCRIPTOR));
-      }
-
-      NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NextMemoryMapEntry, 
DescriptorSize);
-    }
-
-    MemoryMapEntry     = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, 
DescriptorSize);
-    NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, 
DescriptorSize);
-  }
-
-  return;
-}
-
 /**
   Merge continous memory map entries whose have same attributes.
 
@@ -471,7 +391,7 @@ MergeMemoryMap (
 
     do {
       MergeGuardPages (NewMemoryMapEntry, NextMemoryMapEntry->PhysicalStart);
-      MemoryBlockLength = (UINT64)(EfiPagesToSize 
(NewMemoryMapEntry->NumberOfPages));
+      MemoryBlockLength = LShiftU64 (NewMemoryMapEntry->NumberOfPages, 
EFI_PAGE_SHIFT);
       if (((UINTN)NextMemoryMapEntry < (UINTN)MemoryMapEnd) &&
           (NewMemoryMapEntry->Type == NextMemoryMapEntry->Type) &&
           (NewMemoryMapEntry->Attribute == NextMemoryMapEntry->Attribute) &&
@@ -538,426 +458,6 @@ EnforceMemoryMapAttribute (
   return;
 }
 
-/**
-  Return the first image record, whose [ImageBase, ImageSize] covered by 
[Buffer, Length].
-
-  @param Buffer  Start Address
-  @param Length  Address length
-
-  @return first image record covered by [buffer, length]
-**/
-STATIC
-IMAGE_PROPERTIES_RECORD *
-GetImageRecordByAddress (
-  IN EFI_PHYSICAL_ADDRESS  Buffer,
-  IN UINT64                Length
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  LIST_ENTRY               *ImageRecordLink;
-  LIST_ENTRY               *ImageRecordList;
-
-  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
-
-  for (ImageRecordLink = ImageRecordList->ForwardLink;
-       ImageRecordLink != ImageRecordList;
-       ImageRecordLink = ImageRecordLink->ForwardLink)
-  {
-    ImageRecord = CR (
-                    ImageRecordLink,
-                    IMAGE_PROPERTIES_RECORD,
-                    Link,
-                    IMAGE_PROPERTIES_RECORD_SIGNATURE
-                    );
-
-    if ((Buffer <= ImageRecord->ImageBase) &&
-        (Buffer + Length >= ImageRecord->ImageBase + ImageRecord->ImageSize))
-    {
-      return ImageRecord;
-    }
-  }
-
-  return NULL;
-}
-
-/**
-  Set the memory map to new entries, according to one old entry,
-  based upon PE code section and data section in image record
-
-  @param  ImageRecord            An image record whose [ImageBase, ImageSize] 
covered
-                                 by old memory map entry.
-  @param  NewRecord              A pointer to several new memory map entries.
-                                 The caller gurantee the buffer size be 1 +
-                                 (SplitRecordCount * DescriptorSize) calculated
-                                 below.
-  @param  OldRecord              A pointer to one old memory map entry.
-  @param  DescriptorSize         Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-**/
-STATIC
-UINTN
-SetNewRecord (
-  IN IMAGE_PROPERTIES_RECORD    *ImageRecord,
-  IN OUT EFI_MEMORY_DESCRIPTOR  *NewRecord,
-  IN EFI_MEMORY_DESCRIPTOR      *OldRecord,
-  IN UINTN                      DescriptorSize
-  )
-{
-  EFI_MEMORY_DESCRIPTOR                 TempRecord;
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
-  LIST_ENTRY                            *ImageRecordCodeSectionLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionEndLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionList;
-  UINTN                                 NewRecordCount;
-  UINT64                                PhysicalEnd;
-  UINT64                                ImageEnd;
-
-  CopyMem (&TempRecord, OldRecord, sizeof (EFI_MEMORY_DESCRIPTOR));
-  PhysicalEnd    = TempRecord.PhysicalStart + EfiPagesToSize 
(TempRecord.NumberOfPages);
-  NewRecordCount = 0;
-
-  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
-
-  ImageRecordCodeSectionLink    = ImageRecordCodeSectionList->ForwardLink;
-  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
-  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-    ImageRecordCodeSection = CR (
-                               ImageRecordCodeSectionLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-
-    if (TempRecord.PhysicalStart <= ImageRecordCodeSection->CodeSegmentBase) {
-      //
-      // DATA
-      //
-      NewRecord->Type          = TempRecord.Type;
-      NewRecord->PhysicalStart = TempRecord.PhysicalStart;
-      NewRecord->VirtualStart  = 0;
-      NewRecord->NumberOfPages = EfiSizeToPages 
(ImageRecordCodeSection->CodeSegmentBase - NewRecord->PhysicalStart);
-      NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
-      if (NewRecord->NumberOfPages != 0) {
-        NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
-        NewRecordCount++;
-      }
-
-      //
-      // CODE
-      //
-      NewRecord->Type          = TempRecord.Type;
-      NewRecord->PhysicalStart = ImageRecordCodeSection->CodeSegmentBase;
-      NewRecord->VirtualStart  = 0;
-      NewRecord->NumberOfPages = EfiSizeToPages 
(ImageRecordCodeSection->CodeSegmentSize);
-      NewRecord->Attribute     = (TempRecord.Attribute & (~EFI_MEMORY_XP)) | 
EFI_MEMORY_RO;
-      if (NewRecord->NumberOfPages != 0) {
-        NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
-        NewRecordCount++;
-      }
-
-      TempRecord.PhysicalStart = ImageRecordCodeSection->CodeSegmentBase + 
EfiPagesToSize (EfiSizeToPages (ImageRecordCodeSection->CodeSegmentSize));
-      TempRecord.NumberOfPages = EfiSizeToPages (PhysicalEnd - 
TempRecord.PhysicalStart);
-      if (TempRecord.NumberOfPages == 0) {
-        break;
-      }
-    }
-  }
-
-  ImageEnd = ImageRecord->ImageBase + ImageRecord->ImageSize;
-
-  //
-  // Final DATA
-  //
-  if (TempRecord.PhysicalStart < ImageEnd) {
-    NewRecord->Type          = TempRecord.Type;
-    NewRecord->PhysicalStart = TempRecord.PhysicalStart;
-    NewRecord->VirtualStart  = 0;
-    NewRecord->NumberOfPages = EfiSizeToPages (ImageEnd - 
TempRecord.PhysicalStart);
-    NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
-    NewRecordCount++;
-  }
-
-  return NewRecordCount;
-}
-
-/**
-  Return the max number of new splitted entries, according to one old entry,
-  based upon PE code section and data section.
-
-  @param  OldRecord              A pointer to one old memory map entry.
-
-  @retval  0 no entry need to be splitted.
-  @return  the max number of new splitted entries
-**/
-STATIC
-UINTN
-GetMaxSplitRecordCount (
-  IN EFI_MEMORY_DESCRIPTOR  *OldRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  UINTN                    SplitRecordCount;
-  UINT64                   PhysicalStart;
-  UINT64                   PhysicalEnd;
-
-  SplitRecordCount = 0;
-  PhysicalStart    = OldRecord->PhysicalStart;
-  PhysicalEnd      = OldRecord->PhysicalStart + EfiPagesToSize 
(OldRecord->NumberOfPages);
-
-  do {
-    ImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - 
PhysicalStart);
-    if (ImageRecord == NULL) {
-      break;
-    }
-
-    SplitRecordCount += (2 * ImageRecord->CodeSegmentCount + 1);
-    PhysicalStart     = ImageRecord->ImageBase + ImageRecord->ImageSize;
-  } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));
-
-  if (SplitRecordCount != 0) {
-    SplitRecordCount--;
-  }
-
-  return SplitRecordCount;
-}
-
-/**
-  Split the memory map to new entries, according to one old entry,
-  based upon PE code section and data section.
-
-  @param  OldRecord              A pointer to one old memory map entry.
-  @param  NewRecord              A pointer to several new memory map entries.
-                                 The caller gurantee the buffer size be 1 +
-                                 (SplitRecordCount * DescriptorSize) calculated
-                                 below.
-  @param  MaxSplitRecordCount    The max number of splitted entries
-  @param  DescriptorSize         Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-
-  @retval  0 no entry is splitted.
-  @return  the real number of splitted record.
-**/
-STATIC
-UINTN
-SplitRecord (
-  IN EFI_MEMORY_DESCRIPTOR      *OldRecord,
-  IN OUT EFI_MEMORY_DESCRIPTOR  *NewRecord,
-  IN UINTN                      MaxSplitRecordCount,
-  IN UINTN                      DescriptorSize
-  )
-{
-  EFI_MEMORY_DESCRIPTOR    TempRecord;
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  IMAGE_PROPERTIES_RECORD  *NewImageRecord;
-  UINT64                   PhysicalStart;
-  UINT64                   PhysicalEnd;
-  UINTN                    NewRecordCount;
-  UINTN                    TotalNewRecordCount;
-  BOOLEAN                  IsLastRecordData;
-
-  if (MaxSplitRecordCount == 0) {
-    CopyMem (NewRecord, OldRecord, DescriptorSize);
-    return 0;
-  }
-
-  TotalNewRecordCount = 0;
-
-  //
-  // Override previous record
-  //
-  CopyMem (&TempRecord, OldRecord, sizeof (EFI_MEMORY_DESCRIPTOR));
-  PhysicalStart = TempRecord.PhysicalStart;
-  PhysicalEnd   = TempRecord.PhysicalStart + EfiPagesToSize 
(TempRecord.NumberOfPages);
-
-  ImageRecord = NULL;
-  do {
-    NewImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - 
PhysicalStart);
-    if (NewImageRecord == NULL) {
-      //
-      // No more image covered by this range, stop
-      //
-      if ((PhysicalEnd > PhysicalStart) && (ImageRecord != NULL)) {
-        //
-        // If this is still address in this record, need record.
-        //
-        NewRecord        = PREVIOUS_MEMORY_DESCRIPTOR (NewRecord, 
DescriptorSize);
-        IsLastRecordData = FALSE;
-        if ((NewRecord->Attribute & EFI_MEMORY_XP) != 0) {
-          IsLastRecordData = TRUE;
-        }
-
-        if (IsLastRecordData) {
-          //
-          // Last record is DATA, just merge it.
-          //
-          NewRecord->NumberOfPages = EfiSizeToPages (PhysicalEnd - 
NewRecord->PhysicalStart);
-        } else {
-          //
-          // Last record is CODE, create a new DATA entry.
-          //
-          NewRecord                = NEXT_MEMORY_DESCRIPTOR (NewRecord, 
DescriptorSize);
-          NewRecord->Type          = TempRecord.Type;
-          NewRecord->PhysicalStart = TempRecord.PhysicalStart;
-          NewRecord->VirtualStart  = 0;
-          NewRecord->NumberOfPages = TempRecord.NumberOfPages;
-          NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
-          TotalNewRecordCount++;
-        }
-      }
-
-      break;
-    }
-
-    ImageRecord = NewImageRecord;
-
-    //
-    // Set new record
-    //
-    NewRecordCount       = SetNewRecord (ImageRecord, NewRecord, &TempRecord, 
DescriptorSize);
-    TotalNewRecordCount += NewRecordCount;
-    NewRecord            = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)NewRecord + 
NewRecordCount * DescriptorSize);
-
-    //
-    // Update PhysicalStart, in order to exclude the image buffer already 
splitted.
-    //
-    PhysicalStart            = ImageRecord->ImageBase + ImageRecord->ImageSize;
-    TempRecord.PhysicalStart = PhysicalStart;
-    TempRecord.NumberOfPages = EfiSizeToPages (PhysicalEnd - PhysicalStart);
-  } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));
-
-  //
-  // The logic in function SplitTable() ensures that TotalNewRecordCount will 
not be zero if the
-  // code reaches here.
-  //
-  ASSERT (TotalNewRecordCount != 0);
-  return TotalNewRecordCount - 1;
-}
-
-/**
-  Split the original memory map, and add more entries to describe PE code 
section and data section.
-  This function will set EfiRuntimeServicesData to be EFI_MEMORY_XP.
-  This function will merge entries with same attributes finally.
-
-  NOTE: It assumes PE code/data section are page aligned.
-  NOTE: It assumes enough entry is prepared for new memory map.
-
-  Split table:
-   +---------------+
-   | Record X      |
-   +---------------+
-   | Record RtCode |
-   +---------------+
-   | Record Y      |
-   +---------------+
-   ==>
-   +---------------+
-   | Record X      |
-   +---------------+ ----
-   | Record RtData |     |
-   +---------------+     |
-   | Record RtCode |     |-> PE/COFF1
-   +---------------+     |
-   | Record RtData |     |
-   +---------------+ ----
-   | Record RtData |     |
-   +---------------+     |
-   | Record RtCode |     |-> PE/COFF2
-   +---------------+     |
-   | Record RtData |     |
-   +---------------+ ----
-   | Record Y      |
-   +---------------+
-
-  @param  MemoryMapSize          A pointer to the size, in bytes, of the
-                                 MemoryMap buffer. On input, this is the size 
of
-                                 old MemoryMap before split. The actual buffer
-                                 size of MemoryMap is MemoryMapSize +
-                                 (AdditionalRecordCount * DescriptorSize) 
calculated
-                                 below. On output, it is the size of new 
MemoryMap
-                                 after split.
-  @param  MemoryMap              A pointer to the buffer in which firmware 
places
-                                 the current memory map.
-  @param  DescriptorSize         Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-**/
-STATIC
-VOID
-SplitTable (
-  IN OUT UINTN                  *MemoryMapSize,
-  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
-  IN UINTN                      DescriptorSize
-  )
-{
-  INTN   IndexOld;
-  INTN   IndexNew;
-  UINTN  MaxSplitRecordCount;
-  UINTN  RealSplitRecordCount;
-  UINTN  TotalSplitRecordCount;
-  UINTN  AdditionalRecordCount;
-
-  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax 
+ 1) * mImagePropertiesPrivateData.ImageRecordCount;
-
-  TotalSplitRecordCount = 0;
-  //
-  // Let old record point to end of valid MemoryMap buffer.
-  //
-  IndexOld = ((*MemoryMapSize) / DescriptorSize) - 1;
-  //
-  // Let new record point to end of full MemoryMap buffer.
-  //
-  IndexNew = ((*MemoryMapSize) / DescriptorSize) - 1 + AdditionalRecordCount;
-  for ( ; IndexOld >= 0; IndexOld--) {
-    MaxSplitRecordCount = GetMaxSplitRecordCount ((EFI_MEMORY_DESCRIPTOR 
*)((UINT8 *)MemoryMap + IndexOld * DescriptorSize));
-    //
-    // Split this MemoryMap record
-    //
-    IndexNew            -= MaxSplitRecordCount;
-    RealSplitRecordCount = SplitRecord (
-                             (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + 
IndexOld * DescriptorSize),
-                             (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + 
IndexNew * DescriptorSize),
-                             MaxSplitRecordCount,
-                             DescriptorSize
-                             );
-    //
-    // Adjust IndexNew according to real split.
-    //
-    CopyMem (
-      ((UINT8 *)MemoryMap + (IndexNew + MaxSplitRecordCount - 
RealSplitRecordCount) * DescriptorSize),
-      ((UINT8 *)MemoryMap + IndexNew * DescriptorSize),
-      RealSplitRecordCount * DescriptorSize
-      );
-    IndexNew               = IndexNew + MaxSplitRecordCount - 
RealSplitRecordCount;
-    TotalSplitRecordCount += RealSplitRecordCount;
-    IndexNew--;
-  }
-
-  //
-  // Move all records to the beginning.
-  //
-  CopyMem (
-    MemoryMap,
-    (UINT8 *)MemoryMap + (AdditionalRecordCount - TotalSplitRecordCount) * 
DescriptorSize,
-    (*MemoryMapSize) + TotalSplitRecordCount * DescriptorSize
-    );
-
-  *MemoryMapSize = (*MemoryMapSize) + DescriptorSize * TotalSplitRecordCount;
-
-  //
-  // Sort from low to high (Just in case)
-  //
-  SortMemoryMap (MemoryMap, *MemoryMapSize, DescriptorSize);
-
-  //
-  // Set RuntimeData to XP
-  //
-  EnforceMemoryMapAttribute (MemoryMap, *MemoryMapSize, DescriptorSize);
-
-  //
-  // Merge same type to save entry size
-  //
-  MergeMemoryMap (MemoryMap, MemoryMapSize, DescriptorSize);
-
-  return;
-}
-
 /**
   This function for GetMemoryMap() with properties table capability.
 
@@ -1035,7 +535,23 @@ CoreGetMemoryMapWithSeparatedImageSection (
       //
       // Split PE code/data
       //
-      SplitTable (MemoryMapSize, MemoryMap, *DescriptorSize);
+      SplitTable (
+        MemoryMapSize,
+        MemoryMap,
+        *DescriptorSize,
+        &mImagePropertiesPrivateData.ImageRecordList,
+        AdditionalRecordCount
+        );
+
+      //
+      // Set RuntimeData to XP
+      //
+      EnforceMemoryMapAttribute (MemoryMap, *MemoryMapSize, *DescriptorSize);
+
+      //
+      // Merge same type to save entry size
+      //
+      MergeMemoryMap (MemoryMap, MemoryMapSize, *DescriptorSize);
     }
   }
 
@@ -1066,218 +582,6 @@ SetMemoryAttributesTableSectionAlignment (
   }
 }
 
-/**
-  Swap two code sections in image record.
-
-  @param  FirstImageRecordCodeSection    first code section in image record
-  @param  SecondImageRecordCodeSection   second code section in image record
-**/
-STATIC
-VOID
-SwapImageRecordCodeSection (
-  IN IMAGE_PROPERTIES_RECORD_CODE_SECTION  *FirstImageRecordCodeSection,
-  IN IMAGE_PROPERTIES_RECORD_CODE_SECTION  *SecondImageRecordCodeSection
-  )
-{
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  TempImageRecordCodeSection;
-
-  TempImageRecordCodeSection.CodeSegmentBase = 
FirstImageRecordCodeSection->CodeSegmentBase;
-  TempImageRecordCodeSection.CodeSegmentSize = 
FirstImageRecordCodeSection->CodeSegmentSize;
-
-  FirstImageRecordCodeSection->CodeSegmentBase = 
SecondImageRecordCodeSection->CodeSegmentBase;
-  FirstImageRecordCodeSection->CodeSegmentSize = 
SecondImageRecordCodeSection->CodeSegmentSize;
-
-  SecondImageRecordCodeSection->CodeSegmentBase = 
TempImageRecordCodeSection.CodeSegmentBase;
-  SecondImageRecordCodeSection->CodeSegmentSize = 
TempImageRecordCodeSection.CodeSegmentSize;
-}
-
-/**
-  Sort code section in image record, based upon CodeSegmentBase from low to 
high.
-
-  @param  ImageRecord    image record to be sorted
-**/
-VOID
-SortImageRecordCodeSection (
-  IN IMAGE_PROPERTIES_RECORD  *ImageRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *NextImageRecordCodeSection;
-  LIST_ENTRY                            *ImageRecordCodeSectionLink;
-  LIST_ENTRY                            *NextImageRecordCodeSectionLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionEndLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionList;
-
-  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
-
-  ImageRecordCodeSectionLink     = ImageRecordCodeSectionList->ForwardLink;
-  NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-  ImageRecordCodeSectionEndLink  = ImageRecordCodeSectionList;
-  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-    ImageRecordCodeSection = CR (
-                               ImageRecordCodeSectionLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    while (NextImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-      NextImageRecordCodeSection = CR (
-                                     NextImageRecordCodeSectionLink,
-                                     IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                                     Link,
-                                     
IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                                     );
-      if (ImageRecordCodeSection->CodeSegmentBase > 
NextImageRecordCodeSection->CodeSegmentBase) {
-        SwapImageRecordCodeSection (ImageRecordCodeSection, 
NextImageRecordCodeSection);
-      }
-
-      NextImageRecordCodeSectionLink = 
NextImageRecordCodeSectionLink->ForwardLink;
-    }
-
-    ImageRecordCodeSectionLink     = ImageRecordCodeSectionLink->ForwardLink;
-    NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-  }
-}
-
-/**
-  Check if code section in image record is valid.
-
-  @param  ImageRecord    image record to be checked
-
-  @retval TRUE  image record is valid
-  @retval FALSE image record is invalid
-**/
-BOOLEAN
-IsImageRecordCodeSectionValid (
-  IN IMAGE_PROPERTIES_RECORD  *ImageRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *LastImageRecordCodeSection;
-  LIST_ENTRY                            *ImageRecordCodeSectionLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionEndLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionList;
-
-  DEBUG ((DEBUG_VERBOSE, "ImageCode SegmentCount - 0x%x\n", 
ImageRecord->CodeSegmentCount));
-
-  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
-
-  ImageRecordCodeSectionLink    = ImageRecordCodeSectionList->ForwardLink;
-  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
-  LastImageRecordCodeSection    = NULL;
-  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-    ImageRecordCodeSection = CR (
-                               ImageRecordCodeSectionLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    if (ImageRecordCodeSection->CodeSegmentSize == 0) {
-      return FALSE;
-    }
-
-    if (ImageRecordCodeSection->CodeSegmentBase < ImageRecord->ImageBase) {
-      return FALSE;
-    }
-
-    if (ImageRecordCodeSection->CodeSegmentBase >= MAX_ADDRESS - 
ImageRecordCodeSection->CodeSegmentSize) {
-      return FALSE;
-    }
-
-    if ((ImageRecordCodeSection->CodeSegmentBase + 
ImageRecordCodeSection->CodeSegmentSize) > (ImageRecord->ImageBase + 
ImageRecord->ImageSize)) {
-      return FALSE;
-    }
-
-    if (LastImageRecordCodeSection != NULL) {
-      if ((LastImageRecordCodeSection->CodeSegmentBase + 
LastImageRecordCodeSection->CodeSegmentSize) > 
ImageRecordCodeSection->CodeSegmentBase) {
-        return FALSE;
-      }
-    }
-
-    LastImageRecordCodeSection = ImageRecordCodeSection;
-    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-  }
-
-  return TRUE;
-}
-
-/**
-  Swap two image records.
-
-  @param  FirstImageRecord   first image record.
-  @param  SecondImageRecord  second image record.
-**/
-STATIC
-VOID
-SwapImageRecord (
-  IN IMAGE_PROPERTIES_RECORD  *FirstImageRecord,
-  IN IMAGE_PROPERTIES_RECORD  *SecondImageRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD  TempImageRecord;
-
-  TempImageRecord.ImageBase        = FirstImageRecord->ImageBase;
-  TempImageRecord.ImageSize        = FirstImageRecord->ImageSize;
-  TempImageRecord.CodeSegmentCount = FirstImageRecord->CodeSegmentCount;
-
-  FirstImageRecord->ImageBase        = SecondImageRecord->ImageBase;
-  FirstImageRecord->ImageSize        = SecondImageRecord->ImageSize;
-  FirstImageRecord->CodeSegmentCount = SecondImageRecord->CodeSegmentCount;
-
-  SecondImageRecord->ImageBase        = TempImageRecord.ImageBase;
-  SecondImageRecord->ImageSize        = TempImageRecord.ImageSize;
-  SecondImageRecord->CodeSegmentCount = TempImageRecord.CodeSegmentCount;
-
-  SwapListEntries (&FirstImageRecord->CodeSegmentList, 
&SecondImageRecord->CodeSegmentList);
-}
-
-/**
-  Sort image record based upon the ImageBase from low to high.
-**/
-STATIC
-VOID
-SortImageRecord (
-  VOID
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  IMAGE_PROPERTIES_RECORD  *NextImageRecord;
-  LIST_ENTRY               *ImageRecordLink;
-  LIST_ENTRY               *NextImageRecordLink;
-  LIST_ENTRY               *ImageRecordEndLink;
-  LIST_ENTRY               *ImageRecordList;
-
-  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
-
-  ImageRecordLink     = ImageRecordList->ForwardLink;
-  NextImageRecordLink = ImageRecordLink->ForwardLink;
-  ImageRecordEndLink  = ImageRecordList;
-  while (ImageRecordLink != ImageRecordEndLink) {
-    ImageRecord = CR (
-                    ImageRecordLink,
-                    IMAGE_PROPERTIES_RECORD,
-                    Link,
-                    IMAGE_PROPERTIES_RECORD_SIGNATURE
-                    );
-    while (NextImageRecordLink != ImageRecordEndLink) {
-      NextImageRecord = CR (
-                          NextImageRecordLink,
-                          IMAGE_PROPERTIES_RECORD,
-                          Link,
-                          IMAGE_PROPERTIES_RECORD_SIGNATURE
-                          );
-      if (ImageRecord->ImageBase > NextImageRecord->ImageBase) {
-        SwapImageRecord (ImageRecord, NextImageRecord);
-      }
-
-      NextImageRecordLink = NextImageRecordLink->ForwardLink;
-    }
-
-    ImageRecordLink     = ImageRecordLink->ForwardLink;
-    NextImageRecordLink = ImageRecordLink->ForwardLink;
-  }
-}
-
 /**
   Insert image record.
 
@@ -1456,54 +760,12 @@ InsertImageRecord (
     mImagePropertiesPrivateData.CodeSegmentCountMax = 
ImageRecord->CodeSegmentCount;
   }
 
-  SortImageRecord ();
+  SortImageRecord (&mImagePropertiesPrivateData.ImageRecordList);
 
 Finish:
   return;
 }
 
-/**
-  Find image record according to image base and size.
-
-  @param  ImageBase    Base of PE image
-  @param  ImageSize    Size of PE image
-
-  @return image record
-**/
-STATIC
-IMAGE_PROPERTIES_RECORD *
-FindImageRecord (
-  IN EFI_PHYSICAL_ADDRESS  ImageBase,
-  IN UINT64                ImageSize
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  LIST_ENTRY               *ImageRecordLink;
-  LIST_ENTRY               *ImageRecordList;
-
-  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
-
-  for (ImageRecordLink = ImageRecordList->ForwardLink;
-       ImageRecordLink != ImageRecordList;
-       ImageRecordLink = ImageRecordLink->ForwardLink)
-  {
-    ImageRecord = CR (
-                    ImageRecordLink,
-                    IMAGE_PROPERTIES_RECORD,
-                    Link,
-                    IMAGE_PROPERTIES_RECORD_SIGNATURE
-                    );
-
-    if ((ImageBase == ImageRecord->ImageBase) &&
-        (ImageSize == ImageRecord->ImageSize))
-    {
-      return ImageRecord;
-    }
-  }
-
-  return NULL;
-}
-
 /**
   Remove Image record.
 
@@ -1526,7 +788,11 @@ RemoveImageRecord (
     return;
   }
 
-  ImageRecord = FindImageRecord 
((EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase, RuntimeImage->ImageSize);
+  ImageRecord = FindImageRecord (
+                  (EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase,
+                  RuntimeImage->ImageSize,
+                  &mImagePropertiesPrivateData.ImageRecordList
+                  );
   if (ImageRecord == NULL) {
     DEBUG ((DEBUG_ERROR, "!!!!!!!! ImageRecord not found !!!!!!!!\n"));
     return;
diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c 
b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
index 7cc829b17402..977239d08afc 100644
--- a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
+++ b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
@@ -32,6 +32,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/DxeServicesTableLib.h>
 #include <Library/DebugLib.h>
 #include <Library/UefiLib.h>
+#include <Library/ImagePropertiesRecordLib.h>
 
 #include <Guid/EventGroup.h>
 #include <Guid/MemoryAttributesTable.h>
@@ -66,29 +67,6 @@ extern LIST_ENTRY  mGcdMemorySpaceMap;
 
 STATIC LIST_ENTRY  mProtectedImageRecordList;
 
-/**
-  Sort code section in image record, based upon CodeSegmentBase from low to 
high.
-
-  @param  ImageRecord    image record to be sorted
-**/
-VOID
-SortImageRecordCodeSection (
-  IN IMAGE_PROPERTIES_RECORD  *ImageRecord
-  );
-
-/**
-  Check if code section in image record is valid.
-
-  @param  ImageRecord    image record to be checked
-
-  @retval TRUE  image record is valid
-  @retval FALSE image record is invalid
-**/
-BOOLEAN
-IsImageRecordCodeSectionValid (
-  IN IMAGE_PROPERTIES_RECORD  *ImageRecord
-  );
-
 /**
   Get the image type.
 
diff --git a/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c 
b/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c
index 394fdae50741..2e4aaddef4e5 100644
--- a/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c
+++ b/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c
@@ -14,6 +14,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/SmmServicesTableLib.h>
 #include <Library/DebugLib.h>
 #include <Library/PcdLib.h>
+#include <Library/ImagePropertiesRecordLib.h>
 
 #include <Library/PeCoffLib.h>
 #include <Library/PeCoffGetEntryPointLib.h>
@@ -25,26 +26,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #define PREVIOUS_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \
   ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) - (Size)))
 
-#define IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE  SIGNATURE_32 
('I','P','R','C')
-
-typedef struct {
-  UINT32                  Signature;
-  LIST_ENTRY              Link;
-  EFI_PHYSICAL_ADDRESS    CodeSegmentBase;
-  UINT64                  CodeSegmentSize;
-} IMAGE_PROPERTIES_RECORD_CODE_SECTION;
-
-#define IMAGE_PROPERTIES_RECORD_SIGNATURE  SIGNATURE_32 ('I','P','R','D')
-
-typedef struct {
-  UINT32                  Signature;
-  LIST_ENTRY              Link;
-  EFI_PHYSICAL_ADDRESS    ImageBase;
-  UINT64                  ImageSize;
-  UINTN                   CodeSegmentCount;
-  LIST_ENTRY              CodeSegmentList;
-} IMAGE_PROPERTIES_RECORD;
-
 #define IMAGE_PROPERTIES_PRIVATE_DATA_SIGNATURE  SIGNATURE_32 ('I','P','P','D')
 
 typedef struct {
@@ -69,87 +50,6 @@ UINT64  mMemoryProtectionAttribute = 
EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTEC
 // Below functions are for MemoryMap
 //
 
-/**
-  Converts a number of EFI_PAGEs to a size in bytes.
-
-  NOTE: Do not use EFI_PAGES_TO_SIZE because it handles UINTN only.
-
-  @param[in]  Pages     The number of EFI_PAGES.
-
-  @return  The number of bytes associated with the number of EFI_PAGEs 
specified
-           by Pages.
-**/
-STATIC
-UINT64
-EfiPagesToSize (
-  IN UINT64  Pages
-  )
-{
-  return LShiftU64 (Pages, EFI_PAGE_SHIFT);
-}
-
-/**
-  Converts a size, in bytes, to a number of EFI_PAGESs.
-
-  NOTE: Do not use EFI_SIZE_TO_PAGES because it handles UINTN only.
-
-  @param[in]  Size      A size in bytes.
-
-  @return  The number of EFI_PAGESs associated with the number of bytes 
specified
-           by Size.
-
-**/
-STATIC
-UINT64
-EfiSizeToPages (
-  IN UINT64  Size
-  )
-{
-  return RShiftU64 (Size, EFI_PAGE_SHIFT) + ((((UINTN)Size) & EFI_PAGE_MASK) ? 
1 : 0);
-}
-
-/**
-  Sort memory map entries based upon PhysicalStart, from low to high.
-
-  @param[in,out]  MemoryMap         A pointer to the buffer in which firmware 
places
-                                    the current memory map.
-  @param[in]      MemoryMapSize     Size, in bytes, of the MemoryMap buffer.
-  @param[in]      DescriptorSize    Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-**/
-STATIC
-VOID
-SortMemoryMap (
-  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
-  IN UINTN                      MemoryMapSize,
-  IN UINTN                      DescriptorSize
-  )
-{
-  EFI_MEMORY_DESCRIPTOR  *MemoryMapEntry;
-  EFI_MEMORY_DESCRIPTOR  *NextMemoryMapEntry;
-  EFI_MEMORY_DESCRIPTOR  *MemoryMapEnd;
-  EFI_MEMORY_DESCRIPTOR  TempMemoryMap;
-
-  MemoryMapEntry     = MemoryMap;
-  NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
-  MemoryMapEnd       = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + 
MemoryMapSize);
-  while (MemoryMapEntry < MemoryMapEnd) {
-    while (NextMemoryMapEntry < MemoryMapEnd) {
-      if (MemoryMapEntry->PhysicalStart > NextMemoryMapEntry->PhysicalStart) {
-        CopyMem (&TempMemoryMap, MemoryMapEntry, sizeof 
(EFI_MEMORY_DESCRIPTOR));
-        CopyMem (MemoryMapEntry, NextMemoryMapEntry, sizeof 
(EFI_MEMORY_DESCRIPTOR));
-        CopyMem (NextMemoryMapEntry, &TempMemoryMap, sizeof 
(EFI_MEMORY_DESCRIPTOR));
-      }
-
-      NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NextMemoryMapEntry, 
DescriptorSize);
-    }
-
-    MemoryMapEntry     = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, 
DescriptorSize);
-    NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, 
DescriptorSize);
-  }
-
-  return;
-}
-
 /**
   Merge continuous memory map entries whose have same attributes.
 
@@ -183,7 +83,7 @@ MergeMemoryMap (
     NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, 
DescriptorSize);
 
     do {
-      MemoryBlockLength = (UINT64)(EfiPagesToSize 
(MemoryMapEntry->NumberOfPages));
+      MemoryBlockLength = LShiftU64 (MemoryMapEntry->NumberOfPages, 
EFI_PAGE_SHIFT);
       if (((UINTN)NextMemoryMapEntry < (UINTN)MemoryMapEnd) &&
           (MemoryMapEntry->Type == NextMemoryMapEntry->Type) &&
           (MemoryMapEntry->Attribute == NextMemoryMapEntry->Attribute) &&
@@ -254,423 +154,6 @@ EnforceMemoryMapAttribute (
   return;
 }
 
-/**
-  Return the first image record, whose [ImageBase, ImageSize] covered by 
[Buffer, Length].
-
-  @param[in] Buffer  Start Address
-  @param[in] Length  Address length
-
-  @return first image record covered by [buffer, length]
-**/
-STATIC
-IMAGE_PROPERTIES_RECORD *
-GetImageRecordByAddress (
-  IN EFI_PHYSICAL_ADDRESS  Buffer,
-  IN UINT64                Length
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  LIST_ENTRY               *ImageRecordLink;
-  LIST_ENTRY               *ImageRecordList;
-
-  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
-
-  for (ImageRecordLink = ImageRecordList->ForwardLink;
-       ImageRecordLink != ImageRecordList;
-       ImageRecordLink = ImageRecordLink->ForwardLink)
-  {
-    ImageRecord = CR (
-                    ImageRecordLink,
-                    IMAGE_PROPERTIES_RECORD,
-                    Link,
-                    IMAGE_PROPERTIES_RECORD_SIGNATURE
-                    );
-
-    if ((Buffer <= ImageRecord->ImageBase) &&
-        (Buffer + Length >= ImageRecord->ImageBase + ImageRecord->ImageSize))
-    {
-      return ImageRecord;
-    }
-  }
-
-  return NULL;
-}
-
-/**
-  Set the memory map to new entries, according to one old entry,
-  based upon PE code section and data section in image record
-
-  @param[in]       ImageRecord            An image record whose [ImageBase, 
ImageSize] covered
-                                          by old memory map entry.
-  @param[in, out]  NewRecord              A pointer to several new memory map 
entries.
-                                          The caller guarantee the buffer size 
be 1 +
-                                          (SplitRecordCount * DescriptorSize) 
calculated
-                                          below.
-  @param[in]       OldRecord              A pointer to one old memory map 
entry.
-  @param[in]       DescriptorSize         Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-**/
-STATIC
-UINTN
-SetNewRecord (
-  IN IMAGE_PROPERTIES_RECORD    *ImageRecord,
-  IN OUT EFI_MEMORY_DESCRIPTOR  *NewRecord,
-  IN EFI_MEMORY_DESCRIPTOR      *OldRecord,
-  IN UINTN                      DescriptorSize
-  )
-{
-  EFI_MEMORY_DESCRIPTOR                 TempRecord;
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
-  LIST_ENTRY                            *ImageRecordCodeSectionLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionEndLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionList;
-  UINTN                                 NewRecordCount;
-  UINT64                                PhysicalEnd;
-  UINT64                                ImageEnd;
-
-  CopyMem (&TempRecord, OldRecord, sizeof (EFI_MEMORY_DESCRIPTOR));
-  PhysicalEnd    = TempRecord.PhysicalStart + EfiPagesToSize 
(TempRecord.NumberOfPages);
-  NewRecordCount = 0;
-
-  //
-  // Always create a new entry for non-PE image record
-  //
-  if (ImageRecord->ImageBase > TempRecord.PhysicalStart) {
-    NewRecord->Type          = TempRecord.Type;
-    NewRecord->PhysicalStart = TempRecord.PhysicalStart;
-    NewRecord->VirtualStart  = 0;
-    NewRecord->NumberOfPages = EfiSizeToPages (ImageRecord->ImageBase - 
TempRecord.PhysicalStart);
-    NewRecord->Attribute     = TempRecord.Attribute;
-    NewRecord                = NEXT_MEMORY_DESCRIPTOR (NewRecord, 
DescriptorSize);
-    NewRecordCount++;
-    TempRecord.PhysicalStart = ImageRecord->ImageBase;
-    TempRecord.NumberOfPages = EfiSizeToPages (PhysicalEnd - 
TempRecord.PhysicalStart);
-  }
-
-  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
-
-  ImageRecordCodeSectionLink    = ImageRecordCodeSectionList->ForwardLink;
-  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
-  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-    ImageRecordCodeSection = CR (
-                               ImageRecordCodeSectionLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-
-    if (TempRecord.PhysicalStart <= ImageRecordCodeSection->CodeSegmentBase) {
-      //
-      // DATA
-      //
-      NewRecord->Type          = EfiRuntimeServicesData;
-      NewRecord->PhysicalStart = TempRecord.PhysicalStart;
-      NewRecord->VirtualStart  = 0;
-      NewRecord->NumberOfPages = EfiSizeToPages 
(ImageRecordCodeSection->CodeSegmentBase - NewRecord->PhysicalStart);
-      NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
-      if (NewRecord->NumberOfPages != 0) {
-        NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
-        NewRecordCount++;
-      }
-
-      //
-      // CODE
-      //
-      NewRecord->Type          = EfiRuntimeServicesCode;
-      NewRecord->PhysicalStart = ImageRecordCodeSection->CodeSegmentBase;
-      NewRecord->VirtualStart  = 0;
-      NewRecord->NumberOfPages = EfiSizeToPages 
(ImageRecordCodeSection->CodeSegmentSize);
-      NewRecord->Attribute     = (TempRecord.Attribute & (~EFI_MEMORY_XP)) | 
EFI_MEMORY_RO;
-      if (NewRecord->NumberOfPages != 0) {
-        NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
-        NewRecordCount++;
-      }
-
-      TempRecord.PhysicalStart = ImageRecordCodeSection->CodeSegmentBase + 
EfiPagesToSize (EfiSizeToPages (ImageRecordCodeSection->CodeSegmentSize));
-      TempRecord.NumberOfPages = EfiSizeToPages (PhysicalEnd - 
TempRecord.PhysicalStart);
-      if (TempRecord.NumberOfPages == 0) {
-        break;
-      }
-    }
-  }
-
-  ImageEnd = ImageRecord->ImageBase + ImageRecord->ImageSize;
-
-  //
-  // Final DATA
-  //
-  if (TempRecord.PhysicalStart < ImageEnd) {
-    NewRecord->Type          = EfiRuntimeServicesData;
-    NewRecord->PhysicalStart = TempRecord.PhysicalStart;
-    NewRecord->VirtualStart  = 0;
-    NewRecord->NumberOfPages = EfiSizeToPages (ImageEnd - 
TempRecord.PhysicalStart);
-    NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
-    NewRecordCount++;
-  }
-
-  return NewRecordCount;
-}
-
-/**
-  Return the max number of new splitted entries, according to one old entry,
-  based upon PE code section and data section.
-
-  @param[in]  OldRecord              A pointer to one old memory map entry.
-
-  @retval  0 no entry need to be splitted.
-  @return  the max number of new splitted entries
-**/
-STATIC
-UINTN
-GetMaxSplitRecordCount (
-  IN EFI_MEMORY_DESCRIPTOR  *OldRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  UINTN                    SplitRecordCount;
-  UINT64                   PhysicalStart;
-  UINT64                   PhysicalEnd;
-
-  SplitRecordCount = 0;
-  PhysicalStart    = OldRecord->PhysicalStart;
-  PhysicalEnd      = OldRecord->PhysicalStart + EfiPagesToSize 
(OldRecord->NumberOfPages);
-
-  do {
-    ImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - 
PhysicalStart);
-    if (ImageRecord == NULL) {
-      break;
-    }
-
-    SplitRecordCount += (2 * ImageRecord->CodeSegmentCount + 2);
-    PhysicalStart     = ImageRecord->ImageBase + ImageRecord->ImageSize;
-  } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));
-
-  return SplitRecordCount;
-}
-
-/**
-  Split the memory map to new entries, according to one old entry,
-  based upon PE code section and data section.
-
-  @param[in]       OldRecord              A pointer to one old memory map 
entry.
-  @param[in, out]  NewRecord              A pointer to several new memory map 
entries.
-                                          The caller guarantee the buffer size 
be 1 +
-                                          (SplitRecordCount * DescriptorSize) 
calculated
-                                          below.
-  @param[in]       MaxSplitRecordCount    The max number of splitted entries
-  @param[in]       DescriptorSize         Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-
-  @retval  0 no entry is splitted.
-  @return  the real number of splitted record.
-**/
-STATIC
-UINTN
-SplitRecord (
-  IN EFI_MEMORY_DESCRIPTOR      *OldRecord,
-  IN OUT EFI_MEMORY_DESCRIPTOR  *NewRecord,
-  IN UINTN                      MaxSplitRecordCount,
-  IN UINTN                      DescriptorSize
-  )
-{
-  EFI_MEMORY_DESCRIPTOR    TempRecord;
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  IMAGE_PROPERTIES_RECORD  *NewImageRecord;
-  UINT64                   PhysicalStart;
-  UINT64                   PhysicalEnd;
-  UINTN                    NewRecordCount;
-  UINTN                    TotalNewRecordCount;
-
-  if (MaxSplitRecordCount == 0) {
-    CopyMem (NewRecord, OldRecord, DescriptorSize);
-    return 0;
-  }
-
-  TotalNewRecordCount = 0;
-
-  //
-  // Override previous record
-  //
-  CopyMem (&TempRecord, OldRecord, sizeof (EFI_MEMORY_DESCRIPTOR));
-  PhysicalStart = TempRecord.PhysicalStart;
-  PhysicalEnd   = TempRecord.PhysicalStart + EfiPagesToSize 
(TempRecord.NumberOfPages);
-
-  ImageRecord = NULL;
-  do {
-    NewImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - 
PhysicalStart);
-    if (NewImageRecord == NULL) {
-      //
-      // No more image covered by this range, stop
-      //
-      if (PhysicalEnd > PhysicalStart) {
-        //
-        // Always create a new entry for non-PE image record
-        //
-        NewRecord->Type          = TempRecord.Type;
-        NewRecord->PhysicalStart = TempRecord.PhysicalStart;
-        NewRecord->VirtualStart  = 0;
-        NewRecord->NumberOfPages = TempRecord.NumberOfPages;
-        NewRecord->Attribute     = TempRecord.Attribute;
-        TotalNewRecordCount++;
-      }
-
-      break;
-    }
-
-    ImageRecord = NewImageRecord;
-
-    //
-    // Set new record
-    //
-    NewRecordCount       = SetNewRecord (ImageRecord, NewRecord, &TempRecord, 
DescriptorSize);
-    TotalNewRecordCount += NewRecordCount;
-    NewRecord            = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)NewRecord + 
NewRecordCount * DescriptorSize);
-
-    //
-    // Update PhysicalStart, in order to exclude the image buffer already 
splitted.
-    //
-    PhysicalStart            = ImageRecord->ImageBase + ImageRecord->ImageSize;
-    TempRecord.PhysicalStart = PhysicalStart;
-    TempRecord.NumberOfPages = EfiSizeToPages (PhysicalEnd - PhysicalStart);
-  } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));
-
-  return TotalNewRecordCount - 1;
-}
-
-/**
-  Split the original memory map, and add more entries to describe PE code 
section and data section.
-  This function will set EfiRuntimeServicesData to be EFI_MEMORY_XP.
-  This function will merge entries with same attributes finally.
-
-  NOTE: It assumes PE code/data section are page aligned.
-  NOTE: It assumes enough entry is prepared for new memory map.
-
-  Split table:
-   +---------------+
-   | Record X      |
-   +---------------+
-   | Record RtCode |
-   +---------------+
-   | Record Y      |
-   +---------------+
-   ==>
-   +---------------+
-   | Record X      |
-   +---------------+
-   | Record RtCode |
-   +---------------+ ----
-   | Record RtData |     |
-   +---------------+     |
-   | Record RtCode |     |-> PE/COFF1
-   +---------------+     |
-   | Record RtData |     |
-   +---------------+ ----
-   | Record RtCode |
-   +---------------+ ----
-   | Record RtData |     |
-   +---------------+     |
-   | Record RtCode |     |-> PE/COFF2
-   +---------------+     |
-   | Record RtData |     |
-   +---------------+ ----
-   | Record RtCode |
-   +---------------+
-   | Record Y      |
-   +---------------+
-
-  @param[in, out]  MemoryMapSize          A pointer to the size, in bytes, of 
the
-                                          MemoryMap buffer. On input, this is 
the size of
-                                          old MemoryMap before split. The 
actual buffer
-                                          size of MemoryMap is MemoryMapSize +
-                                          (AdditionalRecordCount * 
DescriptorSize) calculated
-                                          below. On output, it is the size of 
new MemoryMap
-                                          after split.
-  @param[in, out]  MemoryMap              A pointer to the buffer in which 
firmware places
-                                          the current memory map.
-  @param[in]       DescriptorSize         Size, in bytes, of an individual 
EFI_MEMORY_DESCRIPTOR.
-**/
-STATIC
-VOID
-SplitTable (
-  IN OUT UINTN                  *MemoryMapSize,
-  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
-  IN UINTN                      DescriptorSize
-  )
-{
-  INTN   IndexOld;
-  INTN   IndexNew;
-  UINTN  MaxSplitRecordCount;
-  UINTN  RealSplitRecordCount;
-  UINTN  TotalSplitRecordCount;
-  UINTN  AdditionalRecordCount;
-
-  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax 
+ 2) * mImagePropertiesPrivateData.ImageRecordCount;
-
-  TotalSplitRecordCount = 0;
-  //
-  // Let old record point to end of valid MemoryMap buffer.
-  //
-  IndexOld = ((*MemoryMapSize) / DescriptorSize) - 1;
-  //
-  // Let new record point to end of full MemoryMap buffer.
-  //
-  IndexNew = ((*MemoryMapSize) / DescriptorSize) - 1 + AdditionalRecordCount;
-  for ( ; IndexOld >= 0; IndexOld--) {
-    MaxSplitRecordCount = GetMaxSplitRecordCount ((EFI_MEMORY_DESCRIPTOR 
*)((UINT8 *)MemoryMap + IndexOld * DescriptorSize));
-    //
-    // Split this MemoryMap record
-    //
-    IndexNew            -= MaxSplitRecordCount;
-    RealSplitRecordCount = SplitRecord (
-                             (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + 
IndexOld * DescriptorSize),
-                             (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + 
IndexNew * DescriptorSize),
-                             MaxSplitRecordCount,
-                             DescriptorSize
-                             );
-    //
-    // Adjust IndexNew according to real split.
-    //
-    if (MaxSplitRecordCount != RealSplitRecordCount) {
-      CopyMem (
-        ((UINT8 *)MemoryMap + (IndexNew + MaxSplitRecordCount - 
RealSplitRecordCount) * DescriptorSize),
-        ((UINT8 *)MemoryMap + IndexNew * DescriptorSize),
-        (RealSplitRecordCount + 1) * DescriptorSize
-        );
-    }
-
-    IndexNew               = IndexNew + MaxSplitRecordCount - 
RealSplitRecordCount;
-    TotalSplitRecordCount += RealSplitRecordCount;
-    IndexNew--;
-  }
-
-  //
-  // Move all records to the beginning.
-  //
-  CopyMem (
-    MemoryMap,
-    (UINT8 *)MemoryMap + (AdditionalRecordCount - TotalSplitRecordCount) * 
DescriptorSize,
-    (*MemoryMapSize) + TotalSplitRecordCount * DescriptorSize
-    );
-
-  *MemoryMapSize = (*MemoryMapSize) + DescriptorSize * TotalSplitRecordCount;
-
-  //
-  // Sort from low to high (Just in case)
-  //
-  SortMemoryMap (MemoryMap, *MemoryMapSize, DescriptorSize);
-
-  //
-  // Set RuntimeData to XP
-  //
-  EnforceMemoryMapAttribute (MemoryMap, *MemoryMapSize, DescriptorSize);
-
-  //
-  // Merge same type to save entry size
-  //
-  MergeMemoryMap (MemoryMap, MemoryMapSize, DescriptorSize);
-
-  return;
-}
-
 /**
   This function for GetMemoryMap() with memory attributes table.
 
@@ -729,7 +212,7 @@ SmmCoreGetMemoryMapMemoryAttributesTable (
     return EFI_INVALID_PARAMETER;
   }
 
-  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax 
+ 2) * mImagePropertiesPrivateData.ImageRecordCount;
+  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax 
+ 3) * mImagePropertiesPrivateData.ImageRecordCount;
 
   OldMemoryMapSize = *MemoryMapSize;
   Status           = SmmCoreGetMemoryMap (MemoryMapSize, MemoryMap, MapKey, 
DescriptorSize, DescriptorVersion);
@@ -747,7 +230,17 @@ SmmCoreGetMemoryMapMemoryAttributesTable (
       // Split PE code/data
       //
       ASSERT (MemoryMap != NULL);
-      SplitTable (MemoryMapSize, MemoryMap, *DescriptorSize);
+      SplitTable (MemoryMapSize, MemoryMap, *DescriptorSize, 
&mImagePropertiesPrivateData.ImageRecordList, AdditionalRecordCount);
+
+      //
+      // Set RuntimeData to XP
+      //
+      EnforceMemoryMapAttribute (MemoryMap, *MemoryMapSize, *DescriptorSize);
+
+      //
+      // Merge same type to save entry size
+      //
+      MergeMemoryMap (MemoryMap, MemoryMapSize, *DescriptorSize);
     }
   }
 
@@ -777,250 +270,6 @@ SetMemoryAttributesTableSectionAlignment (
   }
 }
 
-/**
-  Swap two code sections in image record.
-
-  @param[in]  FirstImageRecordCodeSection    first code section in image record
-  @param[in]  SecondImageRecordCodeSection   second code section in image 
record
-**/
-STATIC
-VOID
-SwapImageRecordCodeSection (
-  IN IMAGE_PROPERTIES_RECORD_CODE_SECTION  *FirstImageRecordCodeSection,
-  IN IMAGE_PROPERTIES_RECORD_CODE_SECTION  *SecondImageRecordCodeSection
-  )
-{
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  TempImageRecordCodeSection;
-
-  TempImageRecordCodeSection.CodeSegmentBase = 
FirstImageRecordCodeSection->CodeSegmentBase;
-  TempImageRecordCodeSection.CodeSegmentSize = 
FirstImageRecordCodeSection->CodeSegmentSize;
-
-  FirstImageRecordCodeSection->CodeSegmentBase = 
SecondImageRecordCodeSection->CodeSegmentBase;
-  FirstImageRecordCodeSection->CodeSegmentSize = 
SecondImageRecordCodeSection->CodeSegmentSize;
-
-  SecondImageRecordCodeSection->CodeSegmentBase = 
TempImageRecordCodeSection.CodeSegmentBase;
-  SecondImageRecordCodeSection->CodeSegmentSize = 
TempImageRecordCodeSection.CodeSegmentSize;
-}
-
-/**
-  Sort code section in image record, based upon CodeSegmentBase from low to 
high.
-
-  @param[in]  ImageRecord    image record to be sorted
-**/
-STATIC
-VOID
-SortImageRecordCodeSection (
-  IN IMAGE_PROPERTIES_RECORD  *ImageRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *NextImageRecordCodeSection;
-  LIST_ENTRY                            *ImageRecordCodeSectionLink;
-  LIST_ENTRY                            *NextImageRecordCodeSectionLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionEndLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionList;
-
-  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
-
-  ImageRecordCodeSectionLink     = ImageRecordCodeSectionList->ForwardLink;
-  NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-  ImageRecordCodeSectionEndLink  = ImageRecordCodeSectionList;
-  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-    ImageRecordCodeSection = CR (
-                               ImageRecordCodeSectionLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    while (NextImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-      NextImageRecordCodeSection = CR (
-                                     NextImageRecordCodeSectionLink,
-                                     IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                                     Link,
-                                     
IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                                     );
-      if (ImageRecordCodeSection->CodeSegmentBase > 
NextImageRecordCodeSection->CodeSegmentBase) {
-        SwapImageRecordCodeSection (ImageRecordCodeSection, 
NextImageRecordCodeSection);
-      }
-
-      NextImageRecordCodeSectionLink = 
NextImageRecordCodeSectionLink->ForwardLink;
-    }
-
-    ImageRecordCodeSectionLink     = ImageRecordCodeSectionLink->ForwardLink;
-    NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-  }
-}
-
-/**
-  Check if code section in image record is valid.
-
-  @param[in]  ImageRecord    image record to be checked
-
-  @retval TRUE  image record is valid
-  @retval FALSE image record is invalid
-**/
-STATIC
-BOOLEAN
-IsImageRecordCodeSectionValid (
-  IN IMAGE_PROPERTIES_RECORD  *ImageRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *LastImageRecordCodeSection;
-  LIST_ENTRY                            *ImageRecordCodeSectionLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionEndLink;
-  LIST_ENTRY                            *ImageRecordCodeSectionList;
-
-  DEBUG ((DEBUG_VERBOSE, "SMM ImageCode SegmentCount - 0x%x\n", 
ImageRecord->CodeSegmentCount));
-
-  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
-
-  ImageRecordCodeSectionLink    = ImageRecordCodeSectionList->ForwardLink;
-  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
-  LastImageRecordCodeSection    = NULL;
-  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-    ImageRecordCodeSection = CR (
-                               ImageRecordCodeSectionLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    if (ImageRecordCodeSection->CodeSegmentSize == 0) {
-      return FALSE;
-    }
-
-    if (ImageRecordCodeSection->CodeSegmentBase < ImageRecord->ImageBase) {
-      return FALSE;
-    }
-
-    if (ImageRecordCodeSection->CodeSegmentBase >= MAX_ADDRESS - 
ImageRecordCodeSection->CodeSegmentSize) {
-      return FALSE;
-    }
-
-    if ((ImageRecordCodeSection->CodeSegmentBase + 
ImageRecordCodeSection->CodeSegmentSize) > (ImageRecord->ImageBase + 
ImageRecord->ImageSize)) {
-      return FALSE;
-    }
-
-    if (LastImageRecordCodeSection != NULL) {
-      if ((LastImageRecordCodeSection->CodeSegmentBase + 
LastImageRecordCodeSection->CodeSegmentSize) > 
ImageRecordCodeSection->CodeSegmentBase) {
-        return FALSE;
-      }
-    }
-
-    LastImageRecordCodeSection = ImageRecordCodeSection;
-    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-  }
-
-  return TRUE;
-}
-
-/**
-  Swap two image records.
-
-  @param[in]  FirstImageRecord   first image record.
-  @param[in]  SecondImageRecord  second image record.
-**/
-STATIC
-VOID
-SwapImageRecord (
-  IN IMAGE_PROPERTIES_RECORD  *FirstImageRecord,
-  IN IMAGE_PROPERTIES_RECORD  *SecondImageRecord
-  )
-{
-  IMAGE_PROPERTIES_RECORD  TempImageRecord;
-
-  TempImageRecord.ImageBase        = FirstImageRecord->ImageBase;
-  TempImageRecord.ImageSize        = FirstImageRecord->ImageSize;
-  TempImageRecord.CodeSegmentCount = FirstImageRecord->CodeSegmentCount;
-
-  FirstImageRecord->ImageBase        = SecondImageRecord->ImageBase;
-  FirstImageRecord->ImageSize        = SecondImageRecord->ImageSize;
-  FirstImageRecord->CodeSegmentCount = SecondImageRecord->CodeSegmentCount;
-
-  SecondImageRecord->ImageBase        = TempImageRecord.ImageBase;
-  SecondImageRecord->ImageSize        = TempImageRecord.ImageSize;
-  SecondImageRecord->CodeSegmentCount = TempImageRecord.CodeSegmentCount;
-
-  SwapListEntries (&FirstImageRecord->CodeSegmentList, 
&SecondImageRecord->CodeSegmentList);
-}
-
-/**
-  Sort image record based upon the ImageBase from low to high.
-**/
-STATIC
-VOID
-SortImageRecord (
-  VOID
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  IMAGE_PROPERTIES_RECORD  *NextImageRecord;
-  LIST_ENTRY               *ImageRecordLink;
-  LIST_ENTRY               *NextImageRecordLink;
-  LIST_ENTRY               *ImageRecordEndLink;
-  LIST_ENTRY               *ImageRecordList;
-
-  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
-
-  ImageRecordLink     = ImageRecordList->ForwardLink;
-  NextImageRecordLink = ImageRecordLink->ForwardLink;
-  ImageRecordEndLink  = ImageRecordList;
-  while (ImageRecordLink != ImageRecordEndLink) {
-    ImageRecord = CR (
-                    ImageRecordLink,
-                    IMAGE_PROPERTIES_RECORD,
-                    Link,
-                    IMAGE_PROPERTIES_RECORD_SIGNATURE
-                    );
-    while (NextImageRecordLink != ImageRecordEndLink) {
-      NextImageRecord = CR (
-                          NextImageRecordLink,
-                          IMAGE_PROPERTIES_RECORD,
-                          Link,
-                          IMAGE_PROPERTIES_RECORD_SIGNATURE
-                          );
-      if (ImageRecord->ImageBase > NextImageRecord->ImageBase) {
-        SwapImageRecord (ImageRecord, NextImageRecord);
-      }
-
-      NextImageRecordLink = NextImageRecordLink->ForwardLink;
-    }
-
-    ImageRecordLink     = ImageRecordLink->ForwardLink;
-    NextImageRecordLink = ImageRecordLink->ForwardLink;
-  }
-}
-
-/**
-  Dump image record.
-**/
-STATIC
-VOID
-DumpImageRecord (
-  VOID
-  )
-{
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  LIST_ENTRY               *ImageRecordLink;
-  LIST_ENTRY               *ImageRecordList;
-  UINTN                    Index;
-
-  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
-
-  for (ImageRecordLink = ImageRecordList->ForwardLink, Index = 0;
-       ImageRecordLink != ImageRecordList;
-       ImageRecordLink = ImageRecordLink->ForwardLink, Index++)
-  {
-    ImageRecord = CR (
-                    ImageRecordLink,
-                    IMAGE_PROPERTIES_RECORD,
-                    Link,
-                    IMAGE_PROPERTIES_RECORD_SIGNATURE
-                    );
-    DEBUG ((DEBUG_VERBOSE, "SMM  Image[%d]: 0x%016lx - 0x%016lx\n", Index, 
ImageRecord->ImageBase, ImageRecord->ImageSize));
-  }
-}
-
 /**
   Insert image record.
 
@@ -1059,7 +308,7 @@ SmmInsertImageRecord (
   // Step 1: record whole region
   //
   ImageRecord->ImageBase = DriverEntry->ImageBuffer;
-  ImageRecord->ImageSize = EfiPagesToSize (DriverEntry->NumberOfPage);
+  ImageRecord->ImageSize = LShiftU64 (DriverEntry->NumberOfPage, 
EFI_PAGE_SHIFT);
 
   ImageAddress = (VOID *)(UINTN)DriverEntry->ImageBuffer;
 
@@ -1193,7 +442,7 @@ SmmInsertImageRecord (
     mImagePropertiesPrivateData.CodeSegmentCountMax = 
ImageRecord->CodeSegmentCount;
   }
 
-  SortImageRecord ();
+  SortImageRecord (&mImagePropertiesPrivateData.ImageRecordList);
 
 Finish:
   return;
@@ -1361,7 +610,7 @@ SmmInstallMemoryAttributesTable (
 
   DEBUG ((DEBUG_VERBOSE, "SMM Total Image Count - 0x%x\n", 
mImagePropertiesPrivateData.ImageRecordCount));
   DEBUG ((DEBUG_VERBOSE, "SMM Dump ImageRecord:\n"));
-  DumpImageRecord ();
+  DumpImageRecord (&mImagePropertiesPrivateData.ImageRecordList);
 
   PublishMemoryAttributesTable ();
 
diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index 43daa037be44..86a7be2f5188 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -228,26 +228,6 @@ typedef struct {
 #define LOADED_IMAGE_PRIVATE_DATA_FROM_THIS(a) \
           CR(a, LOADED_IMAGE_PRIVATE_DATA, Info, 
LOADED_IMAGE_PRIVATE_DATA_SIGNATURE)
 
-#define IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE  SIGNATURE_32 
('I','P','R','C')
-
-typedef struct {
-  UINT32                  Signature;
-  LIST_ENTRY              Link;
-  EFI_PHYSICAL_ADDRESS    CodeSegmentBase;
-  UINT64                  CodeSegmentSize;
-} IMAGE_PROPERTIES_RECORD_CODE_SECTION;
-
-#define IMAGE_PROPERTIES_RECORD_SIGNATURE  SIGNATURE_32 ('I','P','R','D')
-
-typedef struct {
-  UINT32                  Signature;
-  LIST_ENTRY              Link;
-  EFI_PHYSICAL_ADDRESS    ImageBase;
-  UINT64                  ImageSize;
-  UINTN                   CodeSegmentCount;
-  LIST_ENTRY              CodeSegmentList;
-} IMAGE_PROPERTIES_RECORD;
-
 //
 // DXE Core Global Variables
 //
diff --git a/MdeModulePkg/Core/Dxe/DxeMain.inf 
b/MdeModulePkg/Core/Dxe/DxeMain.inf
index 35d5bf0dee6f..c0aa37d118c5 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.inf
+++ b/MdeModulePkg/Core/Dxe/DxeMain.inf
@@ -94,6 +94,7 @@ [LibraryClasses]
   DebugAgentLib
   CpuExceptionHandlerLib
   PcdLib
+  ImagePropertiesRecordLib
 
 [Guids]
   gEfiEventMemoryMapChangeGuid                  ## PRODUCES             ## 
Event
diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf 
b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
index 3df44b38f13c..75a5934f0c1d 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
@@ -61,6 +61,7 @@ [LibraryClasses]
   HobLib
   SmmMemLib
   SafeIntLib
+  ImagePropertiesRecordLib
 
 [Protocols]
   gEfiDxeSmmReadyToLockProtocolGuid             ## UNDEFINED # 
SmiHandlerRegister
-- 
2.41.0.windows.2



-=-=-=-=-=-=-=-=-=-=-=-
Groups.io Links: You receive all messages sent to this group.
View/Reply Online (#107030): https://edk2.groups.io/g/devel/message/107030
Mute This Topic: https://groups.io/mt/100221566/21656
Group Owner: devel+ow...@edk2.groups.io
Unsubscribe: https://edk2.groups.io/g/devel/unsub [arch...@mail-archive.com]
-=-=-=-=-=-=-=-=-=-=-=-


Reply via email to