/** @file

  Copyright (c) 2014 - 2017, Intel Corporation. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <Uefi.h>
#include <PiDxe.h>
#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiLib.h>
#include <Library/UefiApplicationEntryPoint.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/DebugLib.h>
#include <Library/DxeServicesLib.h>
#include <Library/PrintLib.h>

#include <Protocol/SmmCommunication.h>
#include <Protocol/SmmAccess2.h>

#include <Guid/MemoryProfile.h>
#include <Guid/PiSmmCommunicationRegionTable.h>

CHAR8  *mActionString[] = {
  "Unknown",
  "gBS->AllocatePages",
  "gBS->FreePages",
  "gBS->AllocatePool",
  "gBS->FreePool",
};

CHAR8  *mSmmActionString[] = {
  "SmmUnknown",
  "gSmst->SmmAllocatePages",
  "gSmst->SmmFreePages",
  "gSmst->SmmAllocatePool",
  "gSmst->SmmFreePool",
};

typedef struct {
  MEMORY_PROFILE_ACTION    Action;
  CHAR8                    *String;
} ACTION_STRING;

ACTION_STRING  mExtActionString[] = {
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_PAGES,                  "Lib:AllocatePages"                },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_PAGES,          "Lib:AllocateRuntimePages"         },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_PAGES,         "Lib:AllocateReservedPages"        },
  { MEMORY_PROFILE_ACTION_LIB_FREE_PAGES,                      "Lib:FreePages"                    },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_PAGES,          "Lib:AllocateAlignedPages"         },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RUNTIME_PAGES,  "Lib:AllocateAlignedRuntimePages"  },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RESERVED_PAGES, "Lib:AllocateAlignedReservedPages" },
  { MEMORY_PROFILE_ACTION_LIB_FREE_ALIGNED_PAGES,              "Lib:FreeAlignedPages"             },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_POOL,                   "Lib:AllocatePool"                 },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_POOL,           "Lib:AllocateRuntimePool"          },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_POOL,          "Lib:AllocateReservedPool"         },
  { MEMORY_PROFILE_ACTION_LIB_FREE_POOL,                       "Lib:FreePool"                     },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ZERO_POOL,              "Lib:AllocateZeroPool"             },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_ZERO_POOL,      "Lib:AllocateRuntimeZeroPool"      },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_ZERO_POOL,     "Lib:AllocateReservedZeroPool"     },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_COPY_POOL,              "Lib:AllocateCopyPool"             },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_COPY_POOL,      "Lib:AllocateRuntimeCopyPool"      },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_COPY_POOL,     "Lib:AllocateReservedCopyPool"     },
  { MEMORY_PROFILE_ACTION_LIB_REALLOCATE_POOL,                 "Lib:ReallocatePool"               },
  { MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RUNTIME_POOL,         "Lib:ReallocateRuntimePool"        },
  { MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RESERVED_POOL,        "Lib:ReallocateReservedPool"       },
};

CHAR8  mUserDefinedActionString[] = { "UserDefined-0x80000000" };

CHAR8  *mMemoryTypeString[] = {
  "EfiReservedMemoryType",
  "EfiLoaderCode",
  "EfiLoaderData",
  "EfiBootServicesCode",
  "EfiBootServicesData",
  "EfiRuntimeServicesCode",
  "EfiRuntimeServicesData",
  "EfiConventionalMemory",
  "EfiUnusableMemory",
  "EfiACPIReclaimMemory",
  "EfiACPIMemoryNVS",
  "EfiMemoryMappedIO",
  "EfiMemoryMappedIOPortSpace",
  "EfiPalCode",
  "EfiPersistentMemory",
  "EfiOSReserved",
  "EfiOemReserved",
};

CHAR8  *mSubsystemString[] = {
  "Unknown",
  "NATIVE",
  "WINDOWS_GUI",
  "WINDOWS_CUI",
  "Unknown",
  "Unknown",
  "Unknown",
  "POSIX_CUI",
  "Unknown",
  "WINDOWS_CE_GUI",
  "EFI_APPLICATION",
  "EFI_BOOT_SERVICE_DRIVER",
  "EFI_RUNTIME_DRIVER",
  "EFI_ROM",
  "XBOX",
  "Unknown",
};

CHAR8  *mFileTypeString[] = {
  "Unknown",
  "RAW",
  "FREEFORM",
  "SECURITY_CORE",
  "PEI_CORE",
  "DXE_CORE",
  "PEIM",
  "DRIVER",
  "COMBINED_PEIM_DRIVER",
  "APPLICATION",
  "SMM",
  "FIRMWARE_VOLUME_IMAGE",
  "COMBINED_SMM_DXE",
  "SMM_CORE",
};

#define PROFILE_NAME_STRING_LENGTH  64
CHAR8  mNameString[PROFILE_NAME_STRING_LENGTH + 1];

//
// Profile summary information
//
#define MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE  SIGNATURE_32 ('M','P','A','S')
#define MEMORY_PROFILE_ALLOC_SUMMARY_INFO_REVISION   0x0001

typedef struct {
  MEMORY_PROFILE_COMMON_HEADER    Header;
  PHYSICAL_ADDRESS                CallerAddress;
  MEMORY_PROFILE_ACTION           Action;
  CHAR8                           *ActionString;
  UINT32                          AllocateCount;
  UINT64                          TotalSize;
} MEMORY_PROFILE_ALLOC_SUMMARY_INFO;

typedef struct {
  UINT32                               Signature;
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO    AllocSummaryInfo;
  LIST_ENTRY                           Link;
} MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA;

typedef struct {
  UINT32                        Signature;
  MEMORY_PROFILE_DRIVER_INFO    *DriverInfo;
  LIST_ENTRY                    *AllocSummaryInfoList;
  LIST_ENTRY                    Link;
} MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA;

typedef struct {
  UINT32                    Signature;
  MEMORY_PROFILE_CONTEXT    *Context;
  LIST_ENTRY                *DriverSummaryInfoList;
} MEMORY_PROFILE_CONTEXT_SUMMARY_DATA;

LIST_ENTRY                           mImageSummaryQueue = INITIALIZE_LIST_HEAD_VARIABLE (mImageSummaryQueue);
MEMORY_PROFILE_CONTEXT_SUMMARY_DATA  mMemoryProfileContextSummary;

/**
  Get the file name portion of the Pdb File Name.

  The portion of the Pdb File Name between the last backslash and
  either a following period or the end of the string is copied into
  AsciiBuffer.  The name is truncated, if necessary, to ensure that
  AsciiBuffer is not overrun.

  @param[in]  PdbFileName     Pdb file name.
  @param[out] AsciiBuffer     The resultant Ascii File Name.

**/
VOID
GetShortPdbFileName (
  IN  CHAR8  *PdbFileName,
  OUT CHAR8  *AsciiBuffer
  )
{
  UINTN  IndexPdb;    // Current work location within a Pdb string.
  UINTN  IndexBuffer; // Current work location within a Buffer string.
  UINTN  StartIndex;
  UINTN  EndIndex;

  ZeroMem (AsciiBuffer, PROFILE_NAME_STRING_LENGTH + 1);

  if (PdbFileName == NULL) {
    AsciiStrnCpyS (AsciiBuffer, PROFILE_NAME_STRING_LENGTH + 1, " ", 1);
  } else {
    StartIndex = 0;
    for (EndIndex = 0; PdbFileName[EndIndex] != 0; EndIndex++) {
    }

    for (IndexPdb = 0; PdbFileName[IndexPdb] != 0; IndexPdb++) {
      if ((PdbFileName[IndexPdb] == '\\') || (PdbFileName[IndexPdb] == '/')) {
        StartIndex = IndexPdb + 1;
      }

      if (PdbFileName[IndexPdb] == '.') {
        EndIndex = IndexPdb;
      }
    }

    IndexBuffer = 0;
    for (IndexPdb = StartIndex; IndexPdb < EndIndex; IndexPdb++) {
      AsciiBuffer[IndexBuffer] = PdbFileName[IndexPdb];
      IndexBuffer++;
      if (IndexBuffer >= PROFILE_NAME_STRING_LENGTH) {
        AsciiBuffer[PROFILE_NAME_STRING_LENGTH] = 0;
        break;
      }
    }
  }
}

