From e3ef3413c56427add7a9f123882b06cc3352061b Mon Sep 17 00:00:00 2001
From: Koby Chan <koby@meta.com>
Date: Fri, 17 Jan 2025 13:58:23 -0800
Subject: [PATCH] Enable including all betamap classes/methods

Summary: Adds an option to include all classes/methods from betamap into the baseline profile.

Reviewed By: danjin250

Differential Revision: D68174163

fbshipit-source-id: dcc223ddf408b4a3fb1e935784aad278c530ddcd
---
 libredex/BaselineProfileConfig.h              |  4 ++++
 libredex/ConfigFiles.cpp                      |  3 +++
 .../ArtProfileWriterPass.cpp                  | 23 +++++++++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/libredex/BaselineProfileConfig.h b/libredex/BaselineProfileConfig.h
index b675a44305..b88c347c55 100644
--- a/libredex/BaselineProfileConfig.h
+++ b/libredex/BaselineProfileConfig.h
@@ -43,6 +43,10 @@ struct BaselineProfileOptions {
   // NOTE: This requires that include_betamap_20pct_coldstart be set to have any
   // effect
   bool betamap_include_coldstart_1pct;
+
+  // If this is true, then the ArtProfileWriter will insert all methods/classes
+  // from the betamap into the baseline profile.
+  bool include_all_startup_classes;
 };
 
 struct BaselineProfileConfig {
diff --git a/libredex/ConfigFiles.cpp b/libredex/ConfigFiles.cpp
index 563f52de23..0d8c91df2b 100644
--- a/libredex/ConfigFiles.cpp
+++ b/libredex/ConfigFiles.cpp
@@ -637,6 +637,9 @@ ConfigFiles::get_default_baseline_profile_config() {
         "betamap_include_coldstart_1pct", false,
         current_baseline_profile_config.options.betamap_include_coldstart_1pct);
     Json::Value deepdata_interactions_json;
+    baseline_profile_config_jw.get(
+        "include_all_startup_classes", false,
+        current_baseline_profile_config.options.include_all_startup_classes);
     baseline_profile_config_jw.get("deep_data_interaction_config", {},
                                    deepdata_interactions_json);
     always_assert(!deepdata_interactions_json.empty());
diff --git a/opt/art-profile-writer/ArtProfileWriterPass.cpp b/opt/art-profile-writer/ArtProfileWriterPass.cpp
index 12f7f0abd0..51e2116ad7 100644
--- a/opt/art-profile-writer/ArtProfileWriterPass.cpp
+++ b/opt/art-profile-writer/ArtProfileWriterPass.cpp
@@ -550,6 +550,29 @@ void ArtProfileWriterPass::run_pass(DexStoresVector& stores,
                               : baseline_profiles::get_baseline_profile(
                                     conf.get_default_baseline_profile_config(),
                                     method_profiles, &method_refs_without_def);
+  if (conf.get_default_baseline_profile_config()
+          .options.include_all_startup_classes) {
+    const std::vector<std::string>& interdexorder =
+        conf.get_coldstart_classes();
+    std::vector<DexClass*> coldstart_classes;
+    for (const auto& entry : interdexorder) {
+      DexType* type = DexType::get_type(entry);
+      if (type) {
+        auto coldstart_class = type_class(type);
+        if (coldstart_class) {
+          coldstart_classes.push_back(coldstart_class);
+          baseline_profile.classes.insert(coldstart_class);
+        }
+      }
+    }
+    baseline_profiles::MethodFlags flags;
+    flags.hot = true;
+    flags.startup = true;
+    walk::methods(coldstart_classes,
+                  [&baseline_profile, flags](DexMethod* method) {
+                    baseline_profile.methods.emplace(method, flags);
+                  });
+  }
   auto scope = build_class_scope(stores);
   if (m_never_compile_callcount_threshold > -1 ||
       m_never_compile_perf_threshold > -1) {