1 files changed, 98 insertions, 75 deletions
diff --git a/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c b/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c
index f6bb35a5fe..9692e2f6bf 100644
--- a/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c
+++ b/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c
@@ -41,7 +41,7 @@ DiscoverPeimsAndOrderWithApriori (
   UINTN                               PeimCount;
   EFI_GUID                            *Guid;
   EFI_PEI_FILE_HANDLE                 *TempFileHandles;
-  EFI_GUID                            *FileGuid;
+  EFI_GUID                            *TempFileGuid;
   EFI_PEI_FIRMWARE_VOLUME_PPI         *FvPpi;
   EFI_FV_FILE_INFO                    FileInfo;
 
@@ -51,38 +51,81 @@ DiscoverPeimsAndOrderWithApriori (
   // Walk the FV and find all the PEIMs and the Apriori file.
   //
   AprioriFileHandle = NULL;
-  Private->CurrentFvFileHandles[0] = NULL;
+  Private->CurrentFvFileHandles = NULL;
   Guid = NULL;
-  FileHandle = NULL;
-  TempFileHandles = Private->FileHandles;
-  FileGuid        = Private->FileGuid;
 
   //
-  // If the current Fv has been scanned, directly get its cachable record.
+  // If the current Fv has been scanned, directly get its cached records.
   //
-  if (Private->Fv[Private->CurrentPeimFvCount].ScanFv) {
-    CopyMem (Private->CurrentFvFileHandles, Private->Fv[Private->CurrentPeimFvCount].FvFileHandles, sizeof (EFI_PEI_FILE_HANDLE) * PcdGet32 (PcdPeiCoreMaxPeimPerFv));
+  if (CoreFileHandle->ScanFv) {
+    Private->CurrentFvFileHandles = CoreFileHandle->FvFileHandles;
     return;
   }
 
+  TempFileHandles = Private->TempFileHandles;
+  TempFileGuid    = Private->TempFileGuid;
+
   //
-  // Go ahead to scan this Fv, and cache FileHandles within it.
+  // Go ahead to scan this Fv, get PeimCount and cache FileHandles within it to TempFileHandles.
   //
-  Status = EFI_NOT_FOUND;
-  for (PeimCount = 0; PeimCount <= PcdGet32 (PcdPeiCoreMaxPeimPerFv); PeimCount++) {
+  PeimCount = 0;
+  FileHandle = NULL;
+  do {
     Status = FvPpi->FindFileByType (FvPpi, PEI_CORE_INTERNAL_FFS_FILE_DISPATCH_TYPE, CoreFileHandle->FvHandle, &FileHandle);
-    if (Status != EFI_SUCCESS || PeimCount == PcdGet32 (PcdPeiCoreMaxPeimPerFv)) {
-      break;
+    if (!EFI_ERROR (Status)) {
+      if (PeimCount >= Private->TempPeimCount) {
+        //
+        // Run out of room, grow the buffer.
+        //
+        TempFileHandles = AllocatePool (
+                            sizeof (EFI_PEI_FILE_HANDLE) * (Private->TempPeimCount + TEMP_FILE_GROWTH_STEP));
+        ASSERT (TempFileHandles != NULL);
+        CopyMem (
+          TempFileHandles,
+          Private->TempFileHandles,
+          sizeof (EFI_PEI_FILE_HANDLE) * Private->TempPeimCount
+          );
+        Private->TempFileHandles = TempFileHandles;
+        TempFileGuid = AllocatePool (
+                         sizeof (EFI_GUID) * (Private->TempPeimCount + TEMP_FILE_GROWTH_STEP));
+        ASSERT (TempFileGuid != NULL);
+        CopyMem (
+          TempFileGuid,
+          Private->TempFileGuid,
+          sizeof (EFI_GUID) * Private->TempPeimCount
+          );
+        Private->TempFileGuid = TempFileGuid;
+        Private->TempPeimCount = Private->TempPeimCount + TEMP_FILE_GROWTH_STEP;
+      }
+
+      TempFileHandles[PeimCount++] = FileHandle;
     }
+  } while (!EFI_ERROR (Status));
+
+  DEBUG ((
+    DEBUG_INFO,
+    "%a(): Found 0x%x PEI FFS files in the %dth FV\n",
+    __FUNCTION__,
+    PeimCount,
+    Private->CurrentPeimFvCount
+    ));
 
-    Private->CurrentFvFileHandles[PeimCount] = FileHandle;
+  if (PeimCount == 0) {
+    //
+    // No PEIM FFS file is found, set ScanFv flag and return.
+    //
+    CoreFileHandle->ScanFv = TRUE;
+    return;
   }
 
   //
-  // Check whether the count of files exceeds the max support files in a FV image
-  // If more files are required in a FV image, PcdPeiCoreMaxPeimPerFv can be set to a larger value in DSC file.
+  // Record PeimCount, allocate buffer for PeimState and FvFileHandles.
   //
-  ASSERT ((Status != EFI_SUCCESS) || (PeimCount < PcdGet32 (PcdPeiCoreMaxPeimPerFv)));
+  CoreFileHandle->PeimCount = PeimCount;
+  CoreFileHandle->PeimState = AllocateZeroPool (sizeof (UINT8) * PeimCount);
+  ASSERT (CoreFileHandle->PeimState != NULL);
+  CoreFileHandle->FvFileHandles = AllocateZeroPool (sizeof (EFI_PEI_FILE_HANDLE) * PeimCount);
+  ASSERT (CoreFileHandle->FvFileHandles != NULL);
 
   //
   // Get Apriori File handle
@@ -96,7 +139,7 @@ DiscoverPeimsAndOrderWithApriori (
     Status = FvPpi->FindSectionByType (FvPpi, EFI_SECTION_RAW, AprioriFileHandle, (VOID **) &Apriori);
     if (!EFI_ERROR (Status)) {
       //
-      // Calculate the number of PEIMs in the A Priori list
+      // Calculate the number of PEIMs in the Apriori file
       //
       Status = FvPpi->GetFileInfo (FvPpi, AprioriFileHandle, &FileInfo);
       ASSERT_EFI_ERROR (Status);
@@ -113,71 +156,55 @@ DiscoverPeimsAndOrderWithApriori (
         // Make an array of file name guids that matches the FileHandle array so we can convert
         // quickly from file name to file handle
         //
-        Status = FvPpi->GetFileInfo (FvPpi, Private->CurrentFvFileHandles[Index], &FileInfo);
-        CopyMem (&FileGuid[Index], &FileInfo.FileName, sizeof(EFI_GUID));
+        Status = FvPpi->GetFileInfo (FvPpi, TempFileHandles[Index], &FileInfo);
+        ASSERT_EFI_ERROR (Status);
+        CopyMem (&TempFileGuid[Index], &FileInfo.FileName, sizeof(EFI_GUID));
       }
 
       //
-      // Walk through FileGuid array to find out who is invalid PEIM guid in Apriori file.
-      // Add available PEIMs in Apriori file into TempFileHandles array at first.
+      // Walk through TempFileGuid array to find out who is invalid PEIM guid in Apriori file.
+      // Add available PEIMs in Apriori file into FvFileHandles array.
       //
-      Index2 = 0;
-      for (Index = 0; Index2 < Private->AprioriCount; Index++) {
-        while (Index2 < Private->AprioriCount) {
-          Guid = ScanGuid (FileGuid, PeimCount * sizeof (EFI_GUID), &Apriori[Index2++]);
-          if (Guid != NULL) {
-            break;
-          }
-        }
-        if (Guid == NULL) {
-          break;
-        }
-        PeimIndex = ((UINTN)Guid - (UINTN)&FileGuid[0])/sizeof (EFI_GUID);
-        TempFileHandles[Index] = Private->CurrentFvFileHandles[PeimIndex];
+      Index = 0;
+      for (Index2 = 0; Index2 < Private->AprioriCount; Index2++) {
+        Guid = ScanGuid (TempFileGuid, PeimCount * sizeof (EFI_GUID), &Apriori[Index2]);
+        if (Guid != NULL) {
+          PeimIndex = ((UINTN)Guid - (UINTN)&TempFileGuid[0])/sizeof (EFI_GUID);
+          CoreFileHandle->FvFileHandles[Index++] = TempFileHandles[PeimIndex];
 
-        //
-        // Since we have copied the file handle we can remove it from this list.
-        //
-        Private->CurrentFvFileHandles[PeimIndex] = NULL;
+          //
+          // Since we have copied the file handle we can remove it from this list.
+          //
+          TempFileHandles[PeimIndex] = NULL;
+        }
       }
 
       //
-      // Update valid Aprioricount
+      // Update valid AprioriCount
       //
       Private->AprioriCount = Index;
 
       //
       // Add in any PEIMs not in the Apriori file
       //
-      for (;Index < PeimCount; Index++) {
-        for (Index2 = 0; Index2 < PeimCount; Index2++) {
-          if (Private->CurrentFvFileHandles[Index2] != NULL) {
-            TempFileHandles[Index] = Private->CurrentFvFileHandles[Index2];
-            Private->CurrentFvFileHandles[Index2] = NULL;
-            break;
-          }
+      for (Index2 = 0; Index2 < PeimCount; Index2++) {
+        if (TempFileHandles[Index2] != NULL) {
+          CoreFileHandle->FvFileHandles[Index++] = TempFileHandles[Index2];
+          TempFileHandles[Index2] = NULL;
         }
       }
-      //
-      //Index the end of array contains re-range Pei moudle.
-      //
-      TempFileHandles[Index] = NULL;
-
-      //
-      // Private->CurrentFvFileHandles is currently in PEIM in the FV order.
-      // We need to update it to start with files in the A Priori list and
-      // then the remaining files in PEIM order.
-      //
-      CopyMem (Private->CurrentFvFileHandles, TempFileHandles, sizeof (EFI_PEI_FILE_HANDLE) * PcdGet32 (PcdPeiCoreMaxPeimPerFv));
+      ASSERT (Index == PeimCount);
     }
+  } else {
+    CopyMem (CoreFileHandle->FvFileHandles, TempFileHandles, sizeof (EFI_PEI_FILE_HANDLE) * PeimCount);
   }
+
   //
-  // Cache the current Fv File Handle. So that we don't have to scan the Fv again.
-  // Instead, we can retrieve the file handles within this Fv from cachable data.
+  // The current Fv File Handles have been cached. So that we don't have to scan the Fv again.
+  // Instead, we can retrieve the file handles within this Fv from cached records.
   //
-  Private->Fv[Private->CurrentPeimFvCount].ScanFv = TRUE;
-  CopyMem (Private->Fv[Private->CurrentPeimFvCount].FvFileHandles, Private->CurrentFvFileHandles, sizeof (EFI_PEI_FILE_HANDLE) * PcdGet32 (PcdPeiCoreMaxPeimPerFv));
-
+  CoreFileHandle->ScanFv = TRUE;
+  Private->CurrentFvFileHandles = CoreFileHandle->FvFileHandles;
 }
 
 //
@@ -977,7 +1004,7 @@ PeiDispatcher (
     SaveCurrentFileHandle =  Private->CurrentFileHandle;
 
     for (Index1 = 0; Index1 < Private->FvCount; Index1++) {
-      for (Index2 = 0; (Index2 < PcdGet32 (PcdPeiCoreMaxPeimPerFv)) && (Private->Fv[Index1].FvFileHandles[Index2] != NULL); Index2++) {
+      for (Index2 = 0; Index2 < Private->Fv[Index1].PeimCount; Index2++) {
         if (Private->Fv[Index1].PeimState[Index2] == PEIM_STATE_REGISTER_FOR_SHADOW) {
           PeimFileHandle = Private->Fv[Index1].FvFileHandles[Index2];
           Private->CurrentFileHandle   = PeimFileHandle;
@@ -1063,7 +1090,7 @@ PeiDispatcher (
       // Start to dispatch all modules within the current Fv.
       //
       for (PeimCount = Private->CurrentPeimCount;
-           (PeimCount < PcdGet32 (PcdPeiCoreMaxPeimPerFv)) && (Private->CurrentFvFileHandles[PeimCount] != NULL);
+           PeimCount < Private->Fv[FvCount].PeimCount;
            PeimCount++) {
         Private->CurrentPeimCount  = PeimCount;
         PeimFileHandle = Private->CurrentFileHandle = Private->CurrentFvFileHandles[PeimCount];
@@ -1207,21 +1234,17 @@ PeiDispatcher (
       }
 
       //
-      // We set to NULL here to optimize the 2nd entry to this routine after
-      //  memory is found. This reprevents rescanning of the FV. We set to
-      //  NULL here so we start at the begining of the next FV
+      // Before walking through the next FV, we should set them to NULL/0 to
+      // start at the begining of the next FV.
       //
       Private->CurrentFileHandle = NULL;
       Private->CurrentPeimCount = 0;
-      //
-      // Before walking through the next FV,Private->CurrentFvFileHandles[]should set to NULL
-      //
-      SetMem (Private->CurrentFvFileHandles, sizeof (EFI_PEI_FILE_HANDLE) * PcdGet32 (PcdPeiCoreMaxPeimPerFv), 0);
+      Private->CurrentFvFileHandles = NULL;
     }
 
     //
-    // Before making another pass, we should set Private->CurrentPeimFvCount =0 to go
-    // through all the FV.
+    // Before making another pass, we should set it to 0 to
+    // go through all the FVs.
     //
     Private->CurrentPeimFvCount = 0;
 
@@ -1300,7 +1323,7 @@ DepexSatisfied (
 
   if (PeimCount < Private->AprioriCount) {
     //
-    // If its in the A priori file then we set Depex to TRUE
+    // If it's in the Apriori file then we set Depex to TRUE
     //
     DEBUG ((DEBUG_DISPATCH, "  RESULT = TRUE (Apriori)\n"));
     return TRUE;