AAPT2: Add support to specify stable IDs

The --stable-ids flag allows the user to specify a file containing
a list of resource name and resource ID pairs in the form of:

package:type/name = 0xPPTTEEEE

This assigns the given resource the specified ID. It helps ensure
that when adding or removing resources, IDs are assigned in a stable
fashion.

If a package, type, or name is not found, no error or warning is
raised.

Change-Id: Ibc2f4e05cc924be255fedd862d835cb5b18d7584
diff --git a/tools/aapt2/compile/IdAssigner.cpp b/tools/aapt2/compile/IdAssigner.cpp
index 341c9b3..501ae9d 100644
--- a/tools/aapt2/compile/IdAssigner.cpp
+++ b/tools/aapt2/compile/IdAssigner.cpp
@@ -19,87 +19,180 @@
 #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();
-
-        // Type ID 0 is invalid, reserve it.
-        usedTypeIds.set(0);
-
-        // Collect used type IDs.
         for (auto& type : package->types) {
-            if (type->id) {
-                usedEntryIds.clear();
+            for (auto& entry : type->entries) {
+                const ResourceName name(package->name, type->type, entry->name);
 
-                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 (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;
+                        }
+                    }
                 }
 
-                // Mark the type ID as taken.
-                usedTypeIds.set(type->id.value());
-            }
-
-            // 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);
+                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;
                     }
                 }
             }
+        }
+    }
 
-            // Assign unused entry IDs.
-            const auto endUsedEntryIter = usedEntryIds.end();
-            auto nextUsedEntryIter = usedEntryIds.begin();
-            uint16_t nextId = 0;
-            for (auto& entry : type->entries) {
-                if (!entry->id) {
-                    // Assign the next available entryID.
-                    while (nextUsedEntryIter != endUsedEntryIter &&
-                            nextId == *nextUsedEntryIter) {
-                        nextId++;
-                        ++nextUsedEntryIter;
-                    }
-                    entry->id = nextId++;
-                }
+    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;
             }
         }
+    }
 
-        // Assign unused type IDs.
-        size_t nextTypeId = 0;
+    // 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) {
-                while (nextTypeId < usedTypeIds.size() && usedTypeIds[nextTypeId]) {
-                    nextTypeId++;
+                // 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;
                 }
-                type->id = static_cast<uint8_t>(nextTypeId);
-                nextTypeId++;
+
+                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++;
+                }
+            }
+
+            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) {
+                    // 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;
+                    }
+
+                    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++;
+                    }
+                }
             }
         }
     }
diff --git a/tools/aapt2/compile/IdAssigner.h b/tools/aapt2/compile/IdAssigner.h
index 514df3a..06cd5e3 100644
--- a/tools/aapt2/compile/IdAssigner.h
+++ b/tools/aapt2/compile/IdAssigner.h
@@ -17,16 +17,29 @@
 #ifndef AAPT_COMPILE_IDASSIGNER_H
 #define AAPT_COMPILE_IDASSIGNER_H
 
+#include "Resource.h"
 #include "process/IResourceTableConsumer.h"
 
