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
diff --git a/services/std_svc/psci/psci_main.c b/services/std_svc/psci/psci_main.c
index 1bcf216..c0866fb 100644
--- a/services/std_svc/psci/psci_main.c
+++ b/services/std_svc/psci/psci_main.c
@@ -90,23 +90,37 @@
 	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 @@
 	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;
 }