/**
  Get a human readable name for an image.
  The following methods will be tried orderly:
    1. Image PDB
    2. FFS UI section
    3. Image GUID

  @param[in] DriverInfo Pointer to memory profile driver info.

  @return The resulting Ascii name string is stored in the mNameString global array.

**/
CHAR8 *
GetDriverNameString (
  IN MEMORY_PROFILE_DRIVER_INFO  *DriverInfo
  )
{
  EFI_STATUS  Status;
  CHAR16      *NameString;
  UINTN       StringSize;

  //
  // Method 1: Get the name string from image PDB
  //
  if (DriverInfo->PdbStringOffset != 0) {
    GetShortPdbFileName ((CHAR8 *)((UINTN)DriverInfo + DriverInfo->PdbStringOffset), mNameString);
    return mNameString;
  }

  if (!IsZeroGuid (&DriverInfo->FileName)) {
    //
    // Try to get the image's FFS UI section by image GUID
    //
    NameString = NULL;
    StringSize = 0;
    Status     = GetSectionFromAnyFv (
                   &DriverInfo->FileName,
                   EFI_SECTION_USER_INTERFACE,
                   0,
                   (VOID **)&NameString,
                   &StringSize
                   );
    if (!EFI_ERROR (Status)) {
      //
      // Method 2: Get the name string from FFS UI section
      //
      if (StrLen (NameString) > PROFILE_NAME_STRING_LENGTH) {
        NameString[PROFILE_NAME_STRING_LENGTH] = 0;
      }

      UnicodeStrToAsciiStrS (NameString, mNameString, sizeof (mNameString));
      FreePool (NameString);
      return mNameString;
    }
  }

  //
  // Method 3: Get the name string from image GUID
  //
  AsciiSPrint (mNameString, sizeof (mNameString), "%g", &DriverInfo->FileName);
  return mNameString;
}

/**
  Memory type to string.

  @param[in] MemoryType Memory type.

  @return Pointer to string.

**/
CHAR8 *
ProfileMemoryTypeToStr (
  IN EFI_MEMORY_TYPE  MemoryType
  )
{
  UINTN  Index;

  if ((UINT32)MemoryType >= 0x80000000) {
    //
    // OS reserved memory type.
    //
    Index = EfiMaxMemoryType;
  } else if ((UINT32)MemoryType >= 0x70000000) {
    //
    // OEM reserved memory type.
    //
    Index = EfiMaxMemoryType + 1;
  } else {
    Index = MemoryType;
  }

  return mMemoryTypeString[Index];
}

/**
  Action to string.

  @param[in] Action                     Profile action.
  @param[in] UserDefinedActionString    Pointer to user defined action string.
  @param[in] IsForSmm                   TRUE  - SMRAM profile.
                                        FALSE - UEFI memory profile.

  @return Pointer to string.

**/
CHAR8 *
ProfileActionToStr (
  IN MEMORY_PROFILE_ACTION  Action,
  IN CHAR8                  *UserDefinedActionString,
  IN BOOLEAN                IsForSmm
  )
{
  UINTN  Index;
  UINTN  ActionStringCount;
  CHAR8  **ActionString;

  if (IsForSmm) {
    ActionString      = mSmmActionString;
    ActionStringCount = ARRAY_SIZE (mSmmActionString);
  } else {
    ActionString      = mActionString;
    ActionStringCount = ARRAY_SIZE (mActionString);
  }

  if ((UINTN)(UINT32)Action < ActionStringCount) {
    return ActionString[Action];
  }

  for (Index = 0; Index < ARRAY_SIZE (mExtActionString); Index++) {
    if (mExtActionString[Index].Action == Action) {
      return mExtActionString[Index].String;
    }
  }

  if ((Action & MEMORY_PROFILE_ACTION_USER_DEFINED_MASK) != 0) {
    if (UserDefinedActionString != NULL) {
      return UserDefinedActionString;
    }

    AsciiSPrint (mUserDefinedActionString, sizeof (mUserDefinedActionString), "UserDefined-0x%08x", Action);
    return mUserDefinedActionString;
  }

  return ActionString[0];
}

/**
  Dump memory profile allocate information.

  @param[in] DriverInfo         Pointer to memory profile driver info.
  @param[in] AllocIndex         Memory profile alloc info index.
  @param[in] AllocInfo          Pointer to memory profile alloc info.
  @param[in] IsForSmm           TRUE  - SMRAM profile.
                                FALSE - UEFI memory profile.

  @return Pointer to next memory profile alloc info.

**/
MEMORY_PROFILE_ALLOC_INFO *
DumpMemoryProfileAllocInfo (
  IN MEMORY_PROFILE_DRIVER_INFO  *DriverInfo,
  IN UINTN                       AllocIndex,
  IN MEMORY_PROFILE_ALLOC_INFO   *AllocInfo,
  IN BOOLEAN                     IsForSmm
  )
{
  CHAR8  *ActionString;

  if (AllocInfo->Header.Signature != MEMORY_PROFILE_ALLOC_INFO_SIGNATURE) {
    return NULL;
  }

  if (AllocInfo->ActionStringOffset != 0) {
    ActionString = (CHAR8 *)((UINTN)AllocInfo + AllocInfo->ActionStringOffset);
  } else {
    ActionString = NULL;
  }

  Print (L"    MEMORY_PROFILE_ALLOC_INFO (0x%x)\n", AllocIndex);
  Print (L"      Signature     - 0x%08x\n", AllocInfo->Header.Signature);
  Print (L"      Length        - 0x%04x\n", AllocInfo->Header.Length);
  Print (L"      Revision      - 0x%04x\n", AllocInfo->Header.Revision);
  Print (L"      CallerAddress - 0x%016lx (Offset: 0x%08x)\n", AllocInfo->CallerAddress, (UINTN)(AllocInfo->CallerAddress - DriverInfo->ImageBase));
  Print (L"      SequenceId    - 0x%08x\n", AllocInfo->SequenceId);
  Print (L"      Action        - 0x%08x (%a)\n", AllocInfo->Action, ProfileActionToStr (AllocInfo->Action, ActionString, IsForSmm));
  Print (L"      MemoryType    - 0x%08x (%a)\n", AllocInfo->MemoryType, ProfileMemoryTypeToStr (AllocInfo->MemoryType));
  Print (L"      Buffer        - 0x%016lx\n", AllocInfo->Buffer);
  Print (L"      Size          - 0x%016lx\n", AllocInfo->Size);

  return (MEMORY_PROFILE_ALLOC_INFO *)((UINTN)AllocInfo + AllocInfo->Header.Length);
}

