From 41f3bc23cbf89f543efd98e8695c74c8bbf019d1 Mon Sep 17 00:00:00 2001
From: Kasper Lund <kasper@toit.io>
Date: Mon, 3 Jun 2024 11:52:52 +0200
Subject: [PATCH] Take program assets into account when computing flash
 allocation sizes. (#2379)

---
 src/flash_allocation.cc         | 8 +++++++-
 src/flash_allocation.h          | 5 ++++-
 src/primitive_flash_registry.cc | 4 +---
 src/process.cc                  | 2 +-
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/flash_allocation.cc b/src/flash_allocation.cc
index 74507edc8..c9a17d2f8 100644
--- a/src/flash_allocation.cc
+++ b/src/flash_allocation.cc
@@ -73,6 +73,12 @@ FlashAllocation::Header::Header(const void* memory,
   checksum_ = compute_checksum(memory);
 }
 
+word FlashAllocation::size() const {
+  word size = size_no_assets();
+  if (is_program()) size += program_assets_size(null, null);
+  return size;
+}
+
 bool FlashAllocation::Header::is_valid(bool embedded) const {
   if (marker_ != FORMAT_MARKER || size_in_pages_ == 0) return false;
   if (embedded) {
@@ -100,7 +106,7 @@ bool FlashAllocation::commit(const void* memory, word size, const Header* header
 int FlashAllocation::program_assets_size(uint8** bytes, word* length) const {
   if (!program_has_assets()) return 0;
   uword allocation_address = reinterpret_cast<uword>(this);
-  uword assets_address = allocation_address + size();
+  uword assets_address = allocation_address + size_no_assets();
   int assets_length = *reinterpret_cast<uint32*>(assets_address);
   if (bytes) *bytes = reinterpret_cast<uint8*>(assets_address + sizeof(uint32));
   if (length) *length = assets_length;
diff --git a/src/flash_allocation.h b/src/flash_allocation.h
index 596197235..08743f8dc 100644
--- a/src/flash_allocation.h
+++ b/src/flash_allocation.h
@@ -68,11 +68,14 @@ class FlashAllocation {
   bool is_region() const { return type() == FLASH_ALLOCATION_TYPE_REGION; }
 
   // Simple accessors.
-  word size() const { return header_.size(); }
+  word size_no_assets() const { return header_.size(); }
   uint8 type() const { return header_.type(); }
   const uint8* id() const { return header_.id(); }
   const uint8* metadata() const { return header_.metadata(); }
 
+  // Get the full size of the allocation. For programs, this includes the assets.
+  word size() const;
+
   // Check if the allocation is valid.
   bool is_valid() const { return header_.is_valid(false); }
   bool is_valid_embedded() const { return header_.is_valid(true); }
diff --git a/src/primitive_flash_registry.cc b/src/primitive_flash_registry.cc
index 39b31e714..322aa637d 100644
--- a/src/primitive_flash_registry.cc
+++ b/src/primitive_flash_registry.cc
@@ -105,9 +105,7 @@ PRIMITIVE(get_size) {
   ARGS(word, offset);
   const FlashAllocation* allocation = FlashRegistry::allocation(offset);
   if (allocation == null) FAIL(INVALID_ARGUMENT);
-  word size = allocation->size();
-  if (allocation->is_program()) size += allocation->program_assets_size(null, null);
-  return Smi::from(size);
+  return Smi::from(allocation->size());
 }
 
 PRIMITIVE(get_header_page) {
diff --git a/src/process.cc b/src/process.cc
index 3fa7ff3f1..a7ffb77f1 100644
--- a/src/process.cc
+++ b/src/process.cc
@@ -40,7 +40,7 @@ Process::Process(Program* program, ProcessRunner* runner, ProcessGroup* group, S
     , runner_(runner)
     , group_(group)
     , program_heap_address_(reinterpret_cast<uword>(program))
-    , program_heap_size_(program ? program->size() : 0)
+    , program_heap_size_(program ? program->size_no_assets() : 0)
     , entry_(Method::invalid())
     , spawn_method_(Method::invalid())
     , object_heap_(