Fix broken standby state implementation in PSCI

This patch fixes the broken support for entry into standby states
introduced under commit-id 'd118f9f864' (tf-issues#94). Upon exit from
the platform defined standby state instead of returning to the caller
of the SMC, execution would get stuck in the wfi instruction meant for
entering a power down state. This patch ensures that exit from a
standby state and entry into a power down state do not interfere with
each other.

Fixes ARM-software/tf-issues#154

Change-Id: I56e5df353368e44d6eefc94ffedefe21929f5cfe

2 files changed, 46 insertions, 20 deletions

diff --git a/services/std_svc/psci/psci_entry.S b/services/std_svc/psci/psci_entry.S
index 256c538d..3d0181a8 100644
--- a/services/std_svc/psci/psci_entry.S
+++ b/services/std_svc/psci/psci_entry.S
@@ -37,6 +37,7 @@
 	.globl	psci_aff_suspend_finish_entry
 	.globl	__psci_cpu_off
 	.globl	__psci_cpu_suspend
+	.globl	psci_power_down_wfi
 
 	/* -----------------------------------------------------
 	 * This cpu has been physically powered up. Depending
@@ -120,9 +121,6 @@ func __psci_cpu_off
 	mrs	x0, mpidr_el1
 	bl	platform_set_coherent_stack
 	bl	psci_cpu_off
-	mov	x1, #PSCI_E_SUCCESS
-	cmp	x0, x1
-	b.eq	final_wfi
 	mov	sp, x19
 	ldp	x19, x20, [sp,#0]
 	add	sp, sp, #0x10
@@ -144,9 +142,6 @@ func __psci_cpu_suspend
 	mov	x1, x21
 	mov	x2, x22
 	bl	psci_cpu_suspend
-	mov	x1, #PSCI_E_SUCCESS
-	cmp	x0, x1
-	b.eq	final_wfi
 	mov	sp, x19
 	ldp	x21, x22, [sp,#0x10]
 	ldp	x19, x20, [sp,#0]
@@ -154,7 +149,16 @@ func __psci_cpu_suspend
 	func_epilogue
 	ret
 
-func final_wfi
+	/* --------------------------------------------
+	 * This function is called to indicate to the
+	 * power controller that it is safe to power
+	 * down this cpu. It should not exit the wfi
+	 * and will be released from reset upon power
+	 * up. 'wfi_spill' is used to catch erroneous
+	 * exits from wfi.
+	 * --------------------------------------------
+	 */
+func psci_power_down_wfi
 	dsb	sy		// ensure write buffer empty
 	wfi
 wfi_spill:
diff --git a/services/std_svc/psci/psci_main.c b/services/std_svc/psci/psci_main.c
index 1bcf2166..c0866fb6 100644
--- a/services/std_svc/psci/psci_main.c
+++ b/services/std_svc/psci/psci_main.c
@@ -90,23 +90,37 @@ int psci_cpu_suspend(unsigned int power_state,
 	if (target_afflvl > MPIDR_MAX_AFFLVL)
 		return PSCI_E_INVALID_PARAMS;
 
+	/* Determine the 'state type' in the 'power_state' parameter */
 	pstate_type = psci_get_pstate_type(power_state);
+
+	/*
+	 * Ensure that we have a platform specific handler for entering
+	 * a standby state.
+	 */
 	if (pstate_type == PSTATE_TYPE_STANDBY) {
-		if  (psci_plat_pm_ops->affinst_standby)
-			rc = psci_plat_pm_ops->affinst_standby(power_state);
-		else
+		if  (!psci_plat_pm_ops->affinst_standby)
 			return PSCI_E_INVALID_PARAMS;
-	} else {
-		mpidr = read_mpidr();
-		rc = psci_afflvl_suspend(mpidr,
-					 entrypoint,
-					 context_id,
-					 power_state,
-					 MPIDR_AFFLVL0,
-					 target_afflvl);
+
+		rc = psci_plat_pm_ops->affinst_standby(power_state);
+		assert(rc == PSCI_E_INVALID_PARAMS || rc == PSCI_E_SUCCESS);
+		return rc;
 	}
 
-	assert(rc == PSCI_E_INVALID_PARAMS || rc == PSCI_E_SUCCESS);
+	/*
+	 * Do what is needed to enter the power down state. Upon success,
+	 * enter the final wfi which will power down this cpu else return
+	 * an error.
+	 */
+	mpidr = read_mpidr();
+	rc = psci_afflvl_suspend(mpidr,
+				 entrypoint,
+				 context_id,
+				 power_state,
+				 MPIDR_AFFLVL0,
+				 target_afflvl);
+	if (rc == PSCI_E_SUCCESS)
+		psci_power_down_wfi();
+	assert(rc == PSCI_E_INVALID_PARAMS);
 	return rc;
 }
 
@@ -127,10 +141,18 @@ int psci_cpu_off(void)
 	rc = psci_afflvl_off(mpidr, MPIDR_AFFLVL0, target_afflvl);
 
 	/*
+	 * Check if all actions needed to safely power down this cpu have
+	 * successfully completed. Enter a wfi loop which will allow the
+	 * power controller to physically power down this cpu.
+	 */
+	if (rc == PSCI_E_SUCCESS)
+		psci_power_down_wfi();
+
+	/*
 	 * The only error cpu_off can return is E_DENIED. So check if that's
 	 * indeed the case.
 	 */
-	assert (rc == PSCI_E_SUCCESS || rc == PSCI_E_DENIED);
+	assert (rc == PSCI_E_DENIED);
 
 	return rc;
 }