From 834f4e6c928d5e74202fc4b2e11510bcf4f79d8f Mon Sep 17 00:00:00 2001
From: Yurii Zubrytskyi <zyy@google.com>
Date: Thu, 21 Oct 2021 09:13:18 -0700
Subject: [PATCH] [aapt2] Fix finalized resources removal

The code used to expect all original staged resources to get
parsed before their finalized version. Looks like that's not the
actual case, and parsing code needs to expect those to come
later, and preserve the information that they aren't needed
anymore

Bug: 203779955
Test: manual, aapt2 dump resources <finalized framework SDK>
Change-Id: I6d9efe2eb7b406d05ca49a0a6ed745a1c509e2ef
---
 .../format/binary/BinaryResourceParser.cpp     | 18 +++++++++++++++---
 .../aapt2/format/binary/BinaryResourceParser.h |  4 ++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/tools/aapt2/format/binary/BinaryResourceParser.cpp b/tools/aapt2/format/binary/BinaryResourceParser.cpp
index 72eaa3561a023..a8845ef96bc36 100644
--- a/tools/aapt2/format/binary/BinaryResourceParser.cpp
+++ b/tools/aapt2/format/binary/BinaryResourceParser.cpp
@@ -120,6 +120,13 @@ bool BinaryResourceParser::Parse() {
                                   static_cast<int>(parser.chunk()->type)));
     }
   }
+
+  if (!staged_entries_to_remove_.empty()) {
+    diag_->Error(DiagMessage(source_) << "didn't find " << staged_entries_to_remove_.size()
+                                      << " original staged resources");
+    return false;
+  }
+
   return true;
 }
 
@@ -393,6 +400,12 @@ bool BinaryResourceParser::ParseType(const ResourceTablePackage* package,
       return false;
     }
 
+    if (const auto to_remove_it = staged_entries_to_remove_.find({name, res_id});
+        to_remove_it != staged_entries_to_remove_.end()) {
+      staged_entries_to_remove_.erase(to_remove_it);
+      continue;
+    }
+
     NewResourceBuilder res_builder(name);
     res_builder.SetValue(std::move(resource_value), config)
         .SetId(res_id, OnIdConflict::CREATE_ENTRY)
@@ -533,9 +546,8 @@ bool BinaryResourceParser::ParseStagedAliases(const ResChunk_header* chunk) {
     // Since a the finalized resource entry is cloned and added to the resource table under the
     // staged resource id, remove the cloned resource entry from the table.
     if (!table_->RemoveResource(resource_name, staged_id)) {
-      diag_->Error(DiagMessage(source_) << "failed to find resource entry for staged "
-                                        << " resource ID " << staged_id);
-      return false;
+      // If we haven't seen this resource yet let's add a record to skip it when parsing.
+      staged_entries_to_remove_.insert({resource_name, staged_id});
     }
   }
   return true;
diff --git a/tools/aapt2/format/binary/BinaryResourceParser.h b/tools/aapt2/format/binary/BinaryResourceParser.h
index cd71d160703ad..1c83166c5cce9 100644
--- a/tools/aapt2/format/binary/BinaryResourceParser.h
+++ b/tools/aapt2/format/binary/BinaryResourceParser.h
@@ -119,6 +119,10 @@ class BinaryResourceParser {
 
   // A mapping of resource ID to type spec flags.
   std::unordered_map<ResourceId, uint32_t> entry_type_spec_flags_;
+
+  // A collection of staged resources that got finalized already and we're supposed to prune -
+  // but the original staged resource record hasn't been parsed yet.
+  std::set<std::pair<ResourceName, ResourceId>> staged_entries_to_remove_;
 };
 
 }  // namespace aapt