AAPT2: Ensure style strings are always first in StringPool

Move the styled strings to a separate section of the StringPool so
that sorting can never mess up the order of Styles.

Bug: 63570514
Test: make aapt2_tests
Change-Id: Id2ce1355b92be1bb31ce0daa7e54ae9b5b6c2ffe

1 files changed, 37 insertions, 63 deletions

diff --git a/tools/aapt2/StringPool.h b/tools/aapt2/StringPool.h
index d1232a29b5aa..8350d0d09108 100644
--- a/tools/aapt2/StringPool.h
+++ b/tools/aapt2/StringPool.h
@@ -42,12 +42,16 @@ struct StyleString {
   std::vector<Span> spans;
 };
 
+// A StringPool for storing the value of String and StyledString resources.
+// Styles and Strings are stored separately, since the runtime variant of this
+// class -- ResStringPool -- requires that styled strings *always* appear first, since their
+// style data is stored as an array indexed by the same indices as the main string pool array.
+// Otherwise, the style data array would have to be sparse and take up more space.
 class StringPool {
  public:
   class Context {
    public:
     enum : uint32_t {
-      kStylePriority = 0u,
       kHighPriority = 1u,
       kNormalPriority = 0x7fffffffu,
       kLowPriority = 0xffffffffu,
@@ -58,8 +62,8 @@ class StringPool {
     Context() = default;
     Context(uint32_t p, const ConfigDescription& c) : priority(p), config(c) {}
     explicit Context(uint32_t p) : priority(p) {}
-    explicit Context(const ConfigDescription& c)
-        : priority(kNormalPriority), config(c) {}
+    explicit Context(const ConfigDescription& c) : priority(kNormalPriority), config(c) {
+    }
   };
 
   class Entry;
@@ -116,13 +120,14 @@ class StringPool {
    public:
     std::string value;
     Context context;
-    size_t index;
 
    private:
     friend class StringPool;
     friend class Ref;
 
+    size_t index_;
     int ref_;
+    const StringPool* pool_;
   };
 
   struct Span {
@@ -133,18 +138,18 @@ class StringPool {
 
   class StyleEntry {
    public:
-    Ref str;
+    std::string value;
+    Context context;
     std::vector<Span> spans;
 
    private:
     friend class StringPool;
     friend class StyleRef;
 
+    size_t index_;
     int ref_;
   };
 
-  using const_iterator = std::vector<std::unique_ptr<Entry>>::const_iterator;
-
   static bool FlattenUtf8(BigBuffer* out, const StringPool& pool);
   static bool FlattenUtf16(BigBuffer* out, const StringPool& pool);
 
@@ -152,92 +157,61 @@ class StringPool {
   StringPool(StringPool&&) = default;
   StringPool& operator=(StringPool&&) = default;
 
-  /**
-   * Adds a string to the pool, unless it already exists. Returns
-   * a reference to the string in the pool.
-   */
+  // Adds a string to the pool, unless it already exists. Returns a reference to the string in the
+  // pool.
   Ref MakeRef(const android::StringPiece& str);
 
-  /**
-   * Adds a string to the pool, unless it already exists, with a context
-   * object that can be used when sorting the string pool. Returns
-   * a reference to the string in the pool.
-   */
+  // Adds a string to the pool, unless it already exists, with a context object that can be used
+  // when sorting the string pool. Returns a reference to the string in the pool.
   Ref MakeRef(const android::StringPiece& str, const Context& context);
 
-  /**
-   * Adds a style to the string pool and returns a reference to it.
-   */
+  // Adds a style to the string pool and returns a reference to it.
   StyleRef MakeRef(const StyleString& str);
 
-  /**
-   * Adds a style to the string pool with a context object that
-   * can be used when sorting the string pool. Returns a reference
-   * to the style in the string pool.
-   */
+  // Adds a style to the string pool with a context object that can be used when sorting the string
+  // pool. Returns a reference to the style in the string pool.
   StyleRef MakeRef(const StyleString& str, const Context& context);
 
-  /**
-   * Adds a style from another string pool. Returns a reference to the
-   * style in the string pool.
-   */
+  // Adds a style from another string pool. Returns a reference to the style in the string pool.
   StyleRef MakeRef(const StyleRef& ref);
 
-  /**
-   * Moves pool into this one without coalescing strings. When this
-   * function returns, pool will be empty.
-   */
+  // Moves pool into this one without coalescing strings. When this function returns, pool will be
+  // empty.
   void Merge(StringPool&& pool);
 
-  /**
-   * Returns the number of strings in the table.
-   */
-  inline size_t size() const;
+  inline const std::vector<std::unique_ptr<Entry>>& strings() const {
+    return strings_;
+  }
 
-  /**
-   * Reserves space for strings and styles as an optimization.
-   */
+  // Returns the number of strings in the table.
+  inline size_t size() const {
+    return styles_.size() + strings_.size();
+  }
+
+  // Reserves space for strings and styles as an optimization.
   void HintWillAdd(size_t string_count, size_t style_count);
 
-  /**
-   * Sorts the strings according to some comparison function.
-   */
-  void Sort(const std::function<bool(const Entry&, const Entry&)>& cmp);
+  // Sorts the strings according to their Context using some comparison function.
+  // Equal Contexts are further sorted by string value, lexicographically.
+  // If no comparison function is provided, values are only sorted lexicographically.
+  void Sort(const std::function<int(const Context&, const Context&)>& cmp = nullptr);
 
-  /**
-   * Removes any strings that have no references.
-   */
+  // Removes any strings that have no references.
   void Prune();
 
  private:
   DISALLOW_COPY_AND_ASSIGN(StringPool);
 
-  friend const_iterator begin(const StringPool& pool);
-  friend const_iterator end(const StringPool& pool);
-
   static bool Flatten(BigBuffer* out, const StringPool& pool, bool utf8);
 
   Ref MakeRefImpl(const android::StringPiece& str, const Context& context, bool unique);
+  void ReAssignIndices();
 
   std::vector<std::unique_ptr<Entry>> strings_;
   std::vector<std::unique_ptr<StyleEntry>> styles_;
   std::unordered_multimap<android::StringPiece, Entry*> indexed_strings_;
 };
 
-//
-// Inline implementation
-//
-
-inline size_t StringPool::size() const { return strings_.size(); }
-
-inline StringPool::const_iterator begin(const StringPool& pool) {
-  return pool.strings_.begin();
-}
-
-inline StringPool::const_iterator end(const StringPool& pool) {
-  return pool.strings_.end();
-}
-
 }  // namespace aapt
 
 #endif  // AAPT_STRING_POOL_H