1 files changed, 13 insertions, 10 deletions

diff --git a/src/lib/loaders/load_and_run_ramstage.c b/src/lib/loaders/load_and_run_ramstage.c
index b3728a17cf..153e38e11e 100644
--- a/src/lib/loaders/load_and_run_ramstage.c
+++ b/src/lib/loaders/load_and_run_ramstage.c
@@ -23,6 +23,7 @@
 #include <cbfs.h>
 #include <program_loading.h>
 #include <romstage_handoff.h>
+#include <stage_cache.h>
 #include <timestamp.h>
 
 extern const struct prog_loader_ops cbfs_ramstage_loader;
@@ -35,16 +36,20 @@ static const struct prog_loader_ops *loaders[] = {
 	&cbfs_ramstage_loader,
 };
 
-static void
-load_ramstage(const struct prog_loader_ops *ops,
-		struct romstage_handoff *handoff, struct prog *ramstage)
+void __attribute__((weak)) stage_cache_add(int stage_id, struct prog *stage) {}
+void __attribute__((weak)) stage_cache_load_stage(int stage_id,
+							struct prog *stage) {}
+void __attribute__((weak)) ramstage_cache_invalid(void) {}
+
+static void load_ramstage(const struct prog_loader_ops *ops,
+				struct prog *ramstage)
 {
 	timestamp_add_now(TS_START_COPYRAM);
 
 	if (ops->prepare(ramstage))
 		return;
 
-	cache_loaded_ramstage(handoff, ramstage);
+	stage_cache_add(STAGE_RAMSTAGE, ramstage);
 
 	timestamp_add_now(TS_END_COPYRAM);
 
@@ -56,18 +61,18 @@ static void run_ramstage_from_resume(struct romstage_handoff *handoff,
 {
 	if (handoff != NULL && handoff->s3_resume) {
 		/* Load the cached ramstage to runtime location. */
-		load_cached_ramstage(handoff, ramstage);
+		stage_cache_load_stage(STAGE_RAMSTAGE, ramstage);
 
 		if (prog_entry(ramstage) != NULL) {
 			printk(BIOS_DEBUG, "Jumping to image.\n");
 			prog_run(ramstage);
 		}
+		ramstage_cache_invalid();
 	}
 }
 
 void run_ramstage(void)
 {
-	struct romstage_handoff *handoff;
 	const struct prog_loader_ops *ops;
 	int i;
 	struct prog ramstage = {
@@ -75,14 +80,12 @@ void run_ramstage(void)
 		.type = PROG_RAMSTAGE,
 	};
 
-	handoff = romstage_handoff_find_or_add();
-
-	run_ramstage_from_resume(handoff, &ramstage);
+	run_ramstage_from_resume(romstage_handoff_find_or_add(), &ramstage);
 
 	for (i = 0; i < ARRAY_SIZE(loaders); i++) {
 		ops = loaders[i];
 		printk(BIOS_DEBUG, "Trying %s ramstage loader.\n", ops->name);
-		load_ramstage(ops, handoff, &ramstage);
+		load_ramstage(ops, &ramstage);
 	}
 
 	die("Ramstage was not loaded!\n");