Dynamically combine dynamic partitions and groups

Generate the dynamic partitions list and update groups from a
combination of both "dynamic_partitions_info.txt" files.

This is based on the `merge_target_files.py` logic, and generates a
close enough output for our purposes.

Bug: 322887496
Test: see the bug comment 37 for steps to test this stack of changes

Change-Id: Idc0148ba730751638d5d8542569270b78cc542b6

3 files changed, 115 insertions, 70 deletions

diff --git a/host/commands/assemble_cvd/misc_info.cc b/host/commands/assemble_cvd/misc_info.cc
index 292258b96..1c7f1144c 100644
--- a/host/commands/assemble_cvd/misc_info.cc
+++ b/host/commands/assemble_cvd/misc_info.cc
@@ -15,13 +15,12 @@
 
 #include "misc_info.h"
 
-#include <algorithm>
+#include <set>
 #include <string>
-#include <vector>
 
 #include <android-base/logging.h>
-#include <android-base/stringprintf.h>
 #include <android-base/strings.h>
+#include <fmt/format.h>
 
 #include "common/libs/utils/contains.h"
 #include "common/libs/utils/result.h"
@@ -31,7 +30,40 @@ namespace {
 
 constexpr char kDynamicPartitions[] = "dynamic_partition_list";
 constexpr char kGoogleDynamicPartitions[] = "google_dynamic_partitions";
+constexpr char kSuperBlockDevices[] = "super_block_devices";
 constexpr char kSuperPartitionGroups[] = "super_partition_groups";
+constexpr char kUseDynamicPartitions[] = "use_dynamic_partitions";
+
+Result<std::string> GetExpected(const MiscInfo& misc_info,
+                                const std::string& key) {
+  auto lookup = misc_info.find(key);
+  CF_EXPECTF(lookup != misc_info.end(),
+             "Unable to retrieve expected value from key: {}", key);
+  return lookup->second;
+}
+
+std::string GetOrDefault(const MiscInfo& misc_info, const std::string& key,
+                         const std::string& default_value) {
+  auto result = GetExpected(misc_info, key);
+  return result.ok() ? *result : default_value;
+}
+
+std::string MergePartitionLists(const std::string& vendor,
+                                const std::string& system) {
+  const std::string full_string = fmt::format("{} {}", vendor, system);
+  auto full_list = android::base::Tokenize(full_string, " ");
+  // std::set removes duplicates and orders the elements, which we want
+  std::set<std::string> full_set(full_list.begin(), full_list.end());
+  return android::base::Join(full_set, " ");
+}
+
+std::string GetPartitionList(const MiscInfo& vendor_info,
+                             const MiscInfo& system_info,
+                             const std::string& key) {
+  std::string vendor_list = GetOrDefault(vendor_info, key, "");
+  std::string system_list = GetOrDefault(system_info, key, "");
+  return MergePartitionLists(vendor_list, system_list);
+}
 
 }  // namespace
 
@@ -69,65 +101,66 @@ std::string WriteMiscInfo(const MiscInfo& misc_info) {
   return out.str();
 }
 
