From 98f17b671d3b04e10b2797d7cb64d1b9442bb80b Mon Sep 17 00:00:00 2001
From: vit9696 <vit9696@users.noreply.github.com>
Date: Fri, 3 Apr 2020 02:30:24 +0300
Subject: [PATCH] OcMemoryLib: Add counting split descriptors

---
 Include/Library/OcMemoryLib.h   | 11 ++++++++++-
 Library/OcMemoryLib/MemoryMap.c | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/Include/Library/OcMemoryLib.h b/Include/Library/OcMemoryLib.h
index c2e0bbae..f9fe9012 100644
--- a/Include/Library/OcMemoryLib.h
+++ b/Include/Library/OcMemoryLib.h
@@ -24,7 +24,6 @@
 #define PREV_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \
   ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) - (Size)))
 
-
 /**
   Get last descriptor address.
   It is assumed that the descriptor contains pages.
@@ -266,6 +265,16 @@ OcUpdateAttributes (
   IN UINT64                DropAttributes
   );
 
+/**
+  Count upper bound of split runtime descriptors.
+
+  @retval amount of runtime descriptors.
+**/
+UINTN
+OcCountSplitDescritptors (
+  VOID
+  );
+
 /**
   Return pointer to PML4 table in PageTable and PWT and PCD flags in Flags.
 
diff --git a/Library/OcMemoryLib/MemoryMap.c b/Library/OcMemoryLib/MemoryMap.c
index 1297df18..9b12acef 100644
--- a/Library/OcMemoryLib/MemoryMap.c
+++ b/Library/OcMemoryLib/MemoryMap.c
@@ -540,3 +540,36 @@ OcUpdateAttributes (
 
   return EFI_UNSUPPORTED;
 }
+
+UINTN
+OcCountSplitDescritptors (
+  VOID
+  )
+{
+  UINTN                             Index;
+  UINTN                             DescriptorCount;
+  CONST EFI_MEMORY_ATTRIBUTES_TABLE *MemoryAttributesTable;
+  EFI_MEMORY_DESCRIPTOR             *MemoryAttributesEntry;
+
+  DescriptorCount = 0;
+
+  for (Index = 0; Index < gST->NumberOfTableEntries; ++Index) {
+    if (CompareGuid (&gST->ConfigurationTable[Index].VendorGuid, &gEfiMemoryAttributesTableGuid)) {
+      MemoryAttributesTable = (CONST EFI_MEMORY_ATTRIBUTES_TABLE *) gST->ConfigurationTable[Index].VendorTable;
+      MemoryAttributesEntry = (EFI_MEMORY_DESCRIPTOR *) (MemoryAttributesTable + 1);
+      for (Index = 0; Index < MemoryAttributesTable->NumberOfEntries; ++Index) {
+        if (MemoryAttributesEntry->Type == EfiRuntimeServicesCode
+          || MemoryAttributesEntry->Type == EfiRuntimeServicesData) {
+          ++DescriptorCount;
+        }
+
+        MemoryAttributesEntry = NEXT_MEMORY_DESCRIPTOR (
+          MemoryAttributesEntry,
+          MemoryAttributesTable->DescriptorSize
+          );
+      }
+    }
+  }
+
+  return DescriptorCount;
+}