/**
  Dump memory profile driver information.

  @param[in] DriverIndex        Memory profile driver info index.
  @param[in] DriverInfo         Pointer to memory profile driver info.
  @param[in] IsForSmm           TRUE  - SMRAM profile.
                                FALSE - UEFI memory profile.

  @return Pointer to next memory profile driver info.

**/
MEMORY_PROFILE_DRIVER_INFO *
DumpMemoryProfileDriverInfo (
  IN UINTN                       DriverIndex,
  IN MEMORY_PROFILE_DRIVER_INFO  *DriverInfo,
  IN BOOLEAN                     IsForSmm
  )
{
  UINTN                      TypeIndex;
  MEMORY_PROFILE_ALLOC_INFO  *AllocInfo;
  UINTN                      AllocIndex;
  CHAR8                      *NameString;

  if (DriverInfo->Header.Signature != MEMORY_PROFILE_DRIVER_INFO_SIGNATURE) {
    return NULL;
  }

  Print (L"  MEMORY_PROFILE_DRIVER_INFO (0x%x)\n", DriverIndex);
  Print (L"    Signature               - 0x%08x\n", DriverInfo->Header.Signature);
  Print (L"    Length                  - 0x%04x\n", DriverInfo->Header.Length);
  Print (L"    Revision                - 0x%04x\n", DriverInfo->Header.Revision);
  NameString = GetDriverNameString (DriverInfo);
  Print (L"    FileName                - %a\n", NameString);
  if (DriverInfo->PdbStringOffset != 0) {
    Print (L"    Pdb                     - %a\n", (CHAR8 *)((UINTN)DriverInfo + DriverInfo->PdbStringOffset));
  }

  Print (L"    ImageBase               - 0x%016lx\n", DriverInfo->ImageBase);
  Print (L"    ImageSize               - 0x%016lx\n", DriverInfo->ImageSize);
  Print (L"    EntryPoint              - 0x%016lx\n", DriverInfo->EntryPoint);
  Print (L"    ImageSubsystem          - 0x%04x (%a)\n", DriverInfo->ImageSubsystem, mSubsystemString[(DriverInfo->ImageSubsystem < sizeof (mSubsystemString)/sizeof (mSubsystemString[0])) ? DriverInfo->ImageSubsystem : 0]);
  Print (L"    FileType                - 0x%02x (%a)\n", DriverInfo->FileType, mFileTypeString[(DriverInfo->FileType < sizeof (mFileTypeString)/sizeof (mFileTypeString[0])) ? DriverInfo->FileType : 0]);
  Print (L"    CurrentUsage            - 0x%016lx\n", DriverInfo->CurrentUsage);
  Print (L"    PeakUsage               - 0x%016lx\n", DriverInfo->PeakUsage);
  for (TypeIndex = 0; TypeIndex < sizeof (DriverInfo->CurrentUsageByType) / sizeof (DriverInfo->CurrentUsageByType[0]); TypeIndex++) {
    if ((DriverInfo->CurrentUsageByType[TypeIndex] != 0) ||
        (DriverInfo->PeakUsageByType[TypeIndex] != 0))
    {
      Print (L"    CurrentUsage[0x%02x]      - 0x%016lx (%a)\n", TypeIndex, DriverInfo->CurrentUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);
      Print (L"    PeakUsage[0x%02x]         - 0x%016lx (%a)\n", TypeIndex, DriverInfo->PeakUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);
    }
  }

  Print (L"    AllocRecordCount        - 0x%08x\n", DriverInfo->AllocRecordCount);

  AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *)((UINTN)DriverInfo + DriverInfo->Header.Length);
  for (AllocIndex = 0; AllocIndex < DriverInfo->AllocRecordCount; AllocIndex++) {
    AllocInfo = DumpMemoryProfileAllocInfo (DriverInfo, AllocIndex, AllocInfo, IsForSmm);
    if (AllocInfo == NULL) {
      return NULL;
    }
  }

  return (MEMORY_PROFILE_DRIVER_INFO *)AllocInfo;
}

/**
  Dump memory profile context information.

  @param[in] Context            Pointer to memory profile context.
  @param[in] IsForSmm           TRUE  - SMRAM profile.
                                FALSE - UEFI memory profile.

  @return Pointer to the end of memory profile context buffer.

**/
VOID *
DumpMemoryProfileContext (
  IN MEMORY_PROFILE_CONTEXT  *Context,
  IN BOOLEAN                 IsForSmm
  )
{
  UINTN                       TypeIndex;
  MEMORY_PROFILE_DRIVER_INFO  *DriverInfo;
  UINTN                       DriverIndex;

  if (Context->Header.Signature != MEMORY_PROFILE_CONTEXT_SIGNATURE) {
    return NULL;
  }

  Print (L"MEMORY_PROFILE_CONTEXT\n");
  Print (L"  Signature                     - 0x%08x\n", Context->Header.Signature);
  Print (L"  Length                        - 0x%04x\n", Context->Header.Length);
  Print (L"  Revision                      - 0x%04x\n", Context->Header.Revision);
  Print (L"  CurrentTotalUsage             - 0x%016lx\n", Context->CurrentTotalUsage);
  Print (L"  PeakTotalUsage                - 0x%016lx\n", Context->PeakTotalUsage);
  for (TypeIndex = 0; TypeIndex < sizeof (Context->CurrentTotalUsageByType) / sizeof (Context->CurrentTotalUsageByType[0]); TypeIndex++) {
    if ((Context->CurrentTotalUsageByType[TypeIndex] != 0) ||
        (Context->PeakTotalUsageByType[TypeIndex] != 0))
    {
      Print (L"  CurrentTotalUsage[0x%02x]       - 0x%016lx (%a)\n", TypeIndex, Context->CurrentTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);
      Print (L"  PeakTotalUsage[0x%02x]          - 0x%016lx (%a)\n", TypeIndex, Context->PeakTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);
    }
  }

  Print (L"  TotalImageSize                - 0x%016lx\n", Context->TotalImageSize);
  Print (L"  ImageCount                    - 0x%08x\n", Context->ImageCount);
  Print (L"  SequenceCount                 - 0x%08x\n", Context->SequenceCount);

  DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)((UINTN)Context + Context->Header.Length);
  for (DriverIndex = 0; DriverIndex < Context->ImageCount; DriverIndex++) {
    DriverInfo = DumpMemoryProfileDriverInfo (DriverIndex, DriverInfo, IsForSmm);
    if (DriverInfo == NULL) {
      return NULL;
    }
  }

  return (VOID *)DriverInfo;
}

