diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index 70c50aa..786dcf5 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2019, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -104,25 +104,30 @@
 #define CTX_SPSR_FIQ		U(0xd8)
 #define CTX_DACR32_EL2		U(0xe0)
 #define CTX_IFSR32_EL2		U(0xe8)
-#define CTX_TIMER_SYSREGS_OFF	U(0xf0) /* Align to the next 16 byte boundary */
+#define CTX_AARCH32_END		U(0xf0) /* Align to the next 16 byte boundary */
 #else
-#define CTX_TIMER_SYSREGS_OFF	U(0xc0)  /* Align to the next 16 byte boundary */
-#endif /* __CTX_INCLUDE_AARCH32_REGS__ */
+#define CTX_AARCH32_END		U(0xc0)  /* Align to the next 16 byte boundary */
+#endif /* CTX_INCLUDE_AARCH32_REGS */
 
 /*
  * If the timer registers aren't saved and restored, we don't have to reserve
  * space for them in the context
  */
 #if NS_TIMER_SWITCH
-#define CTX_CNTP_CTL_EL0	(CTX_TIMER_SYSREGS_OFF + U(0x0))
-#define CTX_CNTP_CVAL_EL0	(CTX_TIMER_SYSREGS_OFF + U(0x8))
-#define CTX_CNTV_CTL_EL0	(CTX_TIMER_SYSREGS_OFF + U(0x10))
-#define CTX_CNTV_CVAL_EL0	(CTX_TIMER_SYSREGS_OFF + U(0x18))
-#define CTX_CNTKCTL_EL1		(CTX_TIMER_SYSREGS_OFF + U(0x20))
-#define CTX_SYSREGS_END		(CTX_TIMER_SYSREGS_OFF + U(0x30)) /* Align to the next 16 byte boundary */
+#define CTX_CNTP_CTL_EL0	(CTX_AARCH32_END + U(0x0))
+#define CTX_CNTP_CVAL_EL0	(CTX_AARCH32_END + U(0x8))
+#define CTX_CNTV_CTL_EL0	(CTX_AARCH32_END + U(0x10))
+#define CTX_CNTV_CVAL_EL0	(CTX_AARCH32_END + U(0x18))
+#define CTX_CNTKCTL_EL1		(CTX_AARCH32_END + U(0x20))
+#define CTX_TIMER_SYSREGS_END	(CTX_AARCH32_END + U(0x30)) /* Align to the next 16 byte boundary */
 #else
-#define CTX_SYSREGS_END		CTX_TIMER_SYSREGS_OFF
-#endif /* __NS_TIMER_SWITCH__ */
+#define CTX_TIMER_SYSREGS_END	CTX_AARCH32_END
+#endif /* NS_TIMER_SWITCH */
+
+/*
+ * End of system registers.
+ */
+#define CTX_SYSREGS_END		CTX_TIMER_SYSREGS_END
 
 /*******************************************************************************
  * Constants that allow assembler code to access members of and the 'fp_regs'
@@ -174,6 +179,9 @@
 #define CTX_FPREGS_END		U(0)
 #endif
 
+/*******************************************************************************
+ * Registers related to CVE-2018-3639
+ ******************************************************************************/
 #define CTX_CVE_2018_3639_OFFSET	(CTX_FPREGS_OFFSET + CTX_FPREGS_END)
 #define CTX_CVE_2018_3639_DISABLE	U(0)
 #define CTX_CVE_2018_3639_END		U(0x10) /* Align to the next 16 byte boundary */
@@ -182,8 +190,6 @@
 
 #include <stdint.h>
 
-#include <platform_def.h>	/* for CACHE_WRITEBACK_GRANULE */
-
 #include <lib/cassert.h>
 
 /*
@@ -200,7 +206,7 @@
 #define CTX_GPREG_ALL		(CTX_GPREGS_END >> DWORD_SHIFT)
 #define CTX_SYSREG_ALL		(CTX_SYSREGS_END >> DWORD_SHIFT)
 #if CTX_INCLUDE_FPREGS
-#define CTX_FPREG_ALL		(CTX_FPREGS_END >> DWORD_SHIFT)
+# define CTX_FPREG_ALL		(CTX_FPREGS_END >> DWORD_SHIFT)
 #endif
 #define CTX_EL3STATE_ALL	(CTX_EL3STATE_END >> DWORD_SHIFT)
 #define CTX_CVE_2018_3639_ALL	(CTX_CVE_2018_3639_END >> DWORD_SHIFT)
@@ -269,7 +275,7 @@
 /* Macros to access members of the 'cpu_context_t' structure */
 #define get_el3state_ctx(h)	(&((cpu_context_t *) h)->el3state_ctx)
 #if CTX_INCLUDE_FPREGS
-#define get_fpregs_ctx(h)	(&((cpu_context_t *) h)->fpregs_ctx)
+# define get_fpregs_ctx(h)	(&((cpu_context_t *) h)->fpregs_ctx)
 #endif
 #define get_sysregs_ctx(h)	(&((cpu_context_t *) h)->sysregs_ctx)
 #define get_gpregs_ctx(h)	(&((cpu_context_t *) h)->gpregs_ctx)
@@ -339,14 +345,6 @@
 void fpregs_context_restore(fp_regs_t *regs);
 #endif
 
-
-#undef CTX_SYSREG_ALL
-#if CTX_INCLUDE_FPREGS
-#undef CTX_FPREG_ALL
-#endif
-#undef CTX_GPREG_ALL
-#undef CTX_EL3STATE_ALL
-
 #endif /* __ASSEMBLY__ */
 
 #endif /* CONTEXT_H */
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 707e6db..1f8e23d 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2019, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -332,9 +332,10 @@
 	ret
 endfunc save_gp_registers
 
-/*
+/* -----------------------------------------------------
  * This function restores all general purpose registers except x30 from the
  * CPU context. x30 register must be explicitly restored by the caller.
+ * -----------------------------------------------------
  */
 func restore_gp_registers
 	ldp	x0, x1, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X0]
@@ -357,9 +358,10 @@
 	ret
 endfunc restore_gp_registers
 
-/*
+/* -----------------------------------------------------
  * Restore general purpose registers (including x30), and exit EL3 via. ERET to
  * a lower exception level.
+ * -----------------------------------------------------
  */
 func restore_gp_registers_eret
 	bl	restore_gp_registers
@@ -377,12 +379,12 @@
 	eret
 endfunc	restore_gp_registers_eret
 
-	/* -----------------------------------------------------
-	 * This routine assumes that the SP_EL3 is pointing to
-	 * a valid context structure from where the gp regs and
-	 * other special registers can be retrieved.
-	 * -----------------------------------------------------
-	 */
+/* -----------------------------------------------------
+ * This routine assumes that the SP_EL3 is pointing to
+ * a valid context structure from where the gp regs and
+ * other special registers can be retrieved.
+ * -----------------------------------------------------
+ */
 func el3_exit
 	/* -----------------------------------------------------
 	 * Save the current SP_EL0 i.e. the EL3 runtime stack
@@ -410,9 +412,9 @@
 	cmp	x17, xzr
 	beq	1f
 	blr	x17
+1:
 #endif
 
-1:
 	/* Restore saved general purpose registers and return */
 	b	restore_gp_registers_eret
 endfunc el3_exit