[PATCH] UefiCpuPkg/Library/MpInitLib: Remove BSP index == 0 Assumption.

["Dong, Eric" <eric.dong@intel.com>]

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=2392

Current code implementation assumes BSP index is 0 at the begin.
This code change removes this assumption. It get BSP index from
the saved data structure if it existed.

Cc: Ray Ni <ray.ni@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Signed-off-by: Eric Dong <eric.dong@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/MpLib.c | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 6ec9b172b8..922c87b766 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -636,7 +636,7 @@ ApWakeupFunction (
       //   to initialize AP in InitConfig path.
       // NOTE: IDTR.BASE stored in CpuMpData->CpuData[0].VolatileRegisters points to a different IDT shared by all APs.
       //
-      RestoreVolatileRegisters (&CpuMpData->CpuData[0].VolatileRegisters, FALSE);
+      RestoreVolatileRegisters (&CpuMpData->CpuData[CpuMpData->BspNumber].VolatileRegisters, FALSE);
       InitializeApData (CpuMpData, ProcessorNumber, BistData, ApTopOfStack);
       ApStartupSignalBuffer = CpuMpData->CpuData[ProcessorNumber].StartupApSignal;
 
@@ -1615,6 +1615,7 @@ MpInitLibInitialize (
   UINTN                    ApResetVectorSize;
   UINTN                    BackupBufferAddr;
   UINTN                    ApIdtBase;
+  UINT64                   BspTopOfStack;
 
   OldCpuMpData = GetCpuMpDataFromGuidedHob ();
   if (OldCpuMpData == NULL) {
@@ -1677,7 +1678,7 @@ MpInitLibInitialize (
   CpuMpData->BackupBufferSize = ApResetVectorSize;
   CpuMpData->WakeupBuffer     = (UINTN) -1;
   CpuMpData->CpuCount         = 1;
-  CpuMpData->BspNumber        = 0;
+  CpuMpData->BspNumber        = OldCpuMpData != NULL ? OldCpuMpData->BspNumber : 0;
   CpuMpData->WaitEvent        = NULL;
   CpuMpData->SwitchBspFlag    = FALSE;
   CpuMpData->CpuData          = (CPU_AP_DATA *) (CpuMpData + 1);
@@ -1704,11 +1705,12 @@ MpInitLibInitialize (
   // Don't pass BSP's TR to APs to avoid AP init failure.
   //
   VolatileRegisters.Tr = 0;
-  CopyMem (&CpuMpData->CpuData[0].VolatileRegisters, &VolatileRegisters, sizeof (VolatileRegisters));
+  CopyMem (&CpuMpData->CpuData[CpuMpData->BspNumber].VolatileRegisters, &VolatileRegisters, sizeof (VolatileRegisters));
   //
   // Set BSP basic information
   //
-  InitializeApData (CpuMpData, 0, 0, CpuMpData->Buffer + ApStackSize);
+  BspTopOfStack = CpuMpData->Buffer + (CpuMpData->BspNumber + 1) * CpuMpData->CpuApStackSize;
+  InitializeApData (CpuMpData, CpuMpData->BspNumber, 0, BspTopOfStack);
   //
   // Save assembly code information
   //
-- 
2.23.0.windows.1