/**
  Dump memory profile descriptor information.

  @param[in] DescriptorIndex    Memory profile descriptor index.
  @param[in] Descriptor         Pointer to memory profile descriptor.

  @return Pointer to next memory profile descriptor.

**/
MEMORY_PROFILE_DESCRIPTOR *
DumpMemoryProfileDescriptor (
  IN UINTN                      DescriptorIndex,
  IN MEMORY_PROFILE_DESCRIPTOR  *Descriptor
  )
{
  if (Descriptor->Header.Signature != MEMORY_PROFILE_DESCRIPTOR_SIGNATURE) {
    return NULL;
  }

  Print (L"  MEMORY_PROFILE_DESCRIPTOR (0x%x)\n", DescriptorIndex);
  Print (L"    Signature               - 0x%08x\n", Descriptor->Header.Signature);
  Print (L"    Length                  - 0x%04x\n", Descriptor->Header.Length);
  Print (L"    Revision                - 0x%04x\n", Descriptor->Header.Revision);
  Print (L"    Address                 - 0x%016lx\n", Descriptor->Address);
  Print (L"    Size                    - 0x%016lx\n", Descriptor->Size);

  return (MEMORY_PROFILE_DESCRIPTOR *)((UINTN)Descriptor + Descriptor->Header.Length);
}

/**
  Dump memory profile free memory information.

  @param[in] FreeMemory         Pointer to memory profile free memory.

  @return Pointer to the end of memory profile free memory buffer.

**/
VOID *
DumpMemoryProfileFreeMemory (
  IN MEMORY_PROFILE_FREE_MEMORY  *FreeMemory
  )
{
  MEMORY_PROFILE_DESCRIPTOR  *Descriptor;
  UINTN                      DescriptorIndex;

  if (FreeMemory->Header.Signature != MEMORY_PROFILE_FREE_MEMORY_SIGNATURE) {
    return NULL;
  }

  Print (L"MEMORY_PROFILE_FREE_MEMORY\n");
  Print (L"  Signature                     - 0x%08x\n", FreeMemory->Header.Signature);
  Print (L"  Length                        - 0x%04x\n", FreeMemory->Header.Length);
  Print (L"  Revision                      - 0x%04x\n", FreeMemory->Header.Revision);
  Print (L"  TotalFreeMemoryPages          - 0x%016lx\n", FreeMemory->TotalFreeMemoryPages);
  Print (L"  FreeMemoryEntryCount          - 0x%08x\n", FreeMemory->FreeMemoryEntryCount);

  Descriptor = (MEMORY_PROFILE_DESCRIPTOR *)((UINTN)FreeMemory + FreeMemory->Header.Length);
  for (DescriptorIndex = 0; DescriptorIndex < FreeMemory->FreeMemoryEntryCount; DescriptorIndex++) {
    Descriptor = DumpMemoryProfileDescriptor (DescriptorIndex, Descriptor);
    if (Descriptor == NULL) {
      return NULL;
    }
  }

  return (VOID *)Descriptor;
}

/**
  Dump memory profile memory range information.

  @param[in] MemoryRange        Pointer to memory profile memory range.

  @return Pointer to the end of memory profile memory range buffer.

**/
VOID *
DumpMemoryProfileMemoryRange (
  IN MEMORY_PROFILE_MEMORY_RANGE  *MemoryRange
  )
{
  MEMORY_PROFILE_DESCRIPTOR  *Descriptor;
  UINTN                      DescriptorIndex;

  if (MemoryRange->Header.Signature != MEMORY_PROFILE_MEMORY_RANGE_SIGNATURE) {
    return NULL;
  }

  Print (L"MEMORY_PROFILE_MEMORY_RANGE\n");
  Print (L"  Signature                     - 0x%08x\n", MemoryRange->Header.Signature);
  Print (L"  Length                        - 0x%04x\n", MemoryRange->Header.Length);
  Print (L"  Revision                      - 0x%04x\n", MemoryRange->Header.Revision);
  Print (L"  MemoryRangeCount              - 0x%08x\n", MemoryRange->MemoryRangeCount);

  Descriptor = (MEMORY_PROFILE_DESCRIPTOR *)((UINTN)MemoryRange + MemoryRange->Header.Length);
  for (DescriptorIndex = 0; DescriptorIndex < MemoryRange->MemoryRangeCount; DescriptorIndex++) {
    Descriptor = DumpMemoryProfileDescriptor (DescriptorIndex, Descriptor);
    if (Descriptor == NULL) {
      return NULL;
    }
  }

  return (VOID *)Descriptor;
}

/**
  Scan memory profile by Signature.

  @param[in] ProfileBuffer      Memory profile base address.
  @param[in] ProfileSize        Memory profile size.
  @param[in] Signature          Signature.

  @return Pointer to the structure with the signature.

**/
VOID *
ScanMemoryProfileBySignature (
  IN PHYSICAL_ADDRESS  ProfileBuffer,
  IN UINT64            ProfileSize,
  IN UINT32            Signature
  )
{
  MEMORY_PROFILE_COMMON_HEADER  *CommonHeader;
  UINTN                         ProfileEnd;

  ProfileEnd   = (UINTN)(ProfileBuffer + ProfileSize);
  CommonHeader = (MEMORY_PROFILE_COMMON_HEADER *)(UINTN)ProfileBuffer;
  while ((UINTN)CommonHeader < ProfileEnd) {
    if (CommonHeader->Signature == Signature) {
      //
      // Found it.
      //
      return (VOID *)CommonHeader;
    }

    if (CommonHeader->Length == 0) {
      ASSERT (FALSE);
      return NULL;
    }

    CommonHeader = (MEMORY_PROFILE_COMMON_HEADER *)((UINTN)CommonHeader + CommonHeader->Length);
  }

  return NULL;
}

