decompiler: reduce memory usage during texture pack installation by lazily loading texture replacements and serializing level building (#4278)

Three changes:
1. Add lazy merging/replacing for texture `rgba_bytes`.
2. Use a single thread when decompiling with texture replacements.
3. Drop textures after serializing and before compression.

These changes should enable users to install **massive** texture packs
without issue.

decompiler/data/TextureDB.cpp

@@ -100,57 +100,110 @@ void TextureDB::add_index_texture(u32 tpage,
 }
 
 void TextureDB::merge_textures(const fs::path& base_path) {
-  for (auto& tex : textures) {
-    fs::path full_path = base_path / tpage_names.at(tex.second.page) / (tex.second.name + ".png");
-    if (fs::exists(full_path)) {
-      lg::info("Merging {}", full_path.string().c_str());
-      int w, h;
-      auto merge_data = stbi_load(full_path.string().c_str(), &w, &h, 0, 4);  // rgba channels
-      if (!merge_data) {
-        lg::warn("failed to load PNG file: {}", full_path.string().c_str());
-        continue;
-      } else if (w != tex.second.w || h != tex.second.h) {
-        lg::warn("merge texture does not match the same dimensions: {}, {} != {} || {} != {}",
-                 full_path.string().c_str(), w, tex.second.w, h, tex.second.h);
-        stbi_image_free(merge_data);
-        continue;
-      }
-      // Merge any non-transparent pixels into the existing texture
-      for (int i = 0; i < w * h * 4; i += 4) {
-        const auto merge_pixel_a = merge_data[i + 3];
-        if (merge_pixel_a != 0) {
-          u32 merge_pixel;
-          memcpy(&merge_pixel, &merge_data[i], sizeof(u32));
-          tex.second.rgba_bytes.at(i / 4) = merge_pixel;
-        }
-      }
-      stbi_image_free(merge_data);
-    }
-  }
+  merge_texture_dir = base_path;
 }
 
 void TextureDB::replace_textures(const fs::path& path) {
-  fs::path base_path(path);
-  for (auto& tex : textures) {
-    fs::path full_path = base_path / tpage_names.at(tex.second.page) / (tex.second.name + ".png");
-    if (!fs::exists(full_path)) {
-      full_path = base_path / "_all" / (tex.second.name + ".png");
-      if (!fs::exists(full_path))
-        continue;
-    }
-    lg::info("Replacing {}", tpage_names.at(tex.second.page) + "/" + (tex.second.name));
-    int w, h;
-    auto data = stbi_load(full_path.string().c_str(), &w, &h, 0, 4);  // rgba channels
-    if (!data) {
-      lg::warn("failed to load PNG file: {}", full_path.string().c_str());
-      continue;
-    }
-    tex.second.rgba_bytes.resize(w * h);
-    memcpy(tex.second.rgba_bytes.data(), data, w * h * 4);
-    tex.second.w = w;
-    tex.second.h = h;
-    stbi_image_free(data);
+  replace_texture_dir = path;
+}
+
+void TextureDB::merge_texture(u32 id, std::vector<u32>& rgba) const {
+  if (!merge_texture_dir) {
+    return;
   }
+
+  const auto& tex = textures.at(id);
+  fs::path full_path = *merge_texture_dir / tpage_names.at(tex.page) / (tex.name + ".png");
+
+  if (!fs::exists(full_path)) {
+    return;
+  }
+
+  lg::info("Merging {}", full_path.string().c_str());
+
+  int w, h;
+  auto merge_data = stbi_load(full_path.string().c_str(), &w, &h, 0, 4);
+
+  if (!merge_data) {
+    lg::warn("failed to load PNG file: {}", full_path.string().c_str());
+    return;
+  }
+
+  if (w != tex.w || h != tex.h) {
+    lg::warn("merge texture does not match the same dimensions: {}, {} != {} || {} != {}",
+             full_path.string().c_str(), w, tex.w, h, tex.h);
+    stbi_image_free(merge_data);
+    return;
+  }
+  // Merge any non-transparent pixels into the existing texture
+  for (int i = 0; i < w * h * 4; i += 4) {
+    const auto merge_pixel_a = merge_data[i + 3];
+    if (merge_pixel_a != 0) {
+      u32 merge_pixel;
+      memcpy(&merge_pixel, &merge_data[i], sizeof(u32));
+      rgba.at(i / 4) = merge_pixel;
+    }
+  }
+
+  stbi_image_free(merge_data);
+}
+
+std::optional<ResolvedTextureData> TextureDB::replace_texture(u32 id) const {
+  if (!replace_texture_dir) {
+    return std::nullopt;
+  }
+
+  const auto& tex = textures.at(id);
+  const auto& tpage_name = tpage_names.at(tex.page);
+
+  fs::path full_path = *replace_texture_dir / tpage_name / (tex.name + ".png");
+
+  if (!fs::exists(full_path)) {
+    full_path = *replace_texture_dir / "_all" / (tex.name + ".png");
+
+    if (!fs::exists(full_path)) {
+      return std::nullopt;
+    }
+  }
+
+  lg::info("Replacing {}", tpage_name + "/" + tex.name);
+
+  int w, h;
+  auto data = stbi_load(full_path.string().c_str(), &w, &h, 0, 4);
+
+  if (!data) {
+    lg::warn("failed to load PNG file: {}", full_path.string().c_str());
+    return std::nullopt;
+  }
+
+  ResolvedTextureData result;
+  result.w = static_cast<u16>(w);
+  result.h = static_cast<u16>(h);
+  result.rgba.resize(w * h);
+
+  memcpy(result.rgba.data(), data, w * h * 4);
+
+  stbi_image_free(data);
+
+  return result;
+}
+
+ResolvedTextureData TextureDB::resolve_texture(u32 id) const {
+  const auto& tex = textures.at(id);
+
+  ResolvedTextureData result{
+      .w = tex.w,
+      .h = tex.h,
+      .rgba = tex.rgba_bytes,
+  };
+
+  merge_texture(id, result.rgba);
+
+  if (auto replacement = replace_texture(id)) {
+    result = std::move(*replacement);
+  }
+
+  return result;
 }
 
 /*!

decompiler/data/TextureDB.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <map>
+#include <optional>
 #include <set>
 #include <string>
 #include <unordered_map>
@@ -11,6 +12,12 @@
 #include "common/util/FileUtil.h"
 
 namespace decompiler {
+struct ResolvedTextureData {
+  u16 w;
+  u16 h;
+  std::vector<u32> rgba;
+};
+
 struct TextureDB {
   TextureDB();
   struct TextureData {
@@ -25,12 +32,16 @@ struct TextureDB {
   std::map<u32, TextureData> textures;
   std::unordered_map<u32, std::string> tpage_names;
   std::unordered_map<std::string, std::set<u32>> texture_ids_per_level;
+  std::optional<fs::path> merge_texture_dir;
+  std::optional<fs::path> replace_texture_dir;
 
   // special textures for animation.
   std::map<u32, tfrag3::IndexTexture> index_textures_by_combo_id;
 
   std::unordered_map<std::string, u32> animated_tex_output_to_anim_slot;
 
+  ResolvedTextureData resolve_texture(u32 id) const;
+
   static constexpr int kPlaceholderWhiteTexturePage = INT16_MAX;
   static constexpr int kPlaceholderWhiteTextureId = 0;
 
@@ -57,6 +68,8 @@ struct TextureDB {
 
   void merge_textures(const fs::path& base_path);
   void replace_textures(const fs::path& path);
+  void merge_texture(u32 id, std::vector<u32>& rgba) const;
+  std::optional<ResolvedTextureData> replace_texture(u32 id) const;
 
   std::string generate_texture_dest_adjustment_table() const;
 };

decompiler/level_extractor/extract_level.cpp

@@ -61,18 +61,17 @@ bool is_valid_bsp(const decompiler::LinkedObjectFile& file) {
   return true;
 }
 
-tfrag3::Texture make_texture(u32 id,
-                             const TextureDB::TextureData& tex,
-                             const std::string& tpage_name,
-                             bool pool_load) {
+tfrag3::Texture make_texture(u32 id, const TextureDB& tex_db, bool pool_load) {
+  const auto& tex = tex_db.textures.at(id);
+  auto resolved = tex_db.resolve_texture(id);
+
   tfrag3::Texture new_tex;
   new_tex.combo_id = id;
-  new_tex.w = tex.w;
-  new_tex.h = tex.h;
-  new_tex.debug_tpage_name = tpage_name;
+  new_tex.w = resolved.w;
+  new_tex.h = resolved.h;
+  new_tex.debug_tpage_name = tex_db.tpage_names.at(tex.page);
   new_tex.debug_name = tex.name;
-  new_tex.data = tex.rgba_bytes;
-  new_tex.combo_id = id;
+  new_tex.data = std::move(resolved.rgba);
   new_tex.load_to_pool = pool_load;
   return new_tex;
 }
@@ -80,12 +79,13 @@ tfrag3::Texture make_texture(u32 id,
 void add_all_textures_from_level(tfrag3::Level& lev,
                                  const std::string& level_name,
                                  const TextureDB& tex_db) {
-  const auto& level_it = tex_db.texture_ids_per_level.find(level_name);
-  if (level_it != tex_db.texture_ids_per_level.end()) {
-    for (auto id : level_it->second) {
-      const auto& tex = tex_db.textures.at(id);
-      lev.textures.push_back(make_texture(id, tex, tex_db.tpage_names.at(tex.page), true));
-    }
+  auto level_it = tex_db.texture_ids_per_level.find(level_name);
+  if (level_it == tex_db.texture_ids_per_level.end()) {
+    return;
+  }
+
+  for (auto id : level_it->second) {
+    lev.textures.push_back(make_texture(id, tex_db, true));
   }
 }
 
@@ -313,8 +313,7 @@ void extract_common(const ObjectFileDB& db,
     if (config.common_tpages.count(normal_texture.page) && !textures_we_have_id.count(id)) {
       textures_we_have.insert(normal_texture.name);
       textures_we_have_id.insert(id);
-      tfrag_level.textures.push_back(
-          make_texture(id, normal_texture, tex_db.tpage_names.at(normal_texture.page), true));
+      tfrag_level.textures.push_back(make_texture(id, tex_db, true));
     }
   }
 
@@ -323,13 +322,16 @@ void extract_common(const ObjectFileDB& db,
     if (config.animated_textures.count(normal_texture.name) &&
         !textures_we_have.count(normal_texture.name)) {
       textures_we_have.insert(normal_texture.name);
-      tfrag_level.textures.push_back(
-          make_texture(id, normal_texture, tex_db.tpage_names.at(normal_texture.page), false));
+      tfrag_level.textures.push_back(make_texture(id, tex_db, false));
     }
   }
 
   Serializer ser;
   tfrag_level.serialize(ser);
+  if (!config.rip_levels) {
+    tfrag_level.textures.clear();
+    tfrag_level.textures.shrink_to_fit();
+  }
   auto compressed =
       compression::compress_zstd(ser.get_save_result().first, ser.get_save_result().second);
 
@@ -369,6 +371,10 @@ void extract_from_level(const ObjectFileDB& db,
 
   Serializer ser;
   level_data.serialize(ser);
+  if (!config.rip_levels) {
+    level_data.textures.clear();
+    level_data.textures.shrink_to_fit();
+  }
   auto compressed =
       compression::compress_zstd(ser.get_save_result().first, ser.get_save_result().second);
   lg::info("stats for {}", level_data.level_name);
@@ -408,12 +414,18 @@ void extract_all_levels(const ObjectFileDB& db,
   auto entities_dir = file_util::get_jak_project_dir() / "decompiler_out" /
                       game_version_names[config.game_version] / "entities";
   file_util::create_dir_if_needed(entities_dir);
+
+  int num_workers = dgo_names.size();
+  if (tex_db.replace_texture_dir) {
+    num_workers = 1;
+  }
+
   SimpleThreadGroup threads;
   threads.run(
       [&](int idx) {
         extract_from_level(db, tex_db, dgo_names[idx], config, output_path, entities_dir);
       },
-      dgo_names.size());
+      dgo_names.size(), num_workers);
   threads.join();
 }
 

decompiler/level_extractor/extract_level.h

@@ -61,10 +61,7 @@ void extract_all_levels(const ObjectFileDB& db,
 void add_all_textures_from_level(tfrag3::Level& lev,
                                  const std::string& level_name,
                                  const TextureDB& tex_db);
-tfrag3::Texture make_texture(u32 id,
-                             const TextureDB::TextureData& tex,
-                             const std::string& tpage_name,
-                             bool pool_load);
+tfrag3::Texture make_texture(u32 id, const TextureDB& tex_db, bool pool_load);
 std::vector<level_tools::TextureRemap> extract_tex_remap(const ObjectFileDB& db,
                                                          const std::string& dgo_name);
 std::optional<ObjectFileRecord> get_bsp_file(const std::vector<ObjectFileRecord>& records,

goalc/build_level/jak1/build_level.cpp

@@ -287,7 +287,7 @@ bool run_build_level(const std::string& input_file,
               if (tex_db.tpage_names.at(tex.page) == tpage_name) {
                 lg::info("custom level: adding texture {} (tpage {})", tex.name, tex.page);
                 tex_names.push_back(tex.name);
-                pc_level.textures.push_back(make_texture(id, tex, tpage_name, true));
+                pc_level.textures.push_back(make_texture(id, tex_db, true));
                 processed_textures[tpage_name].push_back(tex.name);
               }
             }
@@ -324,7 +324,7 @@ bool run_build_level(const std::string& input_file,
             if (db_tpage_name == wanted_tpage_name && tex.name == wanted_tex) {
               lg::info("custom level: adding texture {} from tpage {} ({})", tex.name, tex.page,
                        db_tpage_name);
-              pc_level.textures.push_back(make_texture(id, tex, db_tpage_name, true));
+              pc_level.textures.push_back(make_texture(id, tex_db, true));
               processed_textures[db_tpage_name].push_back(tex.name);
             }
         }

goalc/build_level/jak2/build_level.cpp

@@ -201,8 +201,7 @@ bool run_build_level(const std::string& input_file,
           if (tex.name == tex0) {
             lg::info("custom level: adding texture {} from tpage {} ({})", tex.name, tex.page,
                      tex_db.tpage_names.at(tex.page));
-            pc_level.textures.push_back(
-                make_texture(id, tex, tex_db.tpage_names.at(tex.page), true));
+            pc_level.textures.push_back(make_texture(id, tex_db, true));
             processed_textures.push_back(tex.name);
           }
         }

goalc/build_level/jak3/build_level.cpp

@@ -199,8 +199,7 @@ bool run_build_level(const std::string& input_file,
           if (tex.name == tex0) {
             lg::info("custom level: adding texture {} from tpage {} ({})", tex.name, tex.page,
                      tex_db.tpage_names.at(tex.page));
-            pc_level.textures.push_back(
-                make_texture(id, tex, tex_db.tpage_names.at(tex.page), true));
+            pc_level.textures.push_back(make_texture(id, tex_db, true));
             processed_textures.push_back(tex.name);
           }
         }