-std::vector<std::string> SuperPartitionComponents(const MiscInfo& info) {
-  auto value_it = info.find(kDynamicPartitions);
-  if (value_it == info.end()) {
-    return {};
+// based on build/make/tools/releasetools/merge/merge_target_files.py
+Result<MiscInfo> GetCombinedDynamicPartitions(const MiscInfo& vendor_info,
+                                              const MiscInfo& system_info) {
+  CF_EXPECTF(CF_EXPECT(GetExpected(vendor_info, kDynamicPartitions)) == "true",
+             "Vendor build must have {}=true", kUseDynamicPartitions);
+  CF_EXPECTF(CF_EXPECT(GetExpected(system_info, kDynamicPartitions)) == "true",
+             "System build must have {}=true", kUseDynamicPartitions);
+  MiscInfo result;
+  // copy where both files are equal
+  for (const auto& key_val : vendor_info) {
+    const auto value_result = GetExpected(system_info, key_val.first);
+    if (value_result.ok() && *value_result == key_val.second) {
+      result[key_val.first] = key_val.second;
+    }
   }
-  auto components = android::base::Split(value_it->second, " ");
-  for (auto& component : components) {
-    component = android::base::Trim(component);
+
+  result[kDynamicPartitions] =
+      GetPartitionList(vendor_info, system_info, kDynamicPartitions);
+
+  const auto block_devices_result =
+      GetExpected(vendor_info, kSuperBlockDevices);
+  if (block_devices_result.ok()) {
+    result[kSuperBlockDevices] = *block_devices_result;
+    for (const auto& block_device :
+         android::base::Tokenize(result[kSuperBlockDevices], " ")) {
+      const auto key = "super_" + block_device + "_device_size";
+      result[key] = CF_EXPECT(GetExpected(vendor_info, key));
+    }
   }
-  components.erase(std::remove(components.begin(), components.end(), ""),
-                   components.end());
-  return components;
-}
 
-bool SetSuperPartitionComponents(const std::vector<std::string>& components,
-                                 MiscInfo* misc_info) {
-  auto super_partition_groups = misc_info->find(kSuperPartitionGroups);
-  if (super_partition_groups == misc_info->end()) {
-    LOG(ERROR) << "Failed to find super partition groups in misc_info";
-    return false;
+  result[kSuperPartitionGroups] =
+      CF_EXPECT(GetExpected(vendor_info, kSuperPartitionGroups));
+  for (const auto& group :
+       android::base::Tokenize(result[kSuperPartitionGroups], " ")) {
+    const auto group_size_key = "super_" + group + "_group_size";
+    result[group_size_key] =
+        CF_EXPECT(GetExpected(vendor_info, group_size_key));
+
+    const auto partition_list_key = "super_" + group + "_partition_list";
+    result[partition_list_key] =
+        GetPartitionList(vendor_info, system_info, partition_list_key);
   }
 
-  // Remove all existing update groups in misc_info
-  auto update_groups =
-      android::base::Split(super_partition_groups->second, " ");
-  for (const auto& group_name : update_groups) {
-    auto partition_list = android::base::StringPrintf("super_%s_partition_list",
-                                                      group_name.c_str());
-    auto partition_size =
-        android::base::StringPrintf("super_%s_group_size", group_name.c_str());
-    for (const auto& key : {partition_list, partition_size}) {
-      auto it = misc_info->find(key);
-      if (it == misc_info->end()) {
-        LOG(ERROR) << "Failed to find " << key << " in misc_info";
-        return false;
-      }
-      misc_info->erase(it);
+  // TODO(chadreynolds): add vabc_cow_version logic if we need to support older
+  // builds
+  for (const auto& key :
+       {"virtual_ab", "virtual_ab_retrofit", "lpmake", "super_metadata_device",
+        "super_partition_error_limit", "super_partition_size"}) {
+    const auto value_result = GetExpected(vendor_info, key);
+    if (value_result.ok()) {
+      result[key] = *value_result;
     }
   }
+  return std::move(result);
+}
 
-  // For merged target-file, put all dynamic partitions under the
-  // google_dynamic_partitions update group.
-  // TODO(xunchang) use different update groups for system and vendor images.
-  (*misc_info)[kDynamicPartitions] = android::base::Join(components, " ");
-  (*misc_info)[kSuperPartitionGroups] = kGoogleDynamicPartitions;
-  std::string partitions_list_key = android::base::StringPrintf(
-      "super_%s_partition_list", kGoogleDynamicPartitions);
-  (*misc_info)[partitions_list_key] = android::base::Join(components, " ");
-
-  // Use the entire super partition as the group size
-  std::string group_size_key = android::base::StringPrintf(
-      "super_%s_group_size", kGoogleDynamicPartitions);
-  auto super_size_it = misc_info->find("super_partition_size");
-  if (super_size_it == misc_info->end()) {
-    LOG(ERROR) << "Failed to find super partition size";
-    return false;
+void MergeInKeys(const MiscInfo& source, MiscInfo& target) {
+  for (const auto& key_val : source) {
+    target[key_val.first] = key_val.second;
   }
-  (*misc_info)[group_size_key] = super_size_it->second;
-  return true;
 }
 
 } // namespace cuttlefish
diff --git a/host/commands/assemble_cvd/misc_info.h b/host/commands/assemble_cvd/misc_info.h
index f832a4c5f..d0abae896 100644
--- a/host/commands/assemble_cvd/misc_info.h
+++ b/host/commands/assemble_cvd/misc_info.h
@@ -23,13 +23,14 @@
 
 namespace cuttlefish {
 
+// TODO(chadreynolds): rename MiscInfo to more generic KeyValueFile since this
+// logic is processing multiple filetypes now
 using MiscInfo = std::map<std::string, std::string>;
 
 Result<MiscInfo> ParseMiscInfo(const std::string& file_contents);
 std::string WriteMiscInfo(const MiscInfo& info);
-
-std::vector<std::string> SuperPartitionComponents(const MiscInfo&);
-bool SetSuperPartitionComponents(const std::vector<std::string>& components,
-                                 MiscInfo* misc_info);
+Result<MiscInfo> GetCombinedDynamicPartitions(const MiscInfo& vendor_info,
+                                              const MiscInfo& system_info);
+void MergeInKeys(const MiscInfo& source, MiscInfo& target);
 
 } // namespace cuttlefish
diff --git a/host/commands/assemble_cvd/super_image_mixer.cc b/host/commands/assemble_cvd/super_image_mixer.cc
index b59d163fb..a5969080f 100644
--- a/host/commands/assemble_cvd/super_image_mixer.cc
+++ b/host/commands/assemble_cvd/super_image_mixer.cc
@@ -41,6 +41,7 @@ namespace cuttlefish {
 namespace {
 
 constexpr char kMiscInfoPath[] = "META/misc_info.txt";
+constexpr char kDynamicPartitionsPath[] = "META/dynamic_partitions_info.txt";
 constexpr std::array kVendorTargetImages = {
     "IMAGES/boot.img",
     "IMAGES/dtbo.img",
@@ -107,29 +108,39 @@ Result<TargetFiles> GetTargetFiles(const std::string& vendor_zip_path,
   return result;
 }
 
+Result<MiscInfo> CombineDynamicPartitionsInfo(TargetFiles& target_files) {
+  CF_EXPECTF(Contains(target_files.vendor_contents, kDynamicPartitionsPath),
+             "Vendor target files zip does not contain {}",
+             kDynamicPartitionsPath);
+  CF_EXPECTF(Contains(target_files.system_contents, kDynamicPartitionsPath),
+             "System target files zip does not contain {}",
+             kDynamicPartitionsPath);
+
+  const MiscInfo vendor_dp_info = CF_EXPECT(ParseMiscInfo(
+      target_files.vendor_zip.ExtractToMemory(kDynamicPartitionsPath)));
+  const MiscInfo system_dp_info = CF_EXPECT(ParseMiscInfo(
+      target_files.system_zip.ExtractToMemory(kDynamicPartitionsPath)));
+
+  return CF_EXPECT(
+      GetCombinedDynamicPartitions(vendor_dp_info, system_dp_info));
+}
+
 Result<void> CombineMiscInfo(TargetFiles& target_files,
                              const std::string& misc_output_path) {
   CF_EXPECTF(Contains(target_files.vendor_contents, kMiscInfoPath),
-             "Default target files zip does not contain {}", kMiscInfoPath);
+             "Vendor target files zip does not contain {}", kMiscInfoPath);
   CF_EXPECTF(Contains(target_files.system_contents, kMiscInfoPath),
              "System target files zip does not contain {}", kMiscInfoPath);
 
-  const MiscInfo default_misc = CF_EXPECT(
+  // the combined misc info uses the vendor values as defaults
+  MiscInfo output_misc = CF_EXPECT(
       ParseMiscInfo(target_files.vendor_zip.ExtractToMemory(kMiscInfoPath)));
   const MiscInfo system_misc = CF_EXPECT(
       ParseMiscInfo(target_files.system_zip.ExtractToMemory(kMiscInfoPath)));
 
-  auto output_misc = default_misc;
-  auto system_super_partitions = SuperPartitionComponents(system_misc);
-  // Ensure specific skipped partitions end up in the misc_info.txt
-  for (auto partition :
-       {"odm", "odm_dlkm", "vendor", "vendor_dlkm", "system_dlkm"}) {
-    if (!Contains(system_super_partitions, partition)) {
-      system_super_partitions.push_back(partition);
-    }
-  }
-  CF_EXPECT(SetSuperPartitionComponents(system_super_partitions, &output_misc),
-            "Failed to update super partitions components for misc_info");
+  const auto combined_dp_info =
+      CF_EXPECT(CombineDynamicPartitionsInfo(target_files));
+  MergeInKeys(combined_dp_info, output_misc);
 
   SharedFD misc_output_file = SharedFD::Creat(misc_output_path.c_str(), 0644);
   CF_EXPECT(misc_output_file->IsOpen(), "Failed to open output misc file: "