/**
  Dump memory profile information.

  @param[in] ProfileBuffer      Memory profile base address.
  @param[in] ProfileSize        Memory profile size.
  @param[in] IsForSmm           TRUE  - SMRAM profile.
                                FALSE - UEFI memory profile.

**/
VOID
DumpMemoryProfile (
  IN PHYSICAL_ADDRESS  ProfileBuffer,
  IN UINT64            ProfileSize,
  IN BOOLEAN           IsForSmm
  )
{
  MEMORY_PROFILE_CONTEXT       *Context;
  MEMORY_PROFILE_FREE_MEMORY   *FreeMemory;
  MEMORY_PROFILE_MEMORY_RANGE  *MemoryRange;

  Context = (MEMORY_PROFILE_CONTEXT *)ScanMemoryProfileBySignature (ProfileBuffer, ProfileSize, MEMORY_PROFILE_CONTEXT_SIGNATURE);
  if (Context != NULL) {
    DumpMemoryProfileContext (Context, IsForSmm);
  }

  FreeMemory = (MEMORY_PROFILE_FREE_MEMORY *)ScanMemoryProfileBySignature (ProfileBuffer, ProfileSize, MEMORY_PROFILE_FREE_MEMORY_SIGNATURE);
  if (FreeMemory != NULL) {
    DumpMemoryProfileFreeMemory (FreeMemory);
  }

  MemoryRange = (MEMORY_PROFILE_MEMORY_RANGE *)ScanMemoryProfileBySignature (ProfileBuffer, ProfileSize, MEMORY_PROFILE_MEMORY_RANGE_SIGNATURE);
  if (MemoryRange != NULL) {
    DumpMemoryProfileMemoryRange (MemoryRange);
  }
}

/**
  Get Allocate summary information structure by caller address.

  @param[in] CallerAddress          Caller address.
  @param[in] DriverSummaryInfoData  Driver summary information data structure.

  @return Allocate summary information structure by caller address.

**/
MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA *
GetAllocSummaryInfoByCallerAddress (
  IN PHYSICAL_ADDRESS                         CallerAddress,
  IN MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA  *DriverSummaryInfoData
  )
{
  LIST_ENTRY                              *AllocSummaryInfoList;
  LIST_ENTRY                              *AllocSummaryLink;
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO       *AllocSummaryInfo;
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA  *AllocSummaryInfoData;

  AllocSummaryInfoList = DriverSummaryInfoData->AllocSummaryInfoList;

  for (AllocSummaryLink = AllocSummaryInfoList->ForwardLink;
       AllocSummaryLink != AllocSummaryInfoList;
       AllocSummaryLink = AllocSummaryLink->ForwardLink)
  {
    AllocSummaryInfoData = CR (
                             AllocSummaryLink,
                             MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA,
                             Link,
                             MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE
                             );
    AllocSummaryInfo = &AllocSummaryInfoData->AllocSummaryInfo;
    if (AllocSummaryInfo->CallerAddress == CallerAddress) {
      return AllocSummaryInfoData;
    }
  }

  return NULL;
}

/**
  Create Allocate summary information structure and
  link to Driver summary information data structure.

  @param[in, out] DriverSummaryInfoData Driver summary information data structure.
  @param[in]      AllocInfo             Pointer to memory profile alloc info.

  @return Pointer to next memory profile alloc info.

**/
MEMORY_PROFILE_ALLOC_INFO *
CreateAllocSummaryInfo (
  IN OUT MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA  *DriverSummaryInfoData,
  IN MEMORY_PROFILE_ALLOC_INFO                    *AllocInfo
  )
{
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA  *AllocSummaryInfoData;
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO       *AllocSummaryInfo;

  if (AllocInfo->Header.Signature != MEMORY_PROFILE_ALLOC_INFO_SIGNATURE) {
    return NULL;
  }

  AllocSummaryInfoData = GetAllocSummaryInfoByCallerAddress (AllocInfo->CallerAddress, DriverSummaryInfoData);
  if (AllocSummaryInfoData == NULL) {
    AllocSummaryInfoData = AllocatePool (sizeof (*AllocSummaryInfoData));
    if (AllocSummaryInfoData == NULL) {
      return NULL;
    }

    AllocSummaryInfoData->Signature    = MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE;
    AllocSummaryInfo                   = &AllocSummaryInfoData->AllocSummaryInfo;
    AllocSummaryInfo->Header.Signature = MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE;
    AllocSummaryInfo->Header.Length    = sizeof (*AllocSummaryInfo);
    AllocSummaryInfo->Header.Revision  = MEMORY_PROFILE_ALLOC_SUMMARY_INFO_REVISION;
    AllocSummaryInfo->CallerAddress    = AllocInfo->CallerAddress;
    AllocSummaryInfo->Action           = AllocInfo->Action;
    if (AllocInfo->ActionStringOffset != 0) {
      AllocSummaryInfo->ActionString = (CHAR8 *)((UINTN)AllocInfo + AllocInfo->ActionStringOffset);
    } else {
      AllocSummaryInfo->ActionString = NULL;
    }

    AllocSummaryInfo->AllocateCount = 0;
    AllocSummaryInfo->TotalSize     = 0;
    InsertTailList (DriverSummaryInfoData->AllocSummaryInfoList, &AllocSummaryInfoData->Link);
  }

  AllocSummaryInfo = &AllocSummaryInfoData->AllocSummaryInfo;
  AllocSummaryInfo->AllocateCount++;
  AllocSummaryInfo->TotalSize += AllocInfo->Size;

  return (MEMORY_PROFILE_ALLOC_INFO *)((UINTN)AllocInfo + AllocInfo->Header.Length);
}

/**
  Create Driver summary information structure and
  link to Context summary information data structure.

  @param[in, out] ContextSummaryData    Context summary information data structure.
  @param[in]      DriverInfo            Pointer to memory profile driver info.

  @return Pointer to next memory profile driver info.

**/
MEMORY_PROFILE_DRIVER_INFO *
CreateDriverSummaryInfo (
  IN OUT MEMORY_PROFILE_CONTEXT_SUMMARY_DATA  *ContextSummaryData,
  IN MEMORY_PROFILE_DRIVER_INFO               *DriverInfo
  )
{
  MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA  *DriverSummaryInfoData;
  MEMORY_PROFILE_ALLOC_INFO                *AllocInfo;
  UINTN                                    AllocIndex;

  if (DriverInfo->Header.Signature != MEMORY_PROFILE_DRIVER_INFO_SIGNATURE) {
    return NULL;
  }

  DriverSummaryInfoData = AllocatePool (sizeof (*DriverSummaryInfoData) + sizeof (LIST_ENTRY));
  if (DriverSummaryInfoData == NULL) {
    return NULL;
  }

  DriverSummaryInfoData->Signature            = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;
  DriverSummaryInfoData->DriverInfo           = DriverInfo;
  DriverSummaryInfoData->AllocSummaryInfoList = (LIST_ENTRY *)(DriverSummaryInfoData + 1);
  InitializeListHead (DriverSummaryInfoData->AllocSummaryInfoList);
  InsertTailList (ContextSummaryData->DriverSummaryInfoList, &DriverSummaryInfoData->Link);

  AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *)((UINTN)DriverInfo + DriverInfo->Header.Length);
  for (AllocIndex = 0; AllocIndex < DriverInfo->AllocRecordCount; AllocIndex++) {
    AllocInfo = CreateAllocSummaryInfo (DriverSummaryInfoData, AllocInfo);
    if (AllocInfo == NULL) {
      return NULL;
    }
  }

  return (MEMORY_PROFILE_DRIVER_INFO *)AllocInfo;
}

