[PATCH 6/6] UefiCpuPkg/MpInitLib: Update AP information when BSP switched

[Jeff Fan <jeff.fan@intel.com>]

When BSP switched, we need to update some AP information. For example,
ApStartupSignalBuffer and ApTopOfStack.

Cc: Feng Tian <feng.tian@intel.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/MpLib.c | 3 +++
 UefiCpuPkg/Library/MpInitLib/MpLib.h | 1 +
 2 files changed, 4 insertions(+)

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 748c8e7..c74efb0 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -549,6 +549,8 @@ ApWakeupFunction (
             GetProcessorNumber (CpuMpData, &ProcessorNumber);
             CpuMpData->CpuData[ProcessorNumber].ApFunction = 0;
             CpuMpData->CpuData[ProcessorNumber].ApFunctionArgument = 0;
+            ApStartupSignalBuffer = CpuMpData->CpuData[ProcessorNumber].StartupApSignal;
+            CpuInfoInHob[ProcessorNumber].ApTopOfStack = CpuInfoInHob[CpuMpData->NewBspNumber].ApTopOfStack;
           } else {
             //
             // Re-get the CPU APICID and Initial APICID
@@ -1421,6 +1423,7 @@ SwitchBSPWorker (
   CpuMpData->BSPInfo.State = CPU_SWITCH_STATE_IDLE;
   CpuMpData->APInfo.State  = CPU_SWITCH_STATE_IDLE;
   CpuMpData->SwitchBspFlag = TRUE;
+  CpuMpData->NewBspNumber  = ProcessorNumber;
 
   //
   // Clear the BSP bit of MSR_IA32_APIC_BASE
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index f81c819..0ac777a 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -216,6 +216,7 @@ struct _CPU_MP_DATA {
   AP_INIT_STATE                  InitFlag;
   BOOLEAN                        X2ApicEnable;
   BOOLEAN                        SwitchBspFlag;
+  UINTN                          NewBspNumber;
   CPU_EXCHANGE_ROLE_INFO         BSPInfo;
   CPU_EXCHANGE_ROLE_INFO         APInfo;
   MTRR_SETTINGS                  MtrrTable;
-- 
2.9.3.windows.2