feat(psci): add a helper function to ensure that non-boot PEs are offline

Introduce a helper function that ensures that non-boot PEs are offline.
This function will be used by DRTM implementation to ensure that system
is running with only single PE.

Signed-off-by: Manish V Badarkhe <manish.badarkhe@arm.com>
Signed-off-by: Lucian Paul-Trifu <lucian.paultrifu@gmail.com>
Signed-off-by: Manish Pandey <manish.pandey2@arm.com>
Change-Id: I521ebefa49297026b02554629b1710a232148e01
diff --git a/include/lib/psci/psci_lib.h b/include/lib/psci/psci_lib.h
index 1ac45ad..43e2f96 100644
--- a/include/lib/psci/psci_lib.h
+++ b/include/lib/psci/psci_lib.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2022, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -91,6 +91,8 @@
 			  entry_point_info_t *next_image_info);
 int psci_stop_other_cores(unsigned int wait_ms,
 			  void (*stop_func)(u_register_t mpidr));
+bool psci_is_last_on_cpu_safe(void);
+
 #endif /* __ASSEMBLER__ */
 
 #endif /* PSCI_LIB_H */
diff --git a/lib/psci/psci_common.c b/lib/psci/psci_common.c
index 170777f..22d66bf 100644
--- a/lib/psci/psci_common.c
+++ b/lib/psci/psci_common.c
@@ -1022,3 +1022,56 @@
 
 	return PSCI_E_SUCCESS;
 }
+
+/*******************************************************************************
+ * This function verifies that all the other cores in the system have been
+ * turned OFF and the current CPU is the last running CPU in the system.
+ * Returns true if the current CPU is the last ON CPU or false otherwise.
+ *
+ * This API has following differences with psci_is_last_on_cpu
+ *  1. PSCI states are locked
+ *  2. It caters for "forest" topology instead of just "tree"
+ *  TODO : Revisit both API's and unify them
+ ******************************************************************************/
+bool psci_is_last_on_cpu_safe(void)
+{
+	unsigned int this_core = plat_my_core_pos();
+	unsigned int parent_nodes[PLAT_MAX_PWR_LVL] = {0};
+	unsigned int i = 0;
+
+	/*
+	 * Traverse the forest of PSCI nodes, nodes with no parents
+	 * (invalid-nodes) are the root nodes.
+	 */
+	while ((psci_non_cpu_pd_nodes[i].parent_node ==
+	       PSCI_PARENT_NODE_INVALID) &&
+	       (i < PSCI_NUM_NON_CPU_PWR_DOMAINS)) {
+		psci_get_parent_pwr_domain_nodes(
+				psci_non_cpu_pd_nodes[i].cpu_start_idx,
+				PLAT_MAX_PWR_LVL, parent_nodes);
+
+		psci_acquire_pwr_domain_locks(PLAT_MAX_PWR_LVL, parent_nodes);
+
+		for (unsigned int core = 0U;
+		     core < psci_non_cpu_pd_nodes[i].ncpus; core++) {
+			if (core == this_core) {
+				continue;
+			}
+
+			if (psci_get_aff_info_state_by_idx(core) !=
+			    AFF_STATE_OFF) {
+				psci_release_pwr_domain_locks(PLAT_MAX_PWR_LVL,
+							      parent_nodes);
+				VERBOSE("core=%u other than boot core=%u %s\n",
+				       core, this_core, "running in the system");
+
+				return false;
+			}
+		}
+
+		psci_release_pwr_domain_locks(PLAT_MAX_PWR_LVL, parent_nodes);
+		i++;
+	}
+
+	return true;
+}
diff --git a/lib/psci/psci_private.h b/lib/psci/psci_private.h
index 72bd6bd..61bd966 100644
--- a/lib/psci/psci_private.h
+++ b/lib/psci/psci_private.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2022, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -47,6 +47,9 @@
  */
 #define PSCI_MAX_CPUS_INDEX	0xFFFFU
 
+/* Invalid parent */
+#define PSCI_PARENT_NODE_INVALID	0xFFFFFFFFU
+
 /*
  * Helper functions to get/set the fields of PSCI per-cpu data.
  */