/**
  Create Context summary information structure.

  @param[in] ProfileBuffer      Memory profile base address.
  @param[in] ProfileSize        Memory profile size.

  @return Context summary information structure.

**/
MEMORY_PROFILE_CONTEXT_SUMMARY_DATA *
CreateContextSummaryData (
  IN PHYSICAL_ADDRESS  ProfileBuffer,
  IN UINT64            ProfileSize
  )
{
  MEMORY_PROFILE_CONTEXT      *Context;
  MEMORY_PROFILE_DRIVER_INFO  *DriverInfo;
  UINTN                       DriverIndex;

  Context = (MEMORY_PROFILE_CONTEXT *)ScanMemoryProfileBySignature (ProfileBuffer, ProfileSize, MEMORY_PROFILE_CONTEXT_SIGNATURE);
  if (Context == NULL) {
    return NULL;
  }

  mMemoryProfileContextSummary.Signature             = MEMORY_PROFILE_CONTEXT_SIGNATURE;
  mMemoryProfileContextSummary.Context               = Context;
  mMemoryProfileContextSummary.DriverSummaryInfoList = &mImageSummaryQueue;

  DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)((UINTN)Context + Context->Header.Length);
  for (DriverIndex = 0; DriverIndex < Context->ImageCount; DriverIndex++) {
    DriverInfo = CreateDriverSummaryInfo (&mMemoryProfileContextSummary, DriverInfo);
    if (DriverInfo == NULL) {
      return NULL;
    }
  }

  return &mMemoryProfileContextSummary;
}

/**
  Dump Context summary information.

  @param[in] ContextSummaryData Context summary information data.
  @param[in] IsForSmm           TRUE  - SMRAM profile.
                                FALSE - UEFI memory profile.

**/
VOID
DumpContextSummaryData (
  IN MEMORY_PROFILE_CONTEXT_SUMMARY_DATA  *ContextSummaryData,
  IN BOOLEAN                              IsForSmm
  )
{
  MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA  *DriverSummaryInfoData;
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA   *AllocSummaryInfoData;
  LIST_ENTRY                               *DriverSummaryInfoList;
  LIST_ENTRY                               *DriverSummaryLink;
  LIST_ENTRY                               *AllocSummaryInfoList;
  LIST_ENTRY                               *AllocSummaryLink;
  MEMORY_PROFILE_DRIVER_INFO               *DriverInfo;
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO        *AllocSummaryInfo;
  CHAR8                                    *NameString;

  if (ContextSummaryData == NULL) {
    return;
  }

  Print (L"\nSummary Data:\n");

  DriverSummaryInfoList = ContextSummaryData->DriverSummaryInfoList;
  for (DriverSummaryLink = DriverSummaryInfoList->ForwardLink;
       DriverSummaryLink != DriverSummaryInfoList;
       DriverSummaryLink = DriverSummaryLink->ForwardLink)
  {
    DriverSummaryInfoData = CR (
                              DriverSummaryLink,
                              MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA,
                              Link,
                              MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                              );
    DriverInfo = DriverSummaryInfoData->DriverInfo;

    NameString = GetDriverNameString (DriverInfo);
    Print (L"\nDriver - %a (Usage - 0x%08x)", NameString, DriverInfo->CurrentUsage);
    if (DriverInfo->CurrentUsage == 0) {
      Print (L"\n");
      continue;
    }

    if (DriverInfo->PdbStringOffset != 0) {
      Print (L" (Pdb - %a)\n", (CHAR8 *)((UINTN)DriverInfo + DriverInfo->PdbStringOffset));
    } else {
      Print (L"\n");
    }

    Print (L"Caller List:\n");
    Print (L"  Count            Size                   RVA              Action\n");
    Print (L"==========  ==================     ================== (================================)\n");
    AllocSummaryInfoList = DriverSummaryInfoData->AllocSummaryInfoList;
    for (AllocSummaryLink = AllocSummaryInfoList->ForwardLink;
         AllocSummaryLink != AllocSummaryInfoList;
         AllocSummaryLink = AllocSummaryLink->ForwardLink)
    {
      AllocSummaryInfoData = CR (
                               AllocSummaryLink,
                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA,
                               Link,
                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE
                               );
      AllocSummaryInfo = &AllocSummaryInfoData->AllocSummaryInfo;

      Print (
        L"0x%08x  0x%016lx <== 0x%016lx",
        AllocSummaryInfo->AllocateCount,
        AllocSummaryInfo->TotalSize,
        AllocSummaryInfo->CallerAddress - DriverInfo->ImageBase
        );
      Print (L" (%a)\n", ProfileActionToStr (AllocSummaryInfo->Action, AllocSummaryInfo->ActionString, IsForSmm));
    }
  }

  return;
}

/**
  Destroy Context summary information.

  @param[in, out] ContextSummaryData    Context summary information data.

**/
VOID
DestroyContextSummaryData (
  IN OUT MEMORY_PROFILE_CONTEXT_SUMMARY_DATA  *ContextSummaryData
  )
{
  MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA  *DriverSummaryInfoData;
  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA   *AllocSummaryInfoData;
  LIST_ENTRY                               *DriverSummaryInfoList;
  LIST_ENTRY                               *DriverSummaryLink;
  LIST_ENTRY                               *AllocSummaryInfoList;
  LIST_ENTRY                               *AllocSummaryLink;

  if (ContextSummaryData == NULL) {
    return;
  }

  DriverSummaryInfoList = ContextSummaryData->DriverSummaryInfoList;
  for (DriverSummaryLink = DriverSummaryInfoList->ForwardLink;
       DriverSummaryLink != DriverSummaryInfoList;
       )
  {
    DriverSummaryInfoData = CR (
                              DriverSummaryLink,
                              MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA,
                              Link,
                              MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                              );
    DriverSummaryLink = DriverSummaryLink->ForwardLink;

    AllocSummaryInfoList = DriverSummaryInfoData->AllocSummaryInfoList;
    for (AllocSummaryLink = AllocSummaryInfoList->ForwardLink;
         AllocSummaryLink != AllocSummaryInfoList;
         )
    {
      AllocSummaryInfoData = CR (
                               AllocSummaryLink,
                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA,
                               Link,
                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE
                               );
      AllocSummaryLink = AllocSummaryLink->ForwardLink;

      RemoveEntryList (&AllocSummaryInfoData->Link);
      FreePool (AllocSummaryInfoData);
    }

    RemoveEntryList (&DriverSummaryInfoData->Link);
    FreePool (DriverSummaryInfoData);
  }

  return;
}

