src/lib: Add memory/time saving special case for ramstage caching

When caching the ramstage for suspend/resume, we copy the entire image
as it resides in RAM. The last part of that, CONFIG_HEAP_SIZE bytes, is
the heap that will be reinitialized when the ramstage is started again.

As such, copying doesn't make sense and complicates HEAP_SIZE
configuration (because it needs to fit the space-constrained cache
location) and costs time and space. Therefore, skip the heap.

Side notes:
- When building with ASAN, program.ld indicates that it will allocate
  some more space after the heap. This is not a problem, we just copy
  an ASAN-sized copy of the heap.
- Heap use is managed in src/lib/malloc with statically allocated
  variables. Because ramstage is cached before it's executed, these
  values will be reset to their compile-time default values, too.

Change-Id: I6553dc8b758196f2476af2e692c0421d0fa2b98e
Signed-off-by: Patrick Georgi <patrick@coreboot.org>
Reviewed-on: https://review.coreboot.org/c/coreboot/+/79525
Reviewed-by: Martin L Roth <gaumless@gmail.com>
Reviewed-by: Paul Menzel <paulepanter@mailbox.org>
Tested-by: build bot (Jenkins) <no-reply@coreboot.org>
Reviewed-by: Arthur Heymans <arthur@aheymans.xyz>
diff --git a/src/lib/cbmem_stage_cache.c b/src/lib/cbmem_stage_cache.c
index b45c950..5e76fb0 100644
--- a/src/lib/cbmem_stage_cache.c
+++ b/src/lib/cbmem_stage_cache.c
@@ -21,14 +21,21 @@
 	meta->entry_addr = (uintptr_t)prog_entry(stage);
 	meta->arg = (uintptr_t)prog_entry_arg(stage);
 
-	c = cbmem_add(CBMEM_ID_STAGEx_CACHE + stage_id, prog_size(stage));
+	unsigned int p_size = prog_size(stage);
+	if (stage_id == STAGE_RAMSTAGE) {
+		/* heap resides at the end of the image and will be
+		   reinitialized, so it doesn't make sense to copy it around. */
+		p_size -= CONFIG_HEAP_SIZE;
+	}
+
+	c = cbmem_add(CBMEM_ID_STAGEx_CACHE + stage_id, p_size);
 	if (c == NULL) {
 		printk(BIOS_ERR, "Can't add stage_cache %x to cbmem\n",
 				CBMEM_ID_STAGEx_CACHE + stage_id);
 		return;
 	}
 
-	memcpy(c, prog_start(stage), prog_size(stage));
+	memcpy(c, prog_start(stage), p_size);
 }
 
 void stage_cache_add_raw(int stage_id, const void *base, const size_t size)
diff --git a/src/lib/ext_stage_cache.c b/src/lib/ext_stage_cache.c
index a35ce37..462a635 100644
--- a/src/lib/ext_stage_cache.c
+++ b/src/lib/ext_stage_cache.c
@@ -59,8 +59,15 @@
 	meta->entry_addr = (uintptr_t)prog_entry(stage);
 	meta->arg = (uintptr_t)prog_entry_arg(stage);
 
-	e = imd_entry_add(imd, CBMEM_ID_STAGEx_CACHE + stage_id,
-				prog_size(stage));
+	unsigned int p_size = prog_size(stage);
+	if (stage_id == STAGE_RAMSTAGE) {
+		/* heap resides at the end of the image and will be
+		 * reinitialized, so it doesn't make sense to copy it around.
+		 */
+		p_size -= CONFIG_HEAP_SIZE;
+	}
+
+	e = imd_entry_add(imd, CBMEM_ID_STAGEx_CACHE + stage_id, p_size);
 
 	if (e == NULL) {
 		printk(BIOS_DEBUG, "Error: Can't add stage_cache %x to imd\n",
@@ -70,7 +77,7 @@
 
 	c = imd_entry_at(imd, e);
 
-	memcpy(c, prog_start(stage), prog_size(stage));
+	memcpy(c, prog_start(stage), p_size);
 }
 
 void stage_cache_add_raw(int stage_id, const void *base, const size_t size)
diff --git a/src/lib/program.ld b/src/lib/program.ld
index f406f9f..68bcab6 100644
--- a/src/lib/program.ld
+++ b/src/lib/program.ld
@@ -149,6 +149,15 @@
 _eprogram = .;
 RECORD_SIZE(program)
 
+/* The stage cache drops CONFIG_HEAP_SIZE bytes from the end of the in-memory
+   image of the ramstage, so ensure that when moving that many bytes backwards
+   from the program end, we're in the heap (or later), in some region that
+   doesn't contain initialized code or data. */
+#if ENV_RAMSTAGE
+_bogus = ASSERT(_eprogram - CONFIG_HEAP_SIZE >= _heap,
+	"HEAP_SIZE and heap misaligned");
+#endif
+
 /* Discard the sections we don't need/want */
 
 zeroptr = 0;