Merge changes I6f6b4d78,I874696cf into rvc-dev

* changes:
  Only send reset state annotation when needed
  Group annotations for the same atom id
diff --git a/tools/stats_log_api_gen/Collation.cpp b/tools/stats_log_api_gen/Collation.cpp
index bf39093..1387bb5 100644
--- a/tools/stats_log_api_gen/Collation.cpp
+++ b/tools/stats_log_api_gen/Collation.cpp
@@ -454,16 +454,12 @@
     return has_attribution_node;
 }
 
-static void populateFieldNumberToAnnotations(const AtomDecl& atomDecl,
-                                             FieldNumberToAnnotations* fieldNumberToAnnotations) {
-    for (FieldNumberToAnnotations::const_iterator it = atomDecl.fieldNumberToAnnotations.begin();
-         it != atomDecl.fieldNumberToAnnotations.end(); it++) {
+static void populateFieldNumberToAtomDeclSet(const shared_ptr<AtomDecl>& atomDecl,
+                                             FieldNumberToAtomDeclSet* fieldNumberToAtomDeclSet) {
+    for (FieldNumberToAnnotations::const_iterator it = atomDecl->fieldNumberToAnnotations.begin();
+         it != atomDecl->fieldNumberToAnnotations.end(); it++) {
         const int fieldNumber = it->first;
-        const set<shared_ptr<Annotation>>& insertAnnotationsSource = it->second;
-        set<shared_ptr<Annotation>>& insertAnnotationsTarget =
-                (*fieldNumberToAnnotations)[fieldNumber];
-        insertAnnotationsTarget.insert(insertAnnotationsSource.begin(),
-                                       insertAnnotationsSource.end());
+        (*fieldNumberToAtomDeclSet)[fieldNumber].insert(atomDecl);
     }
 }
 
@@ -513,18 +509,19 @@
         }
 
         const Descriptor* atom = atomField->message_type();
-        AtomDecl atomDecl(atomField->number(), atomField->name(), atom->name());
+        shared_ptr<AtomDecl> atomDecl =
+                make_shared<AtomDecl>(atomField->number(), atomField->name(), atom->name());
 
         if (atomField->options().GetExtension(os::statsd::allow_from_any_uid) == true) {
-            atomDecl.whitelisted = true;
+            atomDecl->whitelisted = true;
             if (dbg) {
                 printf("%s is whitelisted\n", atomField->name().c_str());
             }
         }
 
-        if (atomDecl.code < PULL_ATOM_START_ID &&
+        if (atomDecl->code < PULL_ATOM_START_ID &&
             atomField->options().GetExtension(os::statsd::truncate_timestamp)) {
-            addAnnotationToAtomDecl(&atomDecl, ATOM_ID_FIELD_NUMBER,
+            addAnnotationToAtomDecl(atomDecl.get(), ATOM_ID_FIELD_NUMBER,
                                     ANNOTATION_ID_TRUNCATE_TIMESTAMP, ANNOTATION_TYPE_BOOL,
                                     AnnotationValue(true));
             if (dbg) {
@@ -533,29 +530,33 @@
         }
 
         vector<java_type_t> signature;
-        errorCount += collate_atom(atom, &atomDecl, &signature);
-        if (atomDecl.primaryFields.size() != 0 && atomDecl.exclusiveField == 0) {
+        errorCount += collate_atom(atom, atomDecl.get(), &signature);
+        if (atomDecl->primaryFields.size() != 0 && atomDecl->exclusiveField == 0) {
             print_error(atomField, "Cannot have a primary field without an exclusive field: %s\n",
                         atomField->name().c_str());
             errorCount++;
             continue;
         }
 
-        atoms->decls.insert(atomDecl);
-        FieldNumberToAnnotations& fieldNumberToAnnotations = atoms->signatureInfoMap[signature];
-        populateFieldNumberToAnnotations(atomDecl, &fieldNumberToAnnotations);
+        FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet = atoms->signatureInfoMap[signature];
+        populateFieldNumberToAtomDeclSet(atomDecl, &fieldNumberToAtomDeclSet);
 
-        AtomDecl nonChainedAtomDecl(atomField->number(), atomField->name(), atom->name());
+        atoms->decls.insert(atomDecl);
+
+        shared_ptr<AtomDecl> nonChainedAtomDecl =
+                make_shared<AtomDecl>(atomField->number(), atomField->name(), atom->name());
         vector<java_type_t> nonChainedSignature;
-        if (get_non_chained_node(atom, &nonChainedAtomDecl, &nonChainedSignature)) {
-            atoms->non_chained_decls.insert(nonChainedAtomDecl);
-            FieldNumberToAnnotations& fieldNumberToAnnotations =
+        if (get_non_chained_node(atom, nonChainedAtomDecl.get(), &nonChainedSignature)) {
+            FieldNumberToAtomDeclSet& nonChainedFieldNumberToAtomDeclSet =
                     atoms->nonChainedSignatureInfoMap[nonChainedSignature];
-            populateFieldNumberToAnnotations(atomDecl, &fieldNumberToAnnotations);
+            populateFieldNumberToAtomDeclSet(nonChainedAtomDecl,
+                                             &nonChainedFieldNumberToAtomDeclSet);
+
+            atoms->non_chained_decls.insert(nonChainedAtomDecl);
         }
 
-        if (atomDecl.code < PULL_ATOM_START_ID && atomDecl.code > maxPushedAtomId) {
-            maxPushedAtomId = atomDecl.code;
+        if (atomDecl->code < PULL_ATOM_START_ID && atomDecl->code > maxPushedAtomId) {
+            maxPushedAtomId = atomDecl->code;
         }
     }
 
@@ -563,8 +564,7 @@
 
     if (dbg) {
         printf("signatures = [\n");
-        for (map<vector<java_type_t>, FieldNumberToAnnotations>::const_iterator it =
-                     atoms->signatureInfoMap.begin();
+        for (SignatureInfoMap::const_iterator it = atoms->signatureInfoMap.begin();
              it != atoms->signatureInfoMap.end(); it++) {
             printf("   ");
             for (vector<java_type_t>::const_iterator jt = it->first.begin(); jt != it->first.end();
diff --git a/tools/stats_log_api_gen/Collation.h b/tools/stats_log_api_gen/Collation.h
index 2aedb21..b1cf134 100644
--- a/tools/stats_log_api_gen/Collation.h
+++ b/tools/stats_log_api_gen/Collation.h
@@ -110,7 +110,16 @@
     }
 };
 
-using FieldNumberToAnnotations = map<int, set<shared_ptr<Annotation>>>;
+struct SharedComparator {
+    template <typename T>
+    inline bool operator()(const shared_ptr<T>& lhs, const shared_ptr<T>& rhs) const {
+        return (*lhs) < (*rhs);
+    }
+};
+
+using AnnotationSet = set<shared_ptr<Annotation>, SharedComparator>;
+
+using FieldNumberToAnnotations = map<int, AnnotationSet>;
 
 /**
  * The name and type for an atom field.
@@ -169,11 +178,19 @@
     }
 };
 
+using AtomDeclSet = set<shared_ptr<AtomDecl>, SharedComparator>;
+
+// Maps a field number to a set of atoms that have annotation(s) for their field with that field
+// number.
+using FieldNumberToAtomDeclSet = map<int, AtomDeclSet>;
+
+using SignatureInfoMap = map<vector<java_type_t>, FieldNumberToAtomDeclSet>;
+
 struct Atoms {
-    map<vector<java_type_t>, FieldNumberToAnnotations> signatureInfoMap;
-    set<AtomDecl> decls;
-    set<AtomDecl> non_chained_decls;
-    map<vector<java_type_t>, FieldNumberToAnnotations> nonChainedSignatureInfoMap;
+    SignatureInfoMap signatureInfoMap;
+    AtomDeclSet decls;
+    AtomDeclSet non_chained_decls;
+    SignatureInfoMap nonChainedSignatureInfoMap;
     int maxPushedAtomId;
 };
 
diff --git a/tools/stats_log_api_gen/atoms_info_writer.cpp b/tools/stats_log_api_gen/atoms_info_writer.cpp
index 6944752..862fed4 100644
--- a/tools/stats_log_api_gen/atoms_info_writer.cpp
+++ b/tools/stats_log_api_gen/atoms_info_writer.cpp
@@ -62,11 +62,11 @@
     fprintf(out,
             "const std::set<int> "
             "AtomsInfo::kTruncatingTimestampAtomBlackList = {\n");
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        if (kTruncatingAtomNames.find(atom->name) != kTruncatingAtomNames.end()) {
-            const string constant = make_constant_name(atom->name);
-            fprintf(out, "    %d, // %s\n", atom->code, constant.c_str());
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        if (kTruncatingAtomNames.find((*atomIt)->name) != kTruncatingAtomNames.end()) {
+            const string constant = make_constant_name((*atomIt)->name);
+            fprintf(out, "    %d, // %s\n", (*atomIt)->code, constant.c_str());
         }
     }
 
@@ -74,13 +74,13 @@
     fprintf(out, "\n");
 
     fprintf(out, "const std::set<int> AtomsInfo::kAtomsWithAttributionChain = {\n");
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        for (vector<AtomField>::const_iterator field = atom->fields.begin();
-             field != atom->fields.end(); field++) {
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        for (vector<AtomField>::const_iterator field = (*atomIt)->fields.begin();
+             field != (*atomIt)->fields.end(); field++) {
             if (field->javaType == JAVA_TYPE_ATTRIBUTION_CHAIN) {
-                const string constant = make_constant_name(atom->name);
-                fprintf(out, "    %d, // %s\n", atom->code, constant.c_str());
+                const string constant = make_constant_name((*atomIt)->name);
+                fprintf(out, "    %d, // %s\n", (*atomIt)->code, constant.c_str());
                 break;
             }
         }
@@ -90,11 +90,11 @@
     fprintf(out, "\n");
 
     fprintf(out, "const std::set<int> AtomsInfo::kWhitelistedAtoms = {\n");
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        if (atom->whitelisted) {
-            const string constant = make_constant_name(atom->name);
-            fprintf(out, "    %d, // %s\n", atom->code, constant.c_str());
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        if ((*atomIt)->whitelisted) {
+            const string constant = make_constant_name((*atomIt)->name);
+            fprintf(out, "    %d, // %s\n", (*atomIt)->code, constant.c_str());
         }
     }
 
@@ -103,17 +103,17 @@
 
     fprintf(out, "static std::map<int, int> getAtomUidField() {\n");
     fprintf(out, "    std::map<int, int> uidField;\n");
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        if (atom->uidField == 0) {
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        if ((*atomIt)->uidField == 0) {
             continue;
         }
         fprintf(out,
                 "\n    // Adding uid field for atom "
                 "(%d)%s\n",
-                atom->code, atom->name.c_str());
-        fprintf(out, "    uidField[%d /* %s */] = %d;\n", atom->code,
-                make_constant_name(atom->name).c_str(), atom->uidField);
+                (*atomIt)->code, (*atomIt)->name.c_str());
+        fprintf(out, "    uidField[%d /* %s */] = %d;\n", (*atomIt)->code,
+                make_constant_name((*atomIt)->name).c_str(), (*atomIt)->uidField);
     }
 
     fprintf(out, "    return uidField;\n");
@@ -128,35 +128,35 @@
             "getStateAtomFieldOptions() {\n");
     fprintf(out, "    std::map<int, StateAtomFieldOptions> options;\n");
     fprintf(out, "    StateAtomFieldOptions* opt;\n");
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        if (atom->primaryFields.size() == 0 && atom->exclusiveField == 0) {
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        if ((*atomIt)->primaryFields.size() == 0 && (*atomIt)->exclusiveField == 0) {
             continue;
         }
         fprintf(out,
                 "\n    // Adding primary and exclusive fields for atom "
                 "(%d)%s\n",
-                atom->code, atom->name.c_str());
-        fprintf(out, "    opt = &(options[%d /* %s */]);\n", atom->code,
-                make_constant_name(atom->name).c_str());
-        fprintf(out, "    opt->primaryFields.reserve(%lu);\n", atom->primaryFields.size());
-        for (const auto& field : atom->primaryFields) {
+                (*atomIt)->code, (*atomIt)->name.c_str());
+        fprintf(out, "    opt = &(options[%d /* %s */]);\n", (*atomIt)->code,
+                make_constant_name((*atomIt)->name).c_str());
+        fprintf(out, "    opt->primaryFields.reserve(%lu);\n", (*atomIt)->primaryFields.size());
+        for (const auto& field : (*atomIt)->primaryFields) {
             fprintf(out, "    opt->primaryFields.push_back(%d);\n", field);
         }
 
-        fprintf(out, "    opt->exclusiveField = %d;\n", atom->exclusiveField);
-        if (atom->defaultState != INT_MAX) {
-            fprintf(out, "    opt->defaultState = %d;\n", atom->defaultState);
+        fprintf(out, "    opt->exclusiveField = %d;\n", (*atomIt)->exclusiveField);
+        if ((*atomIt)->defaultState != INT_MAX) {
+            fprintf(out, "    opt->defaultState = %d;\n", (*atomIt)->defaultState);
         } else {
             fprintf(out, "    opt->defaultState = UNSET_VALUE;\n");
         }
 
-        if (atom->resetState != INT_MAX) {
-            fprintf(out, "    opt->resetState = %d;\n", atom->resetState);
+        if ((*atomIt)->resetState != INT_MAX) {
+            fprintf(out, "    opt->resetState = %d;\n", (*atomIt)->resetState);
         } else {
             fprintf(out, "    opt->resetState = UNSET_VALUE;\n");
         }
-        fprintf(out, "    opt->nested = %d;\n", atom->nested);
+        fprintf(out, "    opt->nested = %d;\n", (*atomIt)->nested);
     }
 
     fprintf(out, "    return options;\n");
diff --git a/tools/stats_log_api_gen/java_writer.cpp b/tools/stats_log_api_gen/java_writer.cpp
index 5a22b5c..3eabb14 100644
--- a/tools/stats_log_api_gen/java_writer.cpp
+++ b/tools/stats_log_api_gen/java_writer.cpp
@@ -22,9 +22,8 @@
 namespace android {
 namespace stats_log_api_gen {
 
-static int write_java_q_logger_class(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
-        const AtomDecl& attributionDecl) {
+static int write_java_q_logger_class(FILE* out, const SignatureInfoMap& signatureInfoMap,
+                                     const AtomDecl& attributionDecl) {
     fprintf(out, "\n");
     fprintf(out, "    // Write logging helper methods for statsd in Q and earlier.\n");
     fprintf(out, "    private static class QLogger {\n");
@@ -41,47 +40,64 @@
 }
 
 static void write_annotations(FILE* out, int argIndex,
-                              const FieldNumberToAnnotations& fieldNumberToAnnotations) {
-    auto it = fieldNumberToAnnotations.find(argIndex);
-    if (it == fieldNumberToAnnotations.end()) {
+                              const FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet) {
+    FieldNumberToAtomDeclSet::const_iterator fieldNumberToAtomDeclSetIt =
+            fieldNumberToAtomDeclSet.find(argIndex);
+    if (fieldNumberToAtomDeclSet.end() == fieldNumberToAtomDeclSetIt) {
         return;
     }
-    const set<shared_ptr<Annotation>>& annotations = it->second;
-    for (auto& annotation : annotations) {
-        // TODO(b/151744250): Group annotations for same atoms.
-        // TODO(b/151786433): Write atom constant name instead of atom id literal.
-        fprintf(out, "        if (code == %d) {\n", annotation->atomId);
-        switch (annotation->type) {
-            case ANNOTATION_TYPE_INT:
+    const AtomDeclSet& atomDeclSet = fieldNumberToAtomDeclSetIt->second;
+    for (const shared_ptr<AtomDecl>& atomDecl : atomDeclSet) {
+        fprintf(out, "        if (code == %d) {\n", atomDecl->code);
+        const AnnotationSet& annotations = atomDecl->fieldNumberToAnnotations.at(argIndex);
+        int resetState = -1;
+        int defaultState = -1;
+        for (const shared_ptr<Annotation>& annotation : annotations) {
+            // TODO(b/151786433): Write atom constant name instead of atom id literal.
+            switch (annotation->type) {
                 // TODO(b/151776731): Check for reset state annotation and only include
                 // reset state when field value == default state annotation value.
-                // TODO(b/151786433): Write annotation constant name instead of
-                // annotation id literal.
-                fprintf(out, "            builder.addIntAnnotation((byte) %d, %d);\n",
-                        annotation->annotationId, annotation->value.intValue);
-                break;
-            case ANNOTATION_TYPE_BOOL:
-                // TODO(b/151786433): Write annotation constant name instead of
-                // annotation id literal.
-                fprintf(out, "            builder.addBooleanAnnotation((byte) %d, %s);\n",
-                        annotation->annotationId, annotation->value.boolValue ? "true" : "false");
-                break;
-            default:
-                break;
+                case ANNOTATION_TYPE_INT:
+                    // TODO(b/151786433): Write annotation constant name instead of
+                    // annotation id literal.
+                    if (ANNOTATION_ID_RESET_STATE == annotation->annotationId) {
+                        resetState = annotation->value.intValue;
+                    } else if (ANNOTATION_ID_DEFAULT_STATE == annotation->annotationId) {
+                        defaultState = annotation->value.intValue;
+                    } else {
+                        fprintf(out, "            builder.addIntAnnotation((byte) %d, %d);\n",
+                                annotation->annotationId, annotation->value.intValue);
+                    }
+                    break;
+                case ANNOTATION_TYPE_BOOL:
+                    // TODO(b/151786433): Write annotation constant name instead of
+                    // annotation id literal.
+                    fprintf(out, "            builder.addBooleanAnnotation((byte) %d, %s);\n",
+                            annotation->annotationId,
+                            annotation->value.boolValue ? "true" : "false");
+                    break;
+                default:
+                    break;
+            }
+        }
+        if (defaultState != -1 && resetState != -1) {
+            fprintf(out, "            if (arg%d == %d) {\n", argIndex, resetState);
+            fprintf(out, "                builder.addIntAnnotation((byte) %d, %d);\n",
+                    ANNOTATION_ID_RESET_STATE, defaultState);
+            fprintf(out, "            }\n");
         }
         fprintf(out, "        }\n");
     }
 }
 
-static int write_java_methods(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
-        const AtomDecl& attributionDecl, const bool supportQ) {
+static int write_java_methods(FILE* out, const SignatureInfoMap& signatureInfoMap,
+                              const AtomDecl& attributionDecl, const bool supportQ) {
     for (auto signatureInfoMapIt = signatureInfoMap.begin();
          signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
         // Print method signature.
         fprintf(out, "    public static void write(int code");
         const vector<java_type_t>& signature = signatureInfoMapIt->first;
-        const FieldNumberToAnnotations& fieldNumberToAnnotations = signatureInfoMapIt->second;
+        const FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet = signatureInfoMapIt->second;
         int argIndex = 1;
         for (vector<java_type_t>::const_iterator arg = signature.begin(); arg != signature.end();
              arg++) {
@@ -120,7 +136,7 @@
 
         // Write atom code.
         fprintf(out, "%s        builder.setAtomId(code);\n", indent.c_str());
-        write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAnnotations);
+        write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAtomDeclSet);
 
         // Write the args.
         argIndex = 1;
@@ -231,7 +247,7 @@
                     fprintf(stderr, "Encountered unsupported type.");
                     return 1;
             }
-            write_annotations(out, argIndex, fieldNumberToAnnotations);
+            write_annotations(out, argIndex, fieldNumberToAtomDeclSet);
             argIndex++;
         }
 
diff --git a/tools/stats_log_api_gen/java_writer_q.cpp b/tools/stats_log_api_gen/java_writer_q.cpp
index 7d22583..d21e270 100644
--- a/tools/stats_log_api_gen/java_writer_q.cpp
+++ b/tools/stats_log_api_gen/java_writer_q.cpp
@@ -53,9 +53,8 @@
     fprintf(out, "%sprivate static final int LIST_TYPE_OVERHEAD = 2;\n", indent.c_str());
 }
 
-int write_java_methods_q_schema(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
-        const AtomDecl& attributionDecl, const string& indent) {
+int write_java_methods_q_schema(FILE* out, const SignatureInfoMap& signatureInfoMap,
+                                const AtomDecl& attributionDecl, const string& indent) {
     int requiredHelpers = 0;
     for (auto signatureInfoMapIt = signatureInfoMap.begin();
          signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
diff --git a/tools/stats_log_api_gen/java_writer_q.h b/tools/stats_log_api_gen/java_writer_q.h
index f1cfc44..c511a84 100644
--- a/tools/stats_log_api_gen/java_writer_q.h
+++ b/tools/stats_log_api_gen/java_writer_q.h
@@ -32,9 +32,8 @@
 
 void write_java_q_logging_constants(FILE* out, const string& indent);
 
-int write_java_methods_q_schema(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
-        const AtomDecl& attributionDecl, const string& indent);
+int write_java_methods_q_schema(FILE* out, const SignatureInfoMap& signatureInfoMap,
+                                const AtomDecl& attributionDecl, const string& indent);
 
 void write_java_helpers_for_q_schema_methods(FILE* out, const AtomDecl& attributionDecl,
                                              const int requiredHelpers, const string& indent);
diff --git a/tools/stats_log_api_gen/native_writer.cpp b/tools/stats_log_api_gen/native_writer.cpp
index 0cf3225..c0d73fa 100644
--- a/tools/stats_log_api_gen/native_writer.cpp
+++ b/tools/stats_log_api_gen/native_writer.cpp
@@ -22,35 +22,53 @@
 namespace stats_log_api_gen {
 
 static void write_annotations(FILE* out, int argIndex,
-                              const FieldNumberToAnnotations& fieldNumberToAnnotations,
+                              const FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet,
                               const string& methodPrefix, const string& methodSuffix) {
-    auto fieldNumberToAnnotationsIt = fieldNumberToAnnotations.find(argIndex);
-    if (fieldNumberToAnnotationsIt == fieldNumberToAnnotations.end()) {
+    FieldNumberToAtomDeclSet::const_iterator fieldNumberToAtomDeclSetIt =
+            fieldNumberToAtomDeclSet.find(argIndex);
+    if (fieldNumberToAtomDeclSet.end() == fieldNumberToAtomDeclSetIt) {
         return;
     }
-    const set<shared_ptr<Annotation>>& annotations = fieldNumberToAnnotationsIt->second;
-    for (const shared_ptr<Annotation>& annotation : annotations) {
-        // TODO(b/151744250): Group annotations for same atoms.
-        // TODO(b/151786433): Write atom constant name instead of atom id literal.
-        fprintf(out, "    if (code == %d) {\n", annotation->atomId);
-        switch (annotation->type) {
-            // TODO(b/151776731): Check for reset state annotation and only include
-            // reset state when field value == default state annotation value.
-            case ANNOTATION_TYPE_INT:
-                // TODO(b/151786433): Write annotation constant name instead of
-                // annotation id literal.
-                fprintf(out, "        %saddInt32Annotation(%s%d, %d);\n", methodPrefix.c_str(),
-                        methodSuffix.c_str(), annotation->annotationId, annotation->value.intValue);
-                break;
-            case ANNOTATION_TYPE_BOOL:
-                // TODO(b/151786433): Write annotation constant name instead of
-                // annotation id literal.
-                fprintf(out, "        %saddBoolAnnotation(%s%d, %s);\n", methodPrefix.c_str(),
-                        methodSuffix.c_str(), annotation->annotationId,
-                        annotation->value.boolValue ? "true" : "false");
-                break;
-            default:
-                break;
+    const AtomDeclSet& atomDeclSet = fieldNumberToAtomDeclSetIt->second;
+    for (const shared_ptr<AtomDecl>& atomDecl : atomDeclSet) {
+        fprintf(out, "    if (code == %d) {\n", atomDecl->code);
+        const AnnotationSet& annotations = atomDecl->fieldNumberToAnnotations.at(argIndex);
+        int resetState = -1;
+        int defaultState = -1;
+        for (const shared_ptr<Annotation>& annotation : annotations) {
+            // TODO(b/151786433): Write atom constant name instead of atom id literal.
+            switch (annotation->type) {
+                // TODO(b/151776731): Check for reset state annotation and only include
+                // reset state when field value == default state annotation value.
+                case ANNOTATION_TYPE_INT:
+                    // TODO(b/151786433): Write annotation constant name instead of
+                    // annotation id literal.
+                    if (ANNOTATION_ID_RESET_STATE == annotation->annotationId) {
+                        resetState = annotation->value.intValue;
+                    } else if (ANNOTATION_ID_DEFAULT_STATE == annotation->annotationId) {
+                        defaultState = annotation->value.intValue;
+                    } else {
+                        fprintf(out, "        %saddInt32Annotation(%s%d, %d);\n",
+                                methodPrefix.c_str(), methodSuffix.c_str(),
+                                annotation->annotationId, annotation->value.intValue);
+                    }
+                    break;
+                case ANNOTATION_TYPE_BOOL:
+                    // TODO(b/151786433): Write annotation constant name instead of
+                    // annotation id literal.
+                    fprintf(out, "        %saddBoolAnnotation(%s%d, %s);\n", methodPrefix.c_str(),
+                            methodSuffix.c_str(), annotation->annotationId,
+                            annotation->value.boolValue ? "true" : "false");
+                    break;
+                default:
+                    break;
+            }
+        }
+        if (defaultState != -1 && resetState != -1) {
+            fprintf(out, "        if (arg%d == %d) {\n", argIndex, resetState);
+            fprintf(out, "            %saddInt32Annotation(%s%d, %d);\n", methodPrefix.c_str(),
+                    methodSuffix.c_str(), ANNOTATION_ID_RESET_STATE, defaultState);
+            fprintf(out, "        }\n");
         }
         fprintf(out, "    }\n");
     }
@@ -62,7 +80,7 @@
     for (auto signatureInfoMapIt = atoms.signatureInfoMap.begin();
          signatureInfoMapIt != atoms.signatureInfoMap.end(); signatureInfoMapIt++) {
         vector<java_type_t> signature = signatureInfoMapIt->first;
-        const FieldNumberToAnnotations& fieldNumberToAnnotations = signatureInfoMapIt->second;
+        const FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet = signatureInfoMapIt->second;
         // Key value pairs not supported in native.
         if (find(signature.begin(), signature.end(), JAVA_TYPE_KEY_VALUE_PAIR) != signature.end()) {
             continue;
@@ -73,7 +91,7 @@
         if (supportQ) {
             fprintf(out, "    StatsEventCompat event;\n");
             fprintf(out, "    event.setAtomId(code);\n");
-            write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAnnotations, "event.", "");
+            write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAtomDeclSet, "event.", "");
             for (vector<java_type_t>::const_iterator arg = signature.begin();
                  arg != signature.end(); arg++) {
                 switch (*arg) {
@@ -109,14 +127,14 @@
                         fprintf(stderr, "Encountered unsupported type.");
                         return 1;
                 }
-                write_annotations(out, argIndex, fieldNumberToAnnotations, "event.", "");
+                write_annotations(out, argIndex, fieldNumberToAtomDeclSet, "event.", "");
                 argIndex++;
             }
             fprintf(out, "    return event.writeToSocket();\n");
         } else {
             fprintf(out, "    AStatsEvent* event = AStatsEvent_obtain();\n");
             fprintf(out, "    AStatsEvent_setAtomId(event, code);\n");
-            write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAnnotations, "AStatsEvent_",
+            write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAtomDeclSet, "AStatsEvent_",
                               "event, ");
             for (vector<java_type_t>::const_iterator arg = signature.begin();
                  arg != signature.end(); arg++) {
@@ -159,7 +177,7 @@
                         fprintf(stderr, "Encountered unsupported type.");
                         return 1;
                 }
-                write_annotations(out, argIndex, fieldNumberToAnnotations, "AStatsEvent_",
+                write_annotations(out, argIndex, fieldNumberToAtomDeclSet, "AStatsEvent_",
                                   "event, ");
                 argIndex++;
             }
@@ -209,10 +227,9 @@
     }
 }
 
-static void write_native_method_header(
-        FILE* out, const string& methodName,
-        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
-        const AtomDecl& attributionDecl) {
+static void write_native_method_header(FILE* out, const string& methodName,
+                                       const SignatureInfoMap& signatureInfoMap,
+                                       const AtomDecl& attributionDecl) {
     for (auto signatureInfoMapIt = signatureInfoMap.begin();
          signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
         vector<java_type_t> signature = signatureInfoMapIt->first;
@@ -278,16 +295,17 @@
     fprintf(out, "//\n");
     fprintf(out, "// Constants for enum values\n");
     fprintf(out, "//\n\n");
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        for (vector<AtomField>::const_iterator field = atom->fields.begin();
-             field != atom->fields.end(); field++) {
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        for (vector<AtomField>::const_iterator field = (*atomIt)->fields.begin();
+             field != (*atomIt)->fields.end(); field++) {
             if (field->javaType == JAVA_TYPE_ENUM) {
-                fprintf(out, "// Values for %s.%s\n", atom->message.c_str(), field->name.c_str());
+                fprintf(out, "// Values for %s.%s\n", (*atomIt)->message.c_str(),
+                        field->name.c_str());
                 for (map<int, string>::const_iterator value = field->enumValues.begin();
                      value != field->enumValues.end(); value++) {
                     fprintf(out, "const int32_t %s__%s__%s = %d;\n",
-                            make_constant_name(atom->message).c_str(),
+                            make_constant_name((*atomIt)->message).c_str(),
                             make_constant_name(field->name).c_str(),
                             make_constant_name(value->second).c_str(), value->first);
                 }
diff --git a/tools/stats_log_api_gen/test_collation.cpp b/tools/stats_log_api_gen/test_collation.cpp
index 9878926..c654a1b 100644
--- a/tools/stats_log_api_gen/test_collation.cpp
+++ b/tools/stats_log_api_gen/test_collation.cpp
@@ -30,8 +30,7 @@
 /**
  * Return whether the map contains a vector of the elements provided.
  */
-static bool map_contains_vector(const map<vector<java_type_t>, FieldNumberToAnnotations>& s,
-                                int count, ...) {
+static bool map_contains_vector(const SignatureInfoMap& s, int count, ...) {
     va_list args;
     vector<java_type_t> v;
 
@@ -110,35 +109,37 @@
                                   JAVA_TYPE_LONG                // sint64
     );
 
-    set<AtomDecl>::const_iterator atom = atoms.decls.begin();
-    EXPECT_EQ(1, atom->code);
-    EXPECT_EQ("int_atom", atom->name);
-    EXPECT_EQ("IntAtom", atom->message);
-    EXPECT_NO_ENUM_FIELD(atom);
-    atom++;
+    EXPECT_EQ(4ul, atoms.decls.size());
 
-    EXPECT_EQ(2, atom->code);
-    EXPECT_EQ("out_of_order_atom", atom->name);
-    EXPECT_EQ("OutOfOrderAtom", atom->message);
-    EXPECT_NO_ENUM_FIELD(atom);
-    atom++;
+    AtomDeclSet::const_iterator atomIt = atoms.decls.begin();
+    EXPECT_EQ(1, (*atomIt)->code);
+    EXPECT_EQ("int_atom", (*atomIt)->name);
+    EXPECT_EQ("IntAtom", (*atomIt)->message);
+    EXPECT_NO_ENUM_FIELD((*atomIt));
+    atomIt++;
 
-    EXPECT_EQ(3, atom->code);
-    EXPECT_EQ("another_int_atom", atom->name);
-    EXPECT_EQ("AnotherIntAtom", atom->message);
-    EXPECT_NO_ENUM_FIELD(atom);
-    atom++;
+    EXPECT_EQ(2, (*atomIt)->code);
+    EXPECT_EQ("out_of_order_atom", (*atomIt)->name);
+    EXPECT_EQ("OutOfOrderAtom", (*atomIt)->message);
+    EXPECT_NO_ENUM_FIELD((*atomIt));
+    atomIt++;
 
-    EXPECT_EQ(4, atom->code);
-    EXPECT_EQ("all_types_atom", atom->name);
-    EXPECT_EQ("AllTypesAtom", atom->message);
+    EXPECT_EQ(3, (*atomIt)->code);
+    EXPECT_EQ("another_int_atom", (*atomIt)->name);
+    EXPECT_EQ("AnotherIntAtom", (*atomIt)->message);
+    EXPECT_NO_ENUM_FIELD((*atomIt));
+    atomIt++;
+
+    EXPECT_EQ(4, (*atomIt)->code);
+    EXPECT_EQ("all_types_atom", (*atomIt)->name);
+    EXPECT_EQ("AllTypesAtom", (*atomIt)->message);
     map<int, string> enumValues;
     enumValues[0] = "VALUE0";
     enumValues[1] = "VALUE1";
-    EXPECT_HAS_ENUM_FIELD(atom, "enum_field", enumValues);
-    atom++;
+    EXPECT_HAS_ENUM_FIELD((*atomIt), "enum_field", enumValues);
+    atomIt++;
 
-    EXPECT_TRUE(atom == atoms.decls.end());
+    EXPECT_EQ(atoms.decls.end(), atomIt);
 }
 
 /**
@@ -235,10 +236,10 @@
     Atoms atoms;
     collate_atoms(ListedAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     for (const auto& atomDecl : atoms.decls) {
-        if (atomDecl.code == 1) {
-            EXPECT_TRUE(atomDecl.whitelisted);
+        if (atomDecl->code == 1) {
+            EXPECT_TRUE(atomDecl->whitelisted);
         } else {
-            EXPECT_FALSE(atomDecl.whitelisted);
+            EXPECT_FALSE(atomDecl->whitelisted);
         }
     }
 }
@@ -258,33 +259,66 @@
     EXPECT_EQ(atoms.signatureInfoMap.size(), 2u);
     EXPECT_MAP_CONTAINS_SIGNATURE(atoms.signatureInfoMap, JAVA_TYPE_INT);
     EXPECT_MAP_CONTAINS_SIGNATURE(atoms.signatureInfoMap, JAVA_TYPE_STRING);
-    for (auto signatureInfoMapIt : atoms.signatureInfoMap) {
-        vector<java_type_t> signature = signatureInfoMapIt.first;
-        const FieldNumberToAnnotations& fieldNumberToAnnotations = signatureInfoMapIt.second;
-        if (signature[0] == JAVA_TYPE_STRING) {
-            EXPECT_EQ(0u, fieldNumberToAnnotations.size());
-        } else if (signature[0] == JAVA_TYPE_INT) {
-            EXPECT_EQ(1u, fieldNumberToAnnotations.size());
-            EXPECT_NE(fieldNumberToAnnotations.end(), fieldNumberToAnnotations.find(1));
-            const set<shared_ptr<Annotation>>& annotations = fieldNumberToAnnotations.at(1);
-            EXPECT_EQ(2u, annotations.size());
-            for (const shared_ptr<Annotation> annotation : annotations) {
-                EXPECT_TRUE(annotation->annotationId == ANNOTATION_ID_IS_UID ||
-                            annotation->annotationId == ANNOTATION_ID_STATE_OPTION);
-                if (ANNOTATION_ID_IS_UID == annotation->annotationId) {
-                    EXPECT_EQ(1, annotation->atomId);
-                    EXPECT_EQ(ANNOTATION_TYPE_BOOL, annotation->type);
-                    EXPECT_TRUE(annotation->value.boolValue);
-                }
 
-                if (ANNOTATION_ID_STATE_OPTION == annotation->annotationId) {
-                    EXPECT_EQ(3, annotation->atomId);
-                    EXPECT_EQ(ANNOTATION_TYPE_INT, annotation->type);
-                    EXPECT_EQ(os::statsd::StateField::EXCLUSIVE_STATE, annotation->value.intValue);
-                }
-            }
-        }
-    }
+    SignatureInfoMap::const_iterator signatureInfoMapIt;
+    const vector<java_type_t>* signature;
+    const FieldNumberToAtomDeclSet* fieldNumberToAtomDeclSet;
+    FieldNumberToAtomDeclSet::const_iterator fieldNumberToAtomDeclSetIt;
+    const AtomDeclSet* atomDeclSet;
+    AtomDeclSet::const_iterator atomDeclSetIt;
+    AtomDecl* atomDecl;
+    FieldNumberToAnnotations* fieldNumberToAnnotations;
+    FieldNumberToAnnotations::const_iterator fieldNumberToAnnotationsIt;
+    const AnnotationSet* annotationSet;
+    AnnotationSet::const_iterator annotationSetIt;
+    Annotation* annotation;
+
+    signatureInfoMapIt = atoms.signatureInfoMap.begin();
+    signature = &(signatureInfoMapIt->first);
+    fieldNumberToAtomDeclSet = &signatureInfoMapIt->second;
+    EXPECT_EQ(1ul, signature->size());
+    EXPECT_EQ(JAVA_TYPE_INT, signature->at(0));
+    EXPECT_EQ(1ul, fieldNumberToAtomDeclSet->size());
+    fieldNumberToAtomDeclSetIt = fieldNumberToAtomDeclSet->begin();
+    EXPECT_EQ(1, fieldNumberToAtomDeclSetIt->first);
+    atomDeclSet = &fieldNumberToAtomDeclSetIt->second;
+    EXPECT_EQ(2ul, atomDeclSet->size());
+    atomDeclSetIt = atomDeclSet->begin();
+    atomDecl = atomDeclSetIt->get();
+    EXPECT_EQ(1, atomDecl->code);
+    fieldNumberToAnnotations = &atomDecl->fieldNumberToAnnotations;
+    fieldNumberToAnnotationsIt = fieldNumberToAnnotations->find(1);
+    EXPECT_NE(fieldNumberToAnnotations->end(), fieldNumberToAnnotationsIt);
+    annotationSet = &fieldNumberToAnnotationsIt->second;
+    EXPECT_EQ(1ul, annotationSet->size());
+    annotationSetIt = annotationSet->begin();
+    annotation = annotationSetIt->get();
+    EXPECT_EQ(ANNOTATION_ID_IS_UID, annotation->annotationId);
+    EXPECT_EQ(1, annotation->atomId);
+    EXPECT_EQ(ANNOTATION_TYPE_BOOL, annotation->type);
+    EXPECT_TRUE(annotation->value.boolValue);
+
+    atomDeclSetIt++;
+    atomDecl = atomDeclSetIt->get();
+    EXPECT_EQ(3, atomDecl->code);
+    fieldNumberToAnnotations = &atomDecl->fieldNumberToAnnotations;
+    fieldNumberToAnnotationsIt = fieldNumberToAnnotations->find(1);
+    EXPECT_NE(fieldNumberToAnnotations->end(), fieldNumberToAnnotationsIt);
+    annotationSet = &fieldNumberToAnnotationsIt->second;
+    EXPECT_EQ(1ul, annotationSet->size());
+    annotationSetIt = annotationSet->begin();
+    annotation = annotationSetIt->get();
+    EXPECT_EQ(ANNOTATION_ID_STATE_OPTION, annotation->annotationId);
+    EXPECT_EQ(3, annotation->atomId);
+    EXPECT_EQ(ANNOTATION_TYPE_INT, annotation->type);
+    EXPECT_EQ(os::statsd::StateField::EXCLUSIVE_STATE, annotation->value.intValue);
+
+    signatureInfoMapIt++;
+    signature = &signatureInfoMapIt->first;
+    fieldNumberToAtomDeclSet = &signatureInfoMapIt->second;
+    EXPECT_EQ(1ul, signature->size());
+    EXPECT_EQ(JAVA_TYPE_STRING, signature->at(0));
+    EXPECT_EQ(0ul, fieldNumberToAtomDeclSet->size());
 }
 
 TEST(CollationTest, RecognizeModule1Atom) {
@@ -295,31 +329,59 @@
     EXPECT_EQ(atoms.decls.size(), 2ul);
     EXPECT_EQ(atoms.signatureInfoMap.size(), 1u);
     EXPECT_MAP_CONTAINS_SIGNATURE(atoms.signatureInfoMap, JAVA_TYPE_INT);
-    for (auto signatureInfoMapIt : atoms.signatureInfoMap) {
-        vector<java_type_t> signature = signatureInfoMapIt.first;
-        const FieldNumberToAnnotations& fieldNumberToAnnotations = signatureInfoMapIt.second;
-        EXPECT_EQ(JAVA_TYPE_INT, signature[0]);
-        EXPECT_EQ(1u, fieldNumberToAnnotations.size());
-        int fieldNumber = 1;
-        EXPECT_NE(fieldNumberToAnnotations.end(), fieldNumberToAnnotations.find(fieldNumber));
-        const set<shared_ptr<Annotation>>& annotations = fieldNumberToAnnotations.at(fieldNumber);
-        EXPECT_EQ(2u, annotations.size());
-        for (const shared_ptr<Annotation> annotation : annotations) {
-            EXPECT_TRUE(annotation->annotationId == ANNOTATION_ID_IS_UID ||
-                        annotation->annotationId == ANNOTATION_ID_STATE_OPTION);
-            if (ANNOTATION_ID_IS_UID == annotation->annotationId) {
-                EXPECT_EQ(1, annotation->atomId);
-                EXPECT_EQ(ANNOTATION_TYPE_BOOL, annotation->type);
-                EXPECT_TRUE(annotation->value.boolValue);
-            }
 
-            if (ANNOTATION_ID_STATE_OPTION == annotation->annotationId) {
-                EXPECT_EQ(3, annotation->atomId);
-                EXPECT_EQ(ANNOTATION_TYPE_INT, annotation->type);
-                EXPECT_EQ(os::statsd::StateField::EXCLUSIVE_STATE, annotation->value.intValue);
-            }
-        }
-    }
+    SignatureInfoMap::const_iterator signatureInfoMapIt;
+    const vector<java_type_t>* signature;
+    const FieldNumberToAtomDeclSet* fieldNumberToAtomDeclSet;
+    FieldNumberToAtomDeclSet::const_iterator fieldNumberToAtomDeclSetIt;
+    const AtomDeclSet* atomDeclSet;
+    AtomDeclSet::const_iterator atomDeclSetIt;
+    AtomDecl* atomDecl;
+    FieldNumberToAnnotations* fieldNumberToAnnotations;
+    FieldNumberToAnnotations::const_iterator fieldNumberToAnnotationsIt;
+    const AnnotationSet* annotationSet;
+    AnnotationSet::const_iterator annotationSetIt;
+    Annotation* annotation;
+
+    signatureInfoMapIt = atoms.signatureInfoMap.begin();
+    signature = &(signatureInfoMapIt->first);
+    fieldNumberToAtomDeclSet = &signatureInfoMapIt->second;
+    EXPECT_EQ(1ul, signature->size());
+    EXPECT_EQ(JAVA_TYPE_INT, signature->at(0));
+    EXPECT_EQ(1ul, fieldNumberToAtomDeclSet->size());
+    fieldNumberToAtomDeclSetIt = fieldNumberToAtomDeclSet->begin();
+    EXPECT_EQ(1, fieldNumberToAtomDeclSetIt->first);
+    atomDeclSet = &fieldNumberToAtomDeclSetIt->second;
+    EXPECT_EQ(2ul, atomDeclSet->size());
+    atomDeclSetIt = atomDeclSet->begin();
+    atomDecl = atomDeclSetIt->get();
+    EXPECT_EQ(1, atomDecl->code);
+    fieldNumberToAnnotations = &atomDecl->fieldNumberToAnnotations;
+    fieldNumberToAnnotationsIt = fieldNumberToAnnotations->find(1);
+    EXPECT_NE(fieldNumberToAnnotations->end(), fieldNumberToAnnotationsIt);
+    annotationSet = &fieldNumberToAnnotationsIt->second;
+    EXPECT_EQ(1ul, annotationSet->size());
+    annotationSetIt = annotationSet->begin();
+    annotation = annotationSetIt->get();
+    EXPECT_EQ(ANNOTATION_ID_IS_UID, annotation->annotationId);
+    EXPECT_EQ(1, annotation->atomId);
+    EXPECT_EQ(ANNOTATION_TYPE_BOOL, annotation->type);
+    EXPECT_TRUE(annotation->value.boolValue);
+
+    atomDeclSetIt++;
+    atomDecl = atomDeclSetIt->get();
+    EXPECT_EQ(3, atomDecl->code);
+    fieldNumberToAnnotations = &atomDecl->fieldNumberToAnnotations;
+    fieldNumberToAnnotationsIt = fieldNumberToAnnotations->find(1);
+    EXPECT_NE(fieldNumberToAnnotations->end(), fieldNumberToAnnotationsIt);
+    annotationSet = &fieldNumberToAnnotationsIt->second;
+    EXPECT_EQ(1ul, annotationSet->size());
+    annotationSetIt = annotationSet->begin();
+    annotation = annotationSetIt->get();
+    EXPECT_EQ(ANNOTATION_ID_STATE_OPTION, annotation->annotationId);
+    EXPECT_EQ(3, annotation->atomId);
+    EXPECT_EQ(ANNOTATION_TYPE_INT, annotation->type);
+    EXPECT_EQ(os::statsd::StateField::EXCLUSIVE_STATE, annotation->value.intValue);
 }
 
 }  // namespace stats_log_api_gen
diff --git a/tools/stats_log_api_gen/utils.cpp b/tools/stats_log_api_gen/utils.cpp
index 0262488..abb8913 100644
--- a/tools/stats_log_api_gen/utils.cpp
+++ b/tools/stats_log_api_gen/utils.cpp
@@ -22,10 +22,10 @@
 namespace stats_log_api_gen {
 
 static void build_non_chained_decl_map(const Atoms& atoms,
-                                       std::map<int, set<AtomDecl>::const_iterator>* decl_map) {
-    for (set<AtomDecl>::const_iterator atom = atoms.non_chained_decls.begin();
-         atom != atoms.non_chained_decls.end(); atom++) {
-        decl_map->insert(std::make_pair(atom->code, atom));
+                                       std::map<int, AtomDeclSet::const_iterator>* decl_map) {
+    for (AtomDeclSet::const_iterator atomIt = atoms.non_chained_decls.begin();
+         atomIt != atoms.non_chained_decls.end(); atomIt++) {
+        decl_map->insert(std::make_pair((*atomIt)->code, atomIt));
     }
 }
 
@@ -117,11 +117,11 @@
 }
 
 static void write_cpp_usage(FILE* out, const string& method_name, const string& atom_code_name,
-                            const AtomDecl& atom, const AtomDecl& attributionDecl) {
+                            const shared_ptr<AtomDecl> atom, const AtomDecl& attributionDecl) {
     fprintf(out, "     * Usage: %s(StatsLog.%s", method_name.c_str(), atom_code_name.c_str());
 
-    for (vector<AtomField>::const_iterator field = atom.fields.begin(); field != atom.fields.end();
-         field++) {
+    for (vector<AtomField>::const_iterator field = atom->fields.begin();
+         field != atom->fields.end(); field++) {
         if (field->javaType == JAVA_TYPE_ATTRIBUTION_CHAIN) {
             for (auto chainField : attributionDecl.fields) {
                 if (chainField.javaType == JAVA_TYPE_STRING) {
@@ -154,27 +154,27 @@
     fprintf(out, " */\n");
     fprintf(out, "enum {\n");
 
-    std::map<int, set<AtomDecl>::const_iterator> atom_code_to_non_chained_decl_map;
+    std::map<int, AtomDeclSet::const_iterator> atom_code_to_non_chained_decl_map;
     build_non_chained_decl_map(atoms, &atom_code_to_non_chained_decl_map);
 
     size_t i = 0;
     // Print atom constants
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        string constant = make_constant_name(atom->name);
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        string constant = make_constant_name((*atomIt)->name);
         fprintf(out, "\n");
         fprintf(out, "    /**\n");
-        fprintf(out, "     * %s %s\n", atom->message.c_str(), atom->name.c_str());
-        write_cpp_usage(out, "stats_write", constant, *atom, attributionDecl);
+        fprintf(out, "     * %s %s\n", (*atomIt)->message.c_str(), (*atomIt)->name.c_str());
+        write_cpp_usage(out, "stats_write", constant, *atomIt, attributionDecl);
 
-        auto non_chained_decl = atom_code_to_non_chained_decl_map.find(atom->code);
+        auto non_chained_decl = atom_code_to_non_chained_decl_map.find((*atomIt)->code);
         if (non_chained_decl != atom_code_to_non_chained_decl_map.end()) {
             write_cpp_usage(out, "stats_write_non_chained", constant, *non_chained_decl->second,
                             attributionDecl);
         }
         fprintf(out, "     */\n");
         char const* const comma = (i == atoms.decls.size() - 1) ? "" : ",";
-        fprintf(out, "    %s = %d%s\n", constant.c_str(), atom->code, comma);
+        fprintf(out, "    %s = %d%s\n", constant.c_str(), (*atomIt)->code, comma);
         i++;
     }
     fprintf(out, "\n");
@@ -245,40 +245,40 @@
 void write_java_atom_codes(FILE* out, const Atoms& atoms) {
     fprintf(out, "    // Constants for atom codes.\n");
 
-    std::map<int, set<AtomDecl>::const_iterator> atom_code_to_non_chained_decl_map;
+    std::map<int, AtomDeclSet::const_iterator> atom_code_to_non_chained_decl_map;
     build_non_chained_decl_map(atoms, &atom_code_to_non_chained_decl_map);
 
     // Print constants for the atom codes.
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        string constant = make_constant_name(atom->name);
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        string constant = make_constant_name((*atomIt)->name);
         fprintf(out, "\n");
         fprintf(out, "    /**\n");
-        fprintf(out, "     * %s %s<br>\n", atom->message.c_str(), atom->name.c_str());
-        write_java_usage(out, "write", constant, *atom);
-        auto non_chained_decl = atom_code_to_non_chained_decl_map.find(atom->code);
+        fprintf(out, "     * %s %s<br>\n", (*atomIt)->message.c_str(), (*atomIt)->name.c_str());
+        write_java_usage(out, "write", constant, **atomIt);
+        auto non_chained_decl = atom_code_to_non_chained_decl_map.find((*atomIt)->code);
         if (non_chained_decl != atom_code_to_non_chained_decl_map.end()) {
-            write_java_usage(out, "write_non_chained", constant, *non_chained_decl->second);
+            write_java_usage(out, "write_non_chained", constant, **(non_chained_decl->second));
         }
         fprintf(out, "     */\n");
-        fprintf(out, "    public static final int %s = %d;\n", constant.c_str(), atom->code);
+        fprintf(out, "    public static final int %s = %d;\n", constant.c_str(), (*atomIt)->code);
     }
     fprintf(out, "\n");
 }
 
 void write_java_enum_values(FILE* out, const Atoms& atoms) {
     fprintf(out, "    // Constants for enum values.\n\n");
-    for (set<AtomDecl>::const_iterator atom = atoms.decls.begin(); atom != atoms.decls.end();
-         atom++) {
-        for (vector<AtomField>::const_iterator field = atom->fields.begin();
-             field != atom->fields.end(); field++) {
+    for (AtomDeclSet::const_iterator atomIt = atoms.decls.begin(); atomIt != atoms.decls.end();
+         atomIt++) {
+        for (vector<AtomField>::const_iterator field = (*atomIt)->fields.begin();
+             field != (*atomIt)->fields.end(); field++) {
             if (field->javaType == JAVA_TYPE_ENUM) {
-                fprintf(out, "    // Values for %s.%s\n", atom->message.c_str(),
+                fprintf(out, "    // Values for %s.%s\n", (*atomIt)->message.c_str(),
                         field->name.c_str());
                 for (map<int, string>::const_iterator value = field->enumValues.begin();
                      value != field->enumValues.end(); value++) {
                     fprintf(out, "    public static final int %s__%s__%s = %d;\n",
-                            make_constant_name(atom->message).c_str(),
+                            make_constant_name((*atomIt)->message).c_str(),
                             make_constant_name(field->name).c_str(),
                             make_constant_name(value->second).c_str(), value->first);
                 }
@@ -307,8 +307,7 @@
     fprintf(out, ");<br>\n");
 }
 
-int write_java_non_chained_methods(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap) {
+int write_java_non_chained_methods(FILE* out, const SignatureInfoMap& signatureInfoMap) {
     for (auto signatureInfoMapIt = signatureInfoMap.begin();
          signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
         // Print method signature.
@@ -351,8 +350,7 @@
     return 0;
 }
 
-int write_java_work_source_methods(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap) {
+int write_java_work_source_methods(FILE* out, const SignatureInfoMap& signatureInfoMap) {
     fprintf(out, "    // WorkSource methods.\n");
     for (auto signatureInfoMapIt = signatureInfoMap.begin();
          signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
diff --git a/tools/stats_log_api_gen/utils.h b/tools/stats_log_api_gen/utils.h
index 468f323..57b6f62 100644
--- a/tools/stats_log_api_gen/utils.h
+++ b/tools/stats_log_api_gen/utils.h
@@ -67,11 +67,9 @@
 void write_java_usage(FILE* out, const string& method_name, const string& atom_code_name,
                       const AtomDecl& atom);
 
-int write_java_non_chained_methods(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap);
+int write_java_non_chained_methods(FILE* out, const SignatureInfoMap& signatureInfoMap);
 
-int write_java_work_source_methods(
-        FILE* out, const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap);
+int write_java_work_source_methods(FILE* out, const SignatureInfoMap& signatureInfoMap);
 
 }  // namespace stats_log_api_gen
 }  // namespace android