/**
  Get and dump UEFI memory profile data.

  @return EFI_SUCCESS   Get the memory profile data successfully.
  @return other         Fail to get the memory profile data.

**/
EFI_STATUS
GetUefiMemoryProfileData (
  VOID
  )
{
  EFI_STATUS                           Status;
  EDKII_MEMORY_PROFILE_PROTOCOL        *ProfileProtocol;
  VOID                                 *Data;
  UINT64                               Size;
  MEMORY_PROFILE_CONTEXT_SUMMARY_DATA  *MemoryProfileContextSummaryData;
  BOOLEAN                              RecordingState;

  Status = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID **)&ProfileProtocol);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "UefiMemoryProfile: Locate MemoryProfile protocol - %r\n", Status));
    return Status;
  }

  //
  // Set recording state if needed.
  //
  RecordingState = MEMORY_PROFILE_RECORDING_DISABLE;
  Status         = ProfileProtocol->GetRecordingState (ProfileProtocol, &RecordingState);
  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {
    ProfileProtocol->SetRecordingState (ProfileProtocol, MEMORY_PROFILE_RECORDING_DISABLE);
  }

  Size   = 0;
  Data   = NULL;
  Status = ProfileProtocol->GetData (
                              ProfileProtocol,
                              &Size,
                              Data
                              );
  if (Status != EFI_BUFFER_TOO_SMALL) {
    Print (L"UefiMemoryProfile: GetData - %r\n", Status);
    goto Done;
  }

  Data = AllocateZeroPool ((UINTN)Size);
  if (Data == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    Print (L"UefiMemoryProfile: AllocateZeroPool (0x%x) - %r\n", Size, Status);
    return Status;
  }

  Status = ProfileProtocol->GetData (
                              ProfileProtocol,
                              &Size,
                              Data
                              );
  if (EFI_ERROR (Status)) {
    Print (L"UefiMemoryProfile: GetData - %r\n", Status);
    goto Done;
  }

  Print (L"UefiMemoryProfileSize - 0x%x\n", Size);
  Print (L"======= UefiMemoryProfile begin =======\n");
  DumpMemoryProfile ((PHYSICAL_ADDRESS)(UINTN)Data, Size, FALSE);

  //
  // Dump summary information
  //
  MemoryProfileContextSummaryData = CreateContextSummaryData ((PHYSICAL_ADDRESS)(UINTN)Data, Size);
  if (MemoryProfileContextSummaryData != NULL) {
    DumpContextSummaryData (MemoryProfileContextSummaryData, FALSE);
    DestroyContextSummaryData (MemoryProfileContextSummaryData);
  }

  Print (L"======= UefiMemoryProfile end =======\n\n\n");

Done:
  if (Data != NULL) {
    FreePool (Data);
  }

  //
  // Restore recording state if needed.
  //
  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {
    ProfileProtocol->SetRecordingState (ProfileProtocol, MEMORY_PROFILE_RECORDING_ENABLE);
  }

  return Status;
}

/**
  Get and dump SMRAM profile data.

  @return EFI_SUCCESS   Get the SMRAM profile data successfully.
  @return other         Fail to get the SMRAM profile data.

**/
EFI_STATUS
GetSmramProfileData (
  VOID
  )
{
  EFI_STATUS                                          Status;
  UINTN                                               CommSize;
  UINT8                                               *CommBuffer;
  EFI_SMM_COMMUNICATE_HEADER                          *CommHeader;
  SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO            *CommGetProfileInfo;
  SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET  *CommGetProfileData;
  SMRAM_PROFILE_PARAMETER_RECORDING_STATE             *CommRecordingState;
  UINTN                                               ProfileSize;
  VOID                                                *ProfileBuffer;
  EFI_SMM_COMMUNICATION_PROTOCOL                      *SmmCommunication;
  UINTN                                               MinimalSizeNeeded;
  EDKII_PI_SMM_COMMUNICATION_REGION_TABLE             *PiSmmCommunicationRegionTable;
  UINT32                                              Index;
  EFI_MEMORY_DESCRIPTOR                               *Entry;
  VOID                                                *Buffer;
  UINTN                                               Size;
  UINTN                                               Offset;
  MEMORY_PROFILE_CONTEXT_SUMMARY_DATA                 *MemoryProfileContextSummaryData;
  BOOLEAN                                             RecordingState;

  ProfileBuffer = NULL;

  Status = gBS->LocateProtocol (&gEfiSmmCommunicationProtocolGuid, NULL, (VOID **)&SmmCommunication);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "SmramProfile: Locate SmmCommunication protocol - %r\n", Status));
    return Status;
  }

  MinimalSizeNeeded = sizeof (EFI_GUID) +
                      sizeof (UINTN) +
                      MAX (
                        sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO),
                        MAX (
                          sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET),
                          sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)
                          )
                        );
  MinimalSizeNeeded += MAX (
                         sizeof (MEMORY_PROFILE_CONTEXT),
                         MAX (
                           sizeof (MEMORY_PROFILE_DRIVER_INFO),
                           MAX (
                             sizeof (MEMORY_PROFILE_ALLOC_INFO),
                             MAX (
                               sizeof (MEMORY_PROFILE_DESCRIPTOR),
                               MAX (
                                 sizeof (MEMORY_PROFILE_FREE_MEMORY),
                                 sizeof (MEMORY_PROFILE_MEMORY_RANGE)
                                 )
                               )
                             )
                           )
                         );

  Status = EfiGetSystemConfigurationTable (
             &gEdkiiPiSmmCommunicationRegionTableGuid,
             (VOID **)&PiSmmCommunicationRegionTable
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "SmramProfile: Get PiSmmCommunicationRegionTable - %r\n", Status));
    return Status;
  }

  ASSERT (PiSmmCommunicationRegionTable != NULL);
  Entry = (EFI_MEMORY_DESCRIPTOR *)(PiSmmCommunicationRegionTable + 1);
  Size  = 0;
  for (Index = 0; Index < PiSmmCommunicationRegionTable->NumberOfEntries; Index++) {
    if (Entry->Type == EfiConventionalMemory) {
      Size = EFI_PAGES_TO_SIZE ((UINTN)Entry->NumberOfPages);
      if (Size >= MinimalSizeNeeded) {
        break;
      }
    }

    Entry = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)Entry + PiSmmCommunicationRegionTable->DescriptorSize);
  }

  ASSERT (Index < PiSmmCommunicationRegionTable->NumberOfEntries);
  CommBuffer = (UINT8 *)(UINTN)Entry->PhysicalStart;

  //
  // Set recording state if needed.
  //
  RecordingState = MEMORY_PROFILE_RECORDING_DISABLE;

  CommHeader = (EFI_SMM_COMMUNICATE_HEADER *)&CommBuffer[0];
  CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));
  CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE);

  CommRecordingState                      = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *)&CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];
  CommRecordingState->Header.Command      = SMRAM_PROFILE_COMMAND_GET_RECORDING_STATE;
  CommRecordingState->Header.DataLength   = sizeof (*CommRecordingState);
  CommRecordingState->Header.ReturnStatus = (UINT64)-1;
  CommRecordingState->RecordingState      = MEMORY_PROFILE_RECORDING_DISABLE;

  CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;
  Status   = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "SmramProfile: SmmCommunication - %r\n", Status));
    return Status;
  }

  if (CommRecordingState->Header.ReturnStatus != 0) {
    Print (L"SmramProfile: GetRecordingState - 0x%0x\n", CommRecordingState->Header.ReturnStatus);
    return EFI_SUCCESS;
  }

  RecordingState = CommRecordingState->RecordingState;
  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {
    CommHeader = (EFI_SMM_COMMUNICATE_HEADER *)&CommBuffer[0];
    CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));
    CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE);

    CommRecordingState                      = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *)&CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];
    CommRecordingState->Header.Command      = SMRAM_PROFILE_COMMAND_SET_RECORDING_STATE;
    CommRecordingState->Header.DataLength   = sizeof (*CommRecordingState);
    CommRecordingState->Header.ReturnStatus = (UINT64)-1;
    CommRecordingState->RecordingState      = MEMORY_PROFILE_RECORDING_DISABLE;

    CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;
    SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);
  }

  //
  // Get Size
  //
  CommHeader = (EFI_SMM_COMMUNICATE_HEADER *)&CommBuffer[0];
  CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));
  CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO);

  CommGetProfileInfo                      = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO *)&CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];
  CommGetProfileInfo->Header.Command      = SMRAM_PROFILE_COMMAND_GET_PROFILE_INFO;
  CommGetProfileInfo->Header.DataLength   = sizeof (*CommGetProfileInfo);
  CommGetProfileInfo->Header.ReturnStatus = (UINT64)-1;
  CommGetProfileInfo->ProfileSize         = 0;

  CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;
  Status   = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);
  ASSERT_EFI_ERROR (Status);

  if (CommGetProfileInfo->Header.ReturnStatus != 0) {
    Status = EFI_SUCCESS;
    Print (L"SmramProfile: GetProfileInfo - 0x%0x\n", CommGetProfileInfo->Header.ReturnStatus);
    goto Done;
  }

  ProfileSize = (UINTN)CommGetProfileInfo->ProfileSize;

  //
  // Get Data
  //
  ProfileBuffer = AllocateZeroPool (ProfileSize);
  if (ProfileBuffer == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    Print (L"SmramProfile: AllocateZeroPool (0x%x) for profile buffer - %r\n", ProfileSize, Status);
    goto Done;
  }

  CommHeader = (EFI_SMM_COMMUNICATE_HEADER *)&CommBuffer[0];
  CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));
  CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET);

  CommGetProfileData                      = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET *)&CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];
  CommGetProfileData->Header.Command      = SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA_BY_OFFSET;
  CommGetProfileData->Header.DataLength   = sizeof (*CommGetProfileData);
  CommGetProfileData->Header.ReturnStatus = (UINT64)-1;

  CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;
  Buffer   = (UINT8 *)CommHeader + CommSize;
  Size    -= CommSize;

  CommGetProfileData->ProfileBuffer = (PHYSICAL_ADDRESS)(UINTN)Buffer;
  CommGetProfileData->ProfileOffset = 0;
  while (CommGetProfileData->ProfileOffset < ProfileSize) {
    Offset = (UINTN)CommGetProfileData->ProfileOffset;
    if (Size <= (ProfileSize - CommGetProfileData->ProfileOffset)) {
      CommGetProfileData->ProfileSize = (UINT64)Size;
    } else {
      CommGetProfileData->ProfileSize = (UINT64)(ProfileSize - CommGetProfileData->ProfileOffset);
    }

    Status = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);
    ASSERT_EFI_ERROR (Status);

    if (CommGetProfileData->Header.ReturnStatus != 0) {
      Status = EFI_SUCCESS;
      Print (L"GetProfileData - 0x%x\n", CommGetProfileData->Header.ReturnStatus);
      goto Done;
    }

    CopyMem ((UINT8 *)ProfileBuffer + Offset, (VOID *)(UINTN)CommGetProfileData->ProfileBuffer, (UINTN)CommGetProfileData->ProfileSize);
  }

  Print (L"SmramProfileSize - 0x%x\n", ProfileSize);
  Print (L"======= SmramProfile begin =======\n");
  DumpMemoryProfile ((PHYSICAL_ADDRESS)(UINTN)ProfileBuffer, ProfileSize, TRUE);

  //
  // Dump summary information
  //
  MemoryProfileContextSummaryData = CreateContextSummaryData ((PHYSICAL_ADDRESS)(UINTN)ProfileBuffer, ProfileSize);
  if (MemoryProfileContextSummaryData != NULL) {
    DumpContextSummaryData (MemoryProfileContextSummaryData, TRUE);
    DestroyContextSummaryData (MemoryProfileContextSummaryData);
  }

  Print (L"======= SmramProfile end =======\n\n\n");

Done:
  if (ProfileBuffer != NULL) {
    FreePool (ProfileBuffer);
  }

  //
  // Restore recording state if needed.
  //
  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {
    CommHeader = (EFI_SMM_COMMUNICATE_HEADER *)&CommBuffer[0];
    CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));
    CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE);

    CommRecordingState                      = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *)&CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];
    CommRecordingState->Header.Command      = SMRAM_PROFILE_COMMAND_SET_RECORDING_STATE;
    CommRecordingState->Header.DataLength   = sizeof (*CommRecordingState);
    CommRecordingState->Header.ReturnStatus = (UINT64)-1;
    CommRecordingState->RecordingState      = MEMORY_PROFILE_RECORDING_ENABLE;

    CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;
    SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);
  }

  return Status;
}

/**
  The user Entry Point for Application. The user code starts with this function
  as the real entry point for the image goes into a library that calls this function.

  @param[in] ImageHandle    The firmware allocated handle for the EFI image.
  @param[in] SystemTable    A pointer to the EFI System Table.

  @retval EFI_SUCCESS       The entry point is executed successfully.
  @retval other             Some error occurs when executing this entry point.

**/
EFI_STATUS
EFIAPI
UefiMain (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS  Status;

  Status = GetUefiMemoryProfileData ();
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "GetUefiMemoryProfileData - %r\n", Status));
  }

  Status = GetSmramProfileData ();
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "GetSmramProfileData - %r\n", Status));
  }

  return EFI_SUCCESS;
}