[PATCH v2] spl: bootstage: move bootstage_stash before jumping to image

[Chanho Park]

Regarding IH_OS_OPENSBI, IH_OS_LINUX and IH_OS_TEE, there is no chance
to stash bootstage record because they do not return to SPL after
jumping to the image.
Hence, this patch separates the final stage bootstage code into
spl_bootstage_finish and call the function before jumping to the image.

Signed-off-by: Chanho Park <chanho61.park at samsung.com>
---
Changes from v1
- Separate the final stage bootstage code into spl_bootstage_finish.
- As Simon suggests, call the function before jumping to the image.

 common/spl/spl.c | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/common/spl/spl.c b/common/spl/spl.c
index 0062f3f45d9a..eaf2f076ddd1 100644
--- a/common/spl/spl.c
+++ b/common/spl/spl.c
@@ -738,6 +738,19 @@ void board_init_f(ulong dummy)
 }
 #endif
 
+static void spl_bootstage_finish(void)
+{
+	int ret;
+
+	bootstage_mark_name(get_bootstage_id(false), "end phase");
+#ifdef CONFIG_BOOTSTAGE_STASH
+	ret = bootstage_stash((void *)CONFIG_BOOTSTAGE_STASH_ADDR,
+			      CONFIG_BOOTSTAGE_STASH_SIZE);
+	if (ret)
+		debug("Failed to stash bootstage: err=%d\n", ret);
+#endif
+}
+
 void board_init_r(gd_t *dummy1, ulong dummy2)
 {
 	u32 spl_boot_list[] = {
@@ -865,12 +878,14 @@ void board_init_r(gd_t *dummy1, ulong dummy2)
 	case IH_OS_TEE:
 		debug("Jumping to U-Boot via OP-TEE\n");
 		spl_board_prepare_for_optee(spl_image.fdt_addr);
+		spl_bootstage_finish();
 		jump_to_image_optee(&spl_image);
 		break;
 #endif
 #if CONFIG_IS_ENABLED(OPENSBI)
 	case IH_OS_OPENSBI:
 		debug("Jumping to U-Boot via RISC-V OpenSBI\n");
+		spl_bootstage_finish();
 		spl_invoke_opensbi(&spl_image);
 		break;
 #endif
@@ -881,6 +896,7 @@ void board_init_r(gd_t *dummy1, ulong dummy2)
 		spl_fixup_fdt((void *)CONFIG_SYS_SPL_ARGS_ADDR);
 #endif
 		spl_board_prepare_for_linux();
+		spl_bootstage_finish();
 		jump_to_image_linux(&spl_image);
 #endif
 	default:
@@ -890,13 +906,7 @@ void board_init_r(gd_t *dummy1, ulong dummy2)
 	debug("SPL malloc() used 0x%lx bytes (%ld KB)\n", gd->malloc_ptr,
 	      gd->malloc_ptr / 1024);
 #endif
-	bootstage_mark_name(get_bootstage_id(false), "end phase");
-#ifdef CONFIG_BOOTSTAGE_STASH
-	ret = bootstage_stash((void *)CONFIG_BOOTSTAGE_STASH_ADDR,
-			      CONFIG_BOOTSTAGE_STASH_SIZE);
-	if (ret)
-		debug("Failed to stash bootstage: err=%d\n", ret);
-#endif
+	spl_bootstage_finish();
 
 	if (IS_ENABLED(CONFIG_SPL_VIDEO_REMOVE)) {
 		struct udevice *dev;
-- 
2.39.2