refactor(context mgmt): refactor EL2 context save and restore functions

This patch splits the el2_sysregs_context_save/restore functions
into multiple functions based on features. This will allow us to
selectively save and restore EL2 context registers based on
features enabled for a particular configuration.

For now feature build flags are used to decide which registers
to save and restore. The long term plan is to dynamically check
for features that are enabled and then save/restore registers
accordingly. Splitting el2_sysregs_context_save/restore functions
into smaller assembly functions makes that task easier. For more
information please take a look at:
https://trustedfirmware-a.readthedocs.io/en/latest/design_documents/context_mgmt_rework.html

Signed-off-by: Zelalem Aweke <zelalem.aweke@arm.com>
Change-Id: I1819a9de8b70fa35c8f45568908025f790c4808c
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index e393493..da610d0 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2022, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -790,11 +790,47 @@
 	if ((security_state != SECURE) ||
 	    ((security_state == SECURE) && ((scr_el3 & SCR_EEL2_BIT) != 0U))) {
 		cpu_context_t *ctx;
+		el2_sysregs_t *el2_sysregs_ctx;
 
 		ctx = cm_get_context(security_state);
 		assert(ctx != NULL);
 
-		el2_sysregs_context_save(get_el2_sysregs_ctx(ctx));
+		el2_sysregs_ctx = get_el2_sysregs_ctx(ctx);
+
+		el2_sysregs_context_save_common(el2_sysregs_ctx);
+#if ENABLE_SPE_FOR_LOWER_ELS
+		el2_sysregs_context_save_spe(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_MTE_REGS
+		el2_sysregs_context_save_mte(el2_sysregs_ctx);
+#endif
+#if ENABLE_MPAM_FOR_LOWER_ELS
+		el2_sysregs_context_save_mpam(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_FGT
+		el2_sysregs_context_save_fgt(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_ECV
+		el2_sysregs_context_save_ecv(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_VHE
+		el2_sysregs_context_save_vhe(el2_sysregs_ctx);
+#endif
+#if RAS_EXTENSION
+		el2_sysregs_context_save_ras(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_NEVE_REGS
+		el2_sysregs_context_save_nv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_TRF_FOR_NS
+		el2_sysregs_context_save_trf(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_CSV2_2
+		el2_sysregs_context_save_csv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_HCX
+		el2_sysregs_context_save_hcx(el2_sysregs_ctx);
+#endif
 	}
 }
 
@@ -812,11 +848,47 @@
 	if ((security_state != SECURE) ||
 	    ((security_state == SECURE) && ((scr_el3 & SCR_EEL2_BIT) != 0U))) {
 		cpu_context_t *ctx;
+		el2_sysregs_t *el2_sysregs_ctx;
 
 		ctx = cm_get_context(security_state);
 		assert(ctx != NULL);
 
-		el2_sysregs_context_restore(get_el2_sysregs_ctx(ctx));
+		el2_sysregs_ctx = get_el2_sysregs_ctx(ctx);
+
+		el2_sysregs_context_restore_common(el2_sysregs_ctx);
+#if ENABLE_SPE_FOR_LOWER_ELS
+		el2_sysregs_context_restore_spe(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_MTE_REGS
+		el2_sysregs_context_restore_mte(el2_sysregs_ctx);
+#endif
+#if ENABLE_MPAM_FOR_LOWER_ELS
+		el2_sysregs_context_restore_mpam(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_FGT
+		el2_sysregs_context_restore_fgt(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_ECV
+		el2_sysregs_context_restore_ecv(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_VHE
+		el2_sysregs_context_restore_vhe(el2_sysregs_ctx);
+#endif
+#if RAS_EXTENSION
+		el2_sysregs_context_restore_ras(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_NEVE_REGS
+		el2_sysregs_context_restore_nv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_TRF_FOR_NS
+		el2_sysregs_context_restore_trf(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_CSV2_2
+		el2_sysregs_context_restore_csv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_HCX
+		el2_sysregs_context_restore_hcx(el2_sysregs_ctx);
+#endif
 	}
 }
 #endif /* CTX_INCLUDE_EL2_REGS */