1 files changed, 150 insertions, 58 deletions

diff --git a/tools/aapt2/compile/IdAssigner.cpp b/tools/aapt2/compile/IdAssigner.cpp
index aa4a5803b8df..4a3f1e10b6db 100644
--- a/tools/aapt2/compile/IdAssigner.cpp
+++ b/tools/aapt2/compile/IdAssigner.cpp
@@ -15,92 +15,184 @@
  */
 
 #include "ResourceTable.h"
-
 #include "compile/IdAssigner.h"
 #include "process/IResourceTableConsumer.h"
 #include "util/Util.h"
 
-#include <bitset>
 #include <cassert>
-#include <set>
+#include <map>
 
 namespace aapt {
 
+/**
+ * Assigns the intended ID to the ResourceTablePackage, ResourceTableType, and ResourceEntry,
+ * as long as there is no existing ID or the ID is the same.
+ */
+static bool assignId(IDiagnostics* diag, const ResourceId& id, const ResourceName& name,
+                     ResourceTablePackage* pkg, ResourceTableType* type, ResourceEntry* entry) {
+    if (pkg->id.value() == id.packageId()) {
+        if (!type->id || type->id.value() == id.typeId()) {
+            type->id = id.typeId();
+
+            if (!entry->id || entry->id.value() == id.entryId()) {
+                entry->id = id.entryId();
+                return true;
+            }
+        }
+    }
+
+    const ResourceId existingId(pkg->id.value(),
+                                type->id ? type->id.value() : 0,
+                                entry->id ? entry->id.value() : 0);
+    diag->error(DiagMessage() << "can't assign ID " << id
+                << " to resource " << name
+                << " with conflicting ID " << existingId);
+    return false;
+}
+
 bool IdAssigner::consume(IAaptContext* context, ResourceTable* table) {
-    std::bitset<256> usedTypeIds;
-    std::set<uint16_t> usedEntryIds;
+    std::map<ResourceId, ResourceName> assignedIds;
 
     for (auto& package : table->packages) {
         assert(package->id && "packages must have manually assigned IDs");
 
-        usedTypeIds.reset();
+        for (auto& type : package->types) {
+            for (auto& entry : type->entries) {
+                const ResourceName name(package->name, type->type, entry->name);
 
-        // Type ID 0 is invalid, reserve it.
-        usedTypeIds.set(0);
+                if (mAssignedIdMap) {
+                    // Assign the pre-assigned stable ID meant for this resource.
+                    const auto iter = mAssignedIdMap->find(name);
+                    if (iter != mAssignedIdMap->end()) {
+                        const ResourceId assignedId = iter->second;
+                        const bool result = assignId(context->getDiagnostics(), assignedId, name,
+                                                     package.get(), type.get(), entry.get());
+                        if (!result) {
+                            return false;
+                        }
+                    }
+                }
 
-        // Collect used type IDs.
-        for (auto& type : package->types) {
-            if (type->id) {
-                usedEntryIds.clear();
-
-                if (usedTypeIds[type->id.value()]) {
-                    // This ID is already taken!
-                    context->getDiagnostics()->error(DiagMessage()
-                                                     << "type '" << type->type << "' in "
-                                                     << "package '" << package->name << "' has "
-                                                     << "duplicate ID "
-                                                     << std::hex << (int) type->id.value()
-                                                     << std::dec);
-                    return false;
+                if (package->id && type->id && entry->id) {
+                    // If the ID is set for this resource, then reserve it.
+                    ResourceId resourceId(package->id.value(), type->id.value(), entry->id.value());
+                    auto result = assignedIds.insert({ resourceId, name });
+                    const ResourceName& existingName = result.first->second;
+                    if (!result.second) {
+                        context->getDiagnostics()->error(DiagMessage() << "resource " << name
+                                                         << " has same ID "
+                                                         << resourceId
+                                                         << " as " << existingName);
+                        return false;
+                    }
                 }
+            }
+        }
+    }
 
-                // Mark the type ID as taken.
-                usedTypeIds.set(type->id.value());
+    if (mAssignedIdMap) {
+        // Reserve all the IDs mentioned in the stable ID map. That way we won't assign
+        // IDs that were listed in the map if they don't exist in the table.
+        for (const auto& stableIdEntry : *mAssignedIdMap) {
+            const ResourceName& preAssignedName = stableIdEntry.first;
+            const ResourceId& preAssignedId = stableIdEntry.second;
+            auto result = assignedIds.insert({ preAssignedId, preAssignedName });
+            const ResourceName& existingName = result.first->second;
+            if (!result.second && existingName != preAssignedName) {
+                context->getDiagnostics()->error(DiagMessage() << "stable ID " << preAssignedId
+                                                 << " for resource " << preAssignedName
+                                                 << " is already taken by resource "
+                                                 << existingName);
+                return false;
             }
+        }
+    }
 
-            // Collect used entry IDs.
-            for (auto& entry : type->entries) {
-                if (entry->id) {
-                    // Mark entry ID as taken.
-                    if (!usedEntryIds.insert(entry->id.value()).second) {
-                        // This ID existed before!
-                        ResourceNameRef nameRef(package->name, type->type, entry->name);
-                        context->getDiagnostics()->error(DiagMessage()
-                                                         << "resource '" << nameRef << "' "
-                                                         << "has duplicate entry ID "
-                                                         << std::hex << (int) entry->id.value()
-                                                         << std::dec);
-                        return false;
+    // Assign any resources without IDs the next available ID. Gaps will be filled if possible,
+    // unless those IDs have been reserved.
+
+    const auto assignedIdsIterEnd = assignedIds.end();
+    for (auto& package : table->packages) {
+        assert(package->id && "packages must have manually assigned IDs");
+
+        // Build a half filled ResourceId object, which will be used to find the closest matching
+        // reserved ID in the assignedId map. From that point the next available type ID can be
+        // found.
+        ResourceId resourceId(package->id.value(), 0, 0);
+        uint8_t nextExpectedTypeId = 1;
+
+        // Find the closest matching ResourceId that is <= the one with only the package set.
+        auto nextTypeIter = assignedIds.lower_bound(resourceId);
+        for (auto& type : package->types) {
+            if (!type->id) {
+                // We need to assign a type ID. Iterate over the reserved IDs until we find
+                // some type ID that is a distance of 2 greater than the last one we've seen.
+                // That means there is an available type ID between these reserved IDs.
+                while (nextTypeIter != assignedIdsIterEnd) {
+                    if (nextTypeIter->first.packageId() != package->id.value()) {
+                        break;
+                    }
+
+                    const uint8_t typeId = nextTypeIter->first.typeId();
+                    if (typeId > nextExpectedTypeId) {
+                        // There is a gap in the type IDs, so use the missing one.
+                        type->id = nextExpectedTypeId++;
+                        break;
                     }
+
+                    // Set our expectation to be the next type ID after the reserved one we
+                    // just saw.
+                    nextExpectedTypeId = typeId + 1;
+
+                    // Move to the next reserved ID.
+                    ++nextTypeIter;
+                }
+
+                if (!type->id) {
+                    // We must have hit the end of the reserved IDs and not found a gap.
+                    // That means the next ID is available.
+                    type->id = nextExpectedTypeId++;
                 }
             }
 
-            // Assign unused entry IDs.
-            const auto endUsedEntryIter = usedEntryIds.end();
-            auto nextUsedEntryIter = usedEntryIds.begin();
-            uint16_t nextId = 0;
+            resourceId = ResourceId(package->id.value(), type->id.value(), 0);
+            uint16_t nextExpectedEntryId = 0;
+
+            // Find the closest matching ResourceId that is <= the one with only the package
+            // and type set.
+            auto nextEntryIter = assignedIds.lower_bound(resourceId);
             for (auto& entry : type->entries) {
                 if (!entry->id) {
-                    // Assign the next available entryID.
-                    while (nextUsedEntryIter != endUsedEntryIter &&
-                            nextId == *nextUsedEntryIter) {
-                        nextId++;
-                        ++nextUsedEntryIter;
+                    // We need to assign an entry ID. Iterate over the reserved IDs until we find
+                    // some entry ID that is a distance of 2 greater than the last one we've seen.
+                    // That means there is an available entry ID between these reserved IDs.
+                    while (nextEntryIter != assignedIdsIterEnd) {
+                        if (nextEntryIter->first.packageId() != package->id.value() ||
+                                nextEntryIter->first.typeId() != type->id.value()) {
+                            break;
+                        }
+
+                        const uint16_t entryId = nextEntryIter->first.entryId();
+                        if (entryId > nextExpectedEntryId) {
+                            // There is a gap in the entry IDs, so use the missing one.
+                            entry->id = nextExpectedEntryId++;
+                            break;
+                        }
+
+                        // Set our expectation to be the next type ID after the reserved one we
+                        // just saw.
+                        nextExpectedEntryId = entryId + 1;
+
+                        // Move to the next reserved entry ID.
+                        ++nextEntryIter;
                     }
-                    entry->id = nextId++;
-                }
-            }
-        }
 
-        // Assign unused type IDs.
-        size_t nextTypeId = 0;
-        for (auto& type : package->types) {
-            if (!type->id) {
-                while (nextTypeId < usedTypeIds.size() && usedTypeIds[nextTypeId]) {
-                    nextTypeId++;
+                    if (!entry->id) {
+                        // We must have hit the end of the reserved IDs and not found a gap.
+                        // That means the next ID is available.
+                        entry->id = nextExpectedEntryId++;
+                    }
                 }
-                type->id = static_cast<uint8_t>(nextTypeId);
-                nextTypeId++;
             }
         }
     }