3 files changed, 82 insertions, 15 deletions

diff --git a/src/tests/test_assets.cc b/src/tests/test_assets.cc
index e8b6318..b7ee8be 100644
--- a/src/tests/test_assets.cc
+++ b/src/tests/test_assets.cc
@@ -16,21 +16,36 @@ int main() {
   printf("Running AssetManager test...\n");
 
   size_t size = 0;
-  const uint8_t* data = GetAsset(AssetId::ASSET_TEST_ASSET, &size);
+  const uint8_t* data1 = GetAsset(AssetId::ASSET_TEST_ASSET, &size);
 
-  assert(data != nullptr);
+  assert(data1 != nullptr);
   assert(size > 0);
 
   const char* expected_prefix = "This is a test asset file.";
-  if (strncmp((const char*)data, expected_prefix, strlen(expected_prefix)) ==
+  if (strncmp((const char*)data1, expected_prefix, strlen(expected_prefix)) ==
       0) {
     printf("Asset content verification: SUCCESS\n");
   } else {
     printf("Asset content verification: FAILED\n");
-    printf("Got: %.*s\n", (int)size, (const char*)data);
+    printf("Got: %.*s\n", (int)size, (const char*)data1);
     return 1;
   }
 
+  // Test caching: request the same asset again and verify pointer is identical
+  size_t size2 = 0;
+  const uint8_t* data2 = GetAsset(AssetId::ASSET_TEST_ASSET, &size2);
+  assert(data2 != nullptr);
+  assert(size2 == size);
+  assert(data1 == data2); // Pointers should be the same for cached static asset
+  printf("Asset caching test: SUCCESS\n");
+
+  // Test ASSET_LAST_ID - should not return a valid asset
+  size_t last_id_size = 0;
+  const uint8_t* last_id_data = GetAsset(AssetId::ASSET_LAST_ID, &last_id_size);
+  assert(last_id_data == nullptr);
+  assert(last_id_size == 0);
+  printf("ASSET_LAST_ID test: SUCCESS\n");
+
   printf("Asset size: %zu bytes\n", size);
   printf("AssetManager test PASSED\n");
 
diff --git a/src/util/asset_manager.cc b/src/util/asset_manager.cc
index 30295b9..3874535 100644
--- a/src/util/asset_manager.cc
+++ b/src/util/asset_manager.cc
@@ -1,30 +1,82 @@
 // This file is part of the 64k demo project.
-// It implements the generic asset retrieval logic.
-// Uses an array lookup for O(1) access to embedded data.
+// It implements the generic asset retrieval logic with runtime caching.
 
 #include "util/asset_manager.h"
+#if defined(USE_TEST_ASSETS)
+#include "generated/test_assets.h"
+#else
 #include "generated/assets.h"
+#endif /* defined(USE_TEST_ASSETS) */
+
+#include <vector> // For potential dynamic allocation for procedural assets
+#include <new>    // For placement new
+#include <cstdlib> // For free
 
 // These are defined in the generated assets_data.cc
+#if defined(USE_TEST_ASSETS)
+extern const AssetRecord g_assets[];
+extern const size_t g_assets_count;
+#else
 extern const AssetRecord g_assets[];
 extern const size_t g_assets_count;
+#endif
+// Array-based cache for assets.
+// Initialized to all zeros (nullptr data, 0 size, false is_procedural)
+// The size is derived from the generated ASSET_LAST_ID enum value.
+static AssetRecord g_asset_cache[(size_t)AssetId::ASSET_LAST_ID] = {};
 
 const uint8_t* GetAsset(AssetId asset_id, size_t* out_size) {
   uint16_t index = (uint16_t)asset_id;
+
+  // Assert that ASSET_LAST_ID is not used for retrieval.
+  // This ensures the size of the cache is correctly used and not accessed out of bounds.
+  // If this assert fails, it means assets.txt has an ID larger than expected
+  // or ASSET_LAST_ID is not correctly generated/updated.
+  // This is a design decision: ASSET_LAST_ID is purely for sizing and range checking,
+  // not for a valid asset retrieval itself.
+  if (index >= (uint16_t)AssetId::ASSET_LAST_ID) {
+    if (out_size)
+      *out_size = 0;
+    return nullptr; // Invalid asset_id
+  }
+  
+  // Check cache first
+  if (g_asset_cache[index].data != nullptr) {
+    if (out_size)
+      *out_size = g_asset_cache[index].size;
+    return g_asset_cache[index].data;
+  }
+
+  // Not in cache, retrieve from static data (packed in binary)
   if (index >= g_assets_count) {
     if (out_size)
       *out_size = 0;
-    return nullptr;
+    return nullptr; // This asset is not in the static packed data either.
   }
 
-  const AssetRecord& record = g_assets[index];
+  // Store static record in cache for future use
+  g_asset_cache[index] = g_assets[index];
+  g_asset_cache[index].is_procedural = false;
+
   if (out_size)
-    *out_size = record.size;
-  return record.data;
+    *out_size = g_asset_cache[index].size;
+  return g_asset_cache[index].data;
 }
 
 void DropAsset(AssetId asset_id, const uint8_t* asset) {
-  (void)asset_id;
-  (void)asset;
-  // Implementation for lazy decompression will go here
+  uint16_t index = (uint16_t)asset_id;
+  if (index >= (uint16_t)AssetId::ASSET_LAST_ID) {
+    return; // Invalid asset_id
+  }
+
+  // Only free memory for procedural assets.
+  if (g_asset_cache[index].is_procedural && g_asset_cache[index].data == asset) {
+    // In a more complex scenario, we might track ref counts.
+    // For this demo, we assume a single owner for dynamically allocated assets.
+    delete[] g_asset_cache[index].data; // Assuming `new uint8_t[]` was used for procedural
+    g_asset_cache[index].data = nullptr;
+    g_asset_cache[index].size = 0;
+    g_asset_cache[index].is_procedural = false;
+  }
+  // For static assets, no dynamic memory to free.
 }
diff --git a/src/util/asset_manager.h b/src/util/asset_manager.h
index 54d0144..6b09430 100644
--- a/src/util/asset_manager.h
+++ b/src/util/asset_manager.h
@@ -6,12 +6,12 @@
 #include <cstddef>
 #include <cstdint>
 
-// Forward declaration of the generated enum
-enum class AssetId : uint16_t;
+enum class AssetId : uint16_t; // Forward declaration
 
 struct AssetRecord {
   const uint8_t* data;
   size_t size;
+  bool is_procedural; // Flag to indicate if memory was allocated dynamically
 };
 
 // Generic interface