[Patch 2/2] UefiCpuPkg/PiSmmCpuDxeSmm: Refine code to avoid duplicated code.

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

Refine code to avoid duplicate code to set processor register.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Eric Dong <eric.dong@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c | 78 ++++++++++-----------------------------
 1 file changed, 20 insertions(+), 58 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c
index ae4b516..500a0e2 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c
@@ -208,18 +208,28 @@ Returns:
 
   This function programs registers for the calling processor.
 
-  @param  RegisterTable Pointer to register table of the running processor.
+  @param  RegisterTableList     Pointer to register table of the running processor.
 
 **/
 VOID
 SetProcessorRegister (
-  IN CPU_REGISTER_TABLE        *RegisterTable
+  IN CPU_REGISTER_TABLE        *RegisterTableList
   )
 {
   CPU_REGISTER_TABLE_ENTRY  *RegisterTableEntry;
   UINTN                     Index;
   UINTN                     Value;
   SPIN_LOCK                 *MsrSpinLock;
+  UINT32                    InitApicId;
+  CPU_REGISTER_TABLE        *RegisterTable;
+
+  InitApicId = GetInitialApicId ();
+  for (Index = 0; Index < mAcpiCpuData.NumberOfCpus; Index++) {
+    if (RegisterTableList[Index].InitialApicId == InitApicId) {
+      RegisterTable =  &RegisterTableList[Index];
+      break;
+    }
+  }
 
   //
   // Traverse Register Table of this logical processor
@@ -347,8 +357,6 @@ SetProcessorRegister (
   }
 }
 
-
-
 /**
   AP initialization before then after SMBASE relocation in the S3 boot path.
 **/
@@ -357,26 +365,12 @@ MPRendezvousProcedure (
   VOID
   )
 {
-  CPU_REGISTER_TABLE         *RegisterTableList;
-  UINT32                     InitApicId;
-  UINTN                      Index;
   UINTN                      TopOfStack;
   UINT8                      Stack[128];
 
   LoadMtrrData (mAcpiCpuData.MtrrTable);
 
-  //
-  // Find processor number for this CPU.
-  //
-  RegisterTableList = (CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.PreSmmInitRegisterTable;
-  InitApicId = GetInitialApicId ();
-  for (Index = 0; Index < mAcpiCpuData.NumberOfCpus; Index++) {
-    if (RegisterTableList[Index].InitialApicId == InitApicId) {
-      SetProcessorRegister (&RegisterTableList[Index]);
-      break;
-    }
-  }
-
+  SetProcessorRegister ((CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.PreSmmInitRegisterTable);
 
   //
   // Count down the number with lock mechanism.
@@ -393,14 +387,7 @@ MPRendezvousProcedure (
   ProgramVirtualWireMode ();
   DisableLvtInterrupts ();
 
-  RegisterTableList = (CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.RegisterTable;
-  InitApicId = GetInitialApicId ();
-  for (Index = 0; Index < mAcpiCpuData.NumberOfCpus; Index++) {
-    if (RegisterTableList[Index].InitialApicId == InitApicId) {
-      SetProcessorRegister (&RegisterTableList[Index]);
-      break;
-    }
-  }
+  SetProcessorRegister ((CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.RegisterTable);
 
   //
   // Place AP into the safe code, count down the number with lock mechanism in the safe code.
@@ -475,27 +462,13 @@ PrepareApStartupVector (
 
 **/
 VOID
-EarlyInitializeCpu (
+InitializeCpuBeforeRebase (
   VOID
   )
 {
-  CPU_REGISTER_TABLE         *RegisterTableList;
-  UINT32                     InitApicId;
-  UINTN                      Index;
-
   LoadMtrrData (mAcpiCpuData.MtrrTable);
 
-  //
-  // Find processor number for this CPU.
-  //
-  RegisterTableList = (CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.PreSmmInitRegisterTable;
-  InitApicId = GetInitialApicId ();
-  for (Index = 0; Index < mAcpiCpuData.NumberOfCpus; Index++) {
-    if (RegisterTableList[Index].InitialApicId == InitApicId) {
-      SetProcessorRegister (&RegisterTableList[Index]);
-      break;
-    }
-  }
+  SetProcessorRegister ((CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.PreSmmInitRegisterTable);
 
   ProgramVirtualWireMode ();
 
@@ -527,22 +500,11 @@ EarlyInitializeCpu (
 
 **/
 VOID
-InitializeCpu (
+InitializeCpuAfterRebase (
   VOID
   )
 {
-  CPU_REGISTER_TABLE         *RegisterTableList;
-  UINT32                     InitApicId;
-  UINTN                      Index;
-
-  RegisterTableList = (CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.RegisterTable;
-  InitApicId = GetInitialApicId ();
-  for (Index = 0; Index < mAcpiCpuData.NumberOfCpus; Index++) {
-    if (RegisterTableList[Index].InitialApicId == InitApicId) {
-      SetProcessorRegister (&RegisterTableList[Index]);
-      break;
-    }
-  }
+  SetProcessorRegister ((CPU_REGISTER_TABLE *) (UINTN) mAcpiCpuData.RegisterTable);
 
   mNumberToFinish = mAcpiCpuData.NumberOfCpus - 1;
 
@@ -660,7 +622,7 @@ SmmRestoreCpu (
     //
     // First time microcode load and restore MTRRs
     //
-    EarlyInitializeCpu ();
+    InitializeCpuBeforeRebase ();
   }
 
   //
@@ -675,7 +637,7 @@ SmmRestoreCpu (
     //
     // Restore MSRs for BSP and all APs
     //
-    InitializeCpu ();
+    InitializeCpuAfterRebase ();
   }
 
   //
-- 
2.7.0.windows.1