+#include <android-base/macros.h>
+#include <unordered_map>
+
 namespace aapt {
 
 /**
  * Assigns IDs to each resource in the table, respecting existing IDs and filling in gaps
  * in between fixed ID assignments.
  */
-struct IdAssigner : public IResourceTableConsumer {
+class IdAssigner : public IResourceTableConsumer {
+public:
+    IdAssigner() = default;
+    explicit IdAssigner(const std::unordered_map<ResourceName, ResourceId>* map) :
+            mAssignedIdMap(map) {
+    }
+
     bool consume(IAaptContext* context, ResourceTable* table) override;
+
+private:
+    const std::unordered_map<ResourceName, ResourceId>* mAssignedIdMap = nullptr;
 };
 
 } // namespace aapt
diff --git a/tools/aapt2/compile/IdAssigner_test.cpp b/tools/aapt2/compile/IdAssigner_test.cpp
index 802e99a..4f43c40 100644
--- a/tools/aapt2/compile/IdAssigner_test.cpp
+++ b/tools/aapt2/compile/IdAssigner_test.cpp
@@ -15,11 +15,7 @@
  */
 
 #include "compile/IdAssigner.h"
-
-#include "test/Context.h"
-#include "test/Builders.h"
-
-#include <gtest/gtest.h>
+#include "test/Test.h"
 
 namespace aapt {
 
@@ -42,9 +38,14 @@
 
 TEST(IdAssignerTest, AssignIdsWithReservedIds) {
     std::unique_ptr<ResourceTable> table = test::ResourceTableBuilder()
+            .addSimple("@android:id/foo", ResourceId(0x01010000))
+            .addSimple("@android:dimen/two")
+            .addSimple("@android:integer/three")
+            .addSimple("@android:string/five")
+            .addSimple("@android:attr/fun", ResourceId(0x01040000))
             .addSimple("@android:attr/foo", ResourceId(0x01040006))
             .addSimple("@android:attr/bar")
-            .addSimple("@android:id/foo")
+            .addSimple("@android:attr/baz")
             .addSimple("@app:id/biz")
             .setPackageId("android", 0x01)
             .setPackageId("app", 0x7f)
@@ -55,6 +56,34 @@
 
     ASSERT_TRUE(assigner.consume(context.get(), table.get()));
     ASSERT_TRUE(verifyIds(table.get()));
+
+    Maybe<ResourceTable::SearchResult> maybeResult;
+
+    // Expect to fill in the gaps between 0x0101XXXX and 0x0104XXXX.
+
+    maybeResult = table->findResource(test::parseNameOrDie("@android:dimen/two"));
+    AAPT_ASSERT_TRUE(maybeResult);
+    EXPECT_EQ(make_value<uint8_t>(2), maybeResult.value().type->id);
+
+    maybeResult = table->findResource(test::parseNameOrDie("@android:integer/three"));
+    AAPT_ASSERT_TRUE(maybeResult);
+    EXPECT_EQ(make_value<uint8_t>(3), maybeResult.value().type->id);
+
+    // Expect to bypass the reserved 0x0104XXXX IDs and use the next 0x0105XXXX IDs.
+
+    maybeResult = table->findResource(test::parseNameOrDie("@android:string/five"));
+    AAPT_ASSERT_TRUE(maybeResult);
+    EXPECT_EQ(make_value<uint8_t>(5), maybeResult.value().type->id);
+
+    // Expect to fill in the gaps between 0x01040000 and 0x01040006.
+
+    maybeResult = table->findResource(test::parseNameOrDie("@android:attr/bar"));
+    AAPT_ASSERT_TRUE(maybeResult);
+    EXPECT_EQ(make_value<uint16_t>(1), maybeResult.value().entry->id);
+
+    maybeResult = table->findResource(test::parseNameOrDie("@android:attr/baz"));
+    AAPT_ASSERT_TRUE(maybeResult);
+    EXPECT_EQ(make_value<uint16_t>(2), maybeResult.value().entry->id);
 }
 
 TEST(IdAssignerTest, FailWhenNonUniqueIdsAssigned) {
@@ -71,6 +100,29 @@
     ASSERT_FALSE(assigner.consume(context.get(), table.get()));
 }
 
+TEST(IdAssignerTest, AssignIdsWithIdMap) {
+    std::unique_ptr<ResourceTable> table = test::ResourceTableBuilder()
+            .addSimple("@android:attr/foo")
+            .addSimple("@android:attr/bar")
+            .setPackageId("android", 0x01)
+            .build();
+
+    std::unique_ptr<IAaptContext> context = test::ContextBuilder().build();
+    std::unordered_map<ResourceName, ResourceId> idMap = {
+            { test::parseNameOrDie("@android:attr/foo"), ResourceId(0x01010002) } };
+    IdAssigner assigner(&idMap);
+    ASSERT_TRUE(assigner.consume(context.get(), table.get()));
+    ASSERT_TRUE(verifyIds(table.get()));
+    Maybe<ResourceTable::SearchResult> result = table->findResource(
+            test::parseNameOrDie("@android:attr/foo"));
+    AAPT_ASSERT_TRUE(result);
+
+    const ResourceTable::SearchResult& searchResult = result.value();
+    EXPECT_EQ(make_value<uint8_t>(0x01), searchResult.package->id);
+    EXPECT_EQ(make_value<uint8_t>(0x01), searchResult.type->id);
+    EXPECT_EQ(make_value<uint16_t>(0x0002), searchResult.entry->id);
+}
+
 ::testing::AssertionResult verifyIds(ResourceTable* table) {
     std::set<uint8_t> packageIds;
     for (auto& package : table->packages) {