1 files changed, 58 insertions, 12 deletions

diff --git a/tools/aapt2/java/ProguardRules.h b/tools/aapt2/java/ProguardRules.h
index 3c349bab1217..46827ee7cf93 100644
--- a/tools/aapt2/java/ProguardRules.h
+++ b/tools/aapt2/java/ProguardRules.h
@@ -22,37 +22,83 @@
 #include <set>
 #include <string>
 
+#include "androidfw/StringPiece.h"
+
 #include "Resource.h"
+#include "ResourceTable.h"
 #include "Source.h"
+#include "ValueVisitor.h"
+#include "io/Io.h"
+#include "process/IResourceTableConsumer.h"
 #include "xml/XmlDom.h"
 
 namespace aapt {
 namespace proguard {
 
+struct UsageLocation {
+  ResourceName name;
+  Source source;
+};
+
 class KeepSet {
  public:
-  inline void AddClass(const Source& source, const std::string& class_name) {
-    keep_set_[class_name].insert(source);
+  KeepSet() = default;
+
+  KeepSet(bool conditional_keep_rules) : conditional_keep_rules_(conditional_keep_rules) {
   }
 
-  inline void AddMethod(const Source& source, const std::string& method_name) {
-    keep_method_set_[method_name].insert(source);
+  inline void AddManifestClass(const UsageLocation& file, const std::string& class_name) {
+    manifest_class_set_[class_name].insert(file);
+  }
+
+  inline void AddConditionalClass(const UsageLocation& file, const std::string& class_name) {
+    conditional_class_set_[class_name].insert(file);
+  }
+
+  inline void AddMethod(const UsageLocation& file, const std::string& method_name) {
+    method_set_[method_name].insert(file);
+  }
+
+  inline void AddReference(const UsageLocation& file, const ResourceName& resource_name) {
+    reference_set_[resource_name].insert(file);
   }
 
  private:
-  friend bool WriteKeepSet(std::ostream* out, const KeepSet& keep_set);
+  friend void WriteKeepSet(const KeepSet& keep_set, io::OutputStream* out);
+
+  friend bool CollectLocations(const UsageLocation& location, const KeepSet& keep_set,
+                               std::set<UsageLocation>* locations);
 
-  std::map<std::string, std::set<Source>> keep_set_;
-  std::map<std::string, std::set<Source>> keep_method_set_;
+  bool conditional_keep_rules_ = false;
+  std::map<std::string, std::set<UsageLocation>> manifest_class_set_;
+  std::map<std::string, std::set<UsageLocation>> method_set_;
+  std::map<std::string, std::set<UsageLocation>> conditional_class_set_;
+  std::map<ResourceName, std::set<UsageLocation>> reference_set_;
 };
 
-bool CollectProguardRulesForManifest(const Source& source,
-                                     xml::XmlResource* res, KeepSet* keep_set,
+bool CollectProguardRulesForManifest(xml::XmlResource* res, KeepSet* keep_set,
                                      bool main_dex_only = false);
-bool CollectProguardRules(const Source& source, xml::XmlResource* res,
-                          KeepSet* keep_set);
 
-bool WriteKeepSet(std::ostream* out, const KeepSet& keep_set);
+bool CollectProguardRules(xml::XmlResource* res, KeepSet* keep_set);
+
+bool CollectResourceReferences(IAaptContext* context, ResourceTable* table, KeepSet* keep_set);
+
+void WriteKeepSet(const KeepSet& keep_set, io::OutputStream* out);
+
+bool CollectLocations(const UsageLocation& location, const KeepSet& keep_set,
+                      std::set<UsageLocation>* locations);
+
+//
+// UsageLocation implementation.
+//
+
+inline bool operator==(const UsageLocation& lhs, const UsageLocation& rhs) {
+  return lhs.name == rhs.name;
+}
+
+inline int operator<(const UsageLocation& lhs, const UsageLocation& rhs) {
+  return lhs.name.compare(rhs.name);
+}
 
 }  // namespace proguard
 }  // namespace aapt