diff --git a/tools/stats_log_api_gen/Collation.cpp b/tools/stats_log_api_gen/Collation.cpp
index 66c4b1488b83e..f31a2af87b2bc 100644
--- a/tools/stats_log_api_gen/Collation.cpp
+++ b/tools/stats_log_api_gen/Collation.cpp
@@ -27,8 +27,11 @@ using google::protobuf::EnumDescriptor;
 using google::protobuf::FieldDescriptor;
 using google::protobuf::FileDescriptor;
 using google::protobuf::SourceLocation;
+using std::make_shared;
 using std::map;
 
+const bool dbg = false;
+
 
 //
 // AtomDecl class
@@ -45,6 +48,7 @@ AtomDecl::AtomDecl(const AtomDecl &that)
         name(that.name),
         message(that.message),
         fields(that.fields),
+        fieldNumberToAnnotations(that.fieldNumberToAnnotations),
         primaryFields(that.primaryFields),
         exclusiveField(that.exclusiveField),
         defaultState(that.defaultState),
@@ -52,8 +56,7 @@ AtomDecl::AtomDecl(const AtomDecl &that)
         nested(that.nested),
         uidField(that.uidField),
         whitelisted(that.whitelisted),
-        binaryFields(that.binaryFields),
-        moduleNames(that.moduleNames) {}
+        binaryFields(that.binaryFields) {}
 
 AtomDecl::AtomDecl(int c, const string& n, const string& m)
     :code(c),
@@ -160,6 +163,17 @@ void collate_enums(const EnumDescriptor &enumDescriptor, AtomField *atomField) {
     }
 }
 
+static void addAnnotationToAtomDecl(AtomDecl* atomDecl, const int fieldNumber,
+        const int annotationId, const AnnotationType annotationType,
+        const AnnotationValue annotationValue) {
+    if (dbg) {
+        printf("   Adding annotation to %s: [%d] = {id: %d, type: %d}\n",
+                atomDecl->name.c_str(), fieldNumber, annotationId, annotationType);
+    }
+    atomDecl->fieldNumberToAnnotations[fieldNumber].insert(make_shared<Annotation>(
+                annotationId, atomDecl->code, annotationType, annotationValue));
+}
+
 /**
  * Gather the info about an atom proto.
  */
@@ -279,7 +293,6 @@ int collate_atom(const Descriptor *atom, AtomDecl *atomDecl,
       if (javaType == JAVA_TYPE_ENUM) {
         // All enums are treated as ints when it comes to function signatures.
         signature->push_back(JAVA_TYPE_INT);
-        collate_enums(*field->enum_type(), &atField);
       } else if (javaType == JAVA_TYPE_OBJECT && isBinaryField) {
           signature->push_back(JAVA_TYPE_BYTE_ARRAY);
       } else {
@@ -292,64 +305,121 @@ int collate_atom(const Descriptor *atom, AtomDecl *atomDecl,
     }
     atomDecl->fields.push_back(atField);
 
-    if (field->options().GetExtension(os::statsd::state_field_option).option() ==
-        os::statsd::StateField::PRIMARY_FIELD) {
-        if (javaType == JAVA_TYPE_UNKNOWN ||
-            javaType == JAVA_TYPE_ATTRIBUTION_CHAIN ||
-            javaType == JAVA_TYPE_OBJECT || javaType == JAVA_TYPE_BYTE_ARRAY) {
-            errorCount++;
+    if (field->options().HasExtension(os::statsd::state_field_option)) {
+        const int option = field->options().GetExtension(os::statsd::state_field_option).option();
+        if (option != STATE_OPTION_UNSET) {
+            addAnnotationToAtomDecl(atomDecl, signature->size(), ANNOTATION_ID_STATE_OPTION,
+                    ANNOTATION_TYPE_INT, AnnotationValue(option));
         }
-        atomDecl->primaryFields.push_back(it->first);
+
+        if (option == STATE_OPTION_PRIMARY) {
+            if (javaType == JAVA_TYPE_UNKNOWN ||
+                    javaType == JAVA_TYPE_ATTRIBUTION_CHAIN ||
+                    javaType == JAVA_TYPE_OBJECT || javaType == JAVA_TYPE_BYTE_ARRAY) {
+                print_error(
+                    field,
+                    "Invalid primary state field: '%s'\n",
+                    atom->name().c_str());
+                errorCount++;
+                continue;
+            }
+            atomDecl->primaryFields.push_back(it->first);
+
+        }
+
+        if (option == STATE_OPTION_PRIMARY_FIELD_FIRST_UID) {
+            if (javaType != JAVA_TYPE_ATTRIBUTION_CHAIN) {
+                print_error(
+                    field,
+                    "PRIMARY_FIELD_FIRST_UID annotation is only for AttributionChains: '%s'\n",
+                    atom->name().c_str());
+                errorCount++;
+                continue;
+            } else {
+                atomDecl->primaryFields.push_back(FIRST_UID_IN_CHAIN_ID);
+            }
+        }
+
+        if (option == STATE_OPTION_EXCLUSIVE) {
+            if (javaType == JAVA_TYPE_UNKNOWN ||
+                    javaType == JAVA_TYPE_ATTRIBUTION_CHAIN ||
+                    javaType == JAVA_TYPE_OBJECT || javaType == JAVA_TYPE_BYTE_ARRAY) {
+                print_error(
+                    field,
+                    "Invalid exclusive state field: '%s'\n",
+                    atom->name().c_str());
+                errorCount++;
+                continue;
+            }
+
+            if (atomDecl->exclusiveField == 0) {
+                atomDecl->exclusiveField = it->first;
+            } else {
+                print_error(
+                    field,
+                    "Cannot have more than one exclusive state field in an atom: '%s'\n",
+                    atom->name().c_str());
+                errorCount++;
+                continue;
+            }
+
+            if (field->options()
+                        .GetExtension(os::statsd::state_field_option)
+                        .has_default_state_value()) {
+                const int defaultState =
+                        field->options().GetExtension(os::statsd::state_field_option)
+                        .default_state_value();
+                atomDecl->defaultState = defaultState;
+
+                addAnnotationToAtomDecl(atomDecl, signature->size(), ANNOTATION_ID_DEFAULT_STATE,
+                        ANNOTATION_TYPE_INT, AnnotationValue(defaultState));
+            }
+
+            if (field->options().GetExtension(os::statsd::state_field_option)
+                    .has_reset_state_value()) {
+                const int resetState = field->options()
+                        .GetExtension(os::statsd::state_field_option)
+                        .reset_state_value();
+
+                atomDecl->resetState = resetState;
+                addAnnotationToAtomDecl(atomDecl, signature->size(), ANNOTATION_ID_RESET_STATE,
+                        ANNOTATION_TYPE_INT, AnnotationValue(resetState));
+            }
+
+            if (field->options().GetExtension(os::statsd::state_field_option)
+                    .has_nested()) {
+                const bool nested =
+                        field->options().GetExtension(os::statsd::state_field_option).nested();
+                atomDecl->nested = nested;
+
+                addAnnotationToAtomDecl(atomDecl, signature->size(), ANNOTATION_ID_STATE_NESTED,
+                        ANNOTATION_TYPE_BOOL, AnnotationValue(nested));
+            }
+        }
+
     }
-
-    if (field->options().GetExtension(os::statsd::state_field_option).option() ==
-        os::statsd::StateField::PRIMARY_FIELD_FIRST_UID) {
-        if (javaType != JAVA_TYPE_ATTRIBUTION_CHAIN) {
-            errorCount++;
-        } else {
-            atomDecl->primaryFields.push_back(FIRST_UID_IN_CHAIN_ID);
-        }
-    }
-
-    if (field->options().GetExtension(os::statsd::state_field_option).option() ==
-        os::statsd::StateField::EXCLUSIVE_STATE) {
-        if (javaType == JAVA_TYPE_UNKNOWN ||
-            javaType == JAVA_TYPE_ATTRIBUTION_CHAIN ||
-            javaType == JAVA_TYPE_OBJECT || javaType == JAVA_TYPE_BYTE_ARRAY) {
-            errorCount++;
-        }
-
-        if (atomDecl->exclusiveField == 0) {
-            atomDecl->exclusiveField = it->first;
-        } else {
-            errorCount++;
-        }
-
-        if (field->options()
-                    .GetExtension(os::statsd::state_field_option)
-                    .has_default_state_value()) {
-            atomDecl->defaultState = field->options()
-                                             .GetExtension(os::statsd::state_field_option)
-                                             .default_state_value();
-        }
-
-        if (field->options().GetExtension(os::statsd::state_field_option).has_reset_state_value()) {
-            atomDecl->resetState = field->options()
-                                           .GetExtension(os::statsd::state_field_option)
-                                           .reset_state_value();
-        }
-        atomDecl->nested = field->options().GetExtension(os::statsd::state_field_option).nested();
-    }
-
     if (field->options().GetExtension(os::statsd::is_uid) == true) {
         if (javaType != JAVA_TYPE_INT) {
+            print_error(
+                field,
+                "is_uid annotation can only be applied to int32 fields: '%s'\n",
+                atom->name().c_str());
             errorCount++;
+            continue;
         }
 
         if (atomDecl->uidField == 0) {
             atomDecl->uidField = it->first;
+
+            addAnnotationToAtomDecl(atomDecl, signature->size(), ANNOTATION_ID_IS_UID,
+                    ANNOTATION_TYPE_BOOL, AnnotationValue(true));
         } else {
+            print_error(
+                field,
+                "Cannot have more than one field in an atom with is_uid annotation: '%s'\n",
+                atom->name().c_str());
             errorCount++;
+            continue;
         }
     }
     // Binary field validity is already checked above.
@@ -408,17 +478,50 @@ bool get_non_chained_node(const Descriptor *atom, AtomDecl *atomDecl,
     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++) {
+        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());
+    }
+}
+
 /**
  * Gather the info about the atoms.
  */
-int collate_atoms(const Descriptor *descriptor, Atoms *atoms) {
+int collate_atoms(const Descriptor *descriptor, const string& moduleName, Atoms *atoms) {
   int errorCount = 0;
-  const bool dbg = false;
 
   int maxPushedAtomId = 2;
   for (int i = 0; i < descriptor->field_count(); i++) {
     const FieldDescriptor *atomField = descriptor->field(i);
 
+    if (moduleName != DEFAULT_MODULE_NAME) {
+        const int moduleCount = atomField->options().ExtensionSize(os::statsd::module);
+        int j;
+        for (j = 0; j < moduleCount; ++j) {
+            const string atomModuleName = atomField->options().GetExtension(os::statsd::module, j);
+            if (atomModuleName == moduleName) {
+                break;
+            }
+        }
+
+        // This atom is not in the module we're interested in; skip it.
+        if (moduleCount == j) {
+            if (dbg) {
+              printf("   Skipping %s (%d)\n", atomField->name().c_str(), atomField->number());
+            }
+            continue;
+        }
+    }
+
     if (dbg) {
       printf("   %s (%d)\n", atomField->name().c_str(), atomField->number());
     }
@@ -441,27 +544,27 @@ int collate_atoms(const Descriptor *descriptor, Atoms *atoms) {
         atomDecl.whitelisted = true;
     }
 
-    for (int j = 0; j < atomField->options().ExtensionSize(os::statsd::module); ++j) {
-        const string moduleName = atomField->options().GetExtension(os::statsd::module, j);
-        atomDecl.moduleNames.insert(moduleName);
-    }
-
     vector<java_type_t> signature;
     errorCount += collate_atom(atom, &atomDecl, &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->signatures_to_modules[signature].insert(
-            atomDecl.moduleNames.begin(), atomDecl.moduleNames.end());
     atoms->decls.insert(atomDecl);
+    FieldNumberToAnnotations& fieldNumberToAnnotations = atoms->signatureInfoMap[signature];
+    populateFieldNumberToAnnotations(atomDecl, &fieldNumberToAnnotations);
 
     AtomDecl nonChainedAtomDecl(atomField->number(), atomField->name(), atom->name());
     vector<java_type_t> nonChainedSignature;
     if (get_non_chained_node(atom, &nonChainedAtomDecl, &nonChainedSignature)) {
-        atoms->non_chained_signatures_to_modules[nonChainedSignature].insert(
-            atomDecl.moduleNames.begin(), atomDecl.moduleNames.end());
         atoms->non_chained_decls.insert(nonChainedAtomDecl);
+        FieldNumberToAnnotations& fieldNumberToAnnotations =
+                atoms->nonChainedSignatureInfoMap[nonChainedSignature];
+        populateFieldNumberToAnnotations(atomDecl, &fieldNumberToAnnotations);
     }
 
     if (atomDecl.code < PULL_ATOM_START_ID && atomDecl.code > maxPushedAtomId) {
@@ -473,12 +576,12 @@ int collate_atoms(const Descriptor *descriptor, Atoms *atoms) {
 
   if (dbg) {
     printf("signatures = [\n");
-    for (map<vector<java_type_t>, set<string>>::const_iterator it =
-             atoms->signatures_to_modules.begin();
-         it != atoms->signatures_to_modules.end(); it++) {
+    for (map<vector<java_type_t>, FieldNumberToAnnotations>::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(); jt++) {
+           jt != it->first.end(); jt++){
         printf(" %d", (int)*jt);
       }
       printf("\n");
diff --git a/tools/stats_log_api_gen/Collation.h b/tools/stats_log_api_gen/Collation.h
index ace85e06f12df..d99b931338f2a 100644
--- a/tools/stats_log_api_gen/Collation.h
+++ b/tools/stats_log_api_gen/Collation.h
@@ -19,6 +19,7 @@
 
 
 #include <google/protobuf/descriptor.h>
+#include "frameworks/base/cmds/statsd/src/atom_field_options.pb.h"
 
 #include <set>
 #include <vector>
@@ -29,6 +30,7 @@ namespace stats_log_api_gen {
 
 using std::map;
 using std::set;
+using std::shared_ptr;
 using std::string;
 using std::vector;
 using google::protobuf::Descriptor;
@@ -38,6 +40,20 @@ const int PULL_ATOM_START_ID = 10000;
 
 const int FIRST_UID_IN_CHAIN_ID = 0;
 
+const unsigned char ANNOTATION_ID_IS_UID = 1;
+const unsigned char ANNOTATION_ID_TRUNCATE_TIMESTAMP = 2;
+const unsigned char ANNOTATION_ID_STATE_OPTION = 3;
+const unsigned char ANNOTATION_ID_DEFAULT_STATE = 4;
+const unsigned char ANNOTATION_ID_RESET_STATE = 5;
+const unsigned char ANNOTATION_ID_STATE_NESTED = 6;
+
+const int STATE_OPTION_UNSET = os::statsd::StateField::STATE_FIELD_UNSET;
+const int STATE_OPTION_EXCLUSIVE = os::statsd::StateField::EXCLUSIVE_STATE;
+const int STATE_OPTION_PRIMARY_FIELD_FIRST_UID = os::statsd::StateField::PRIMARY_FIELD_FIRST_UID;
+const int STATE_OPTION_PRIMARY = os::statsd::StateField::PRIMARY_FIELD;
+
+const string DEFAULT_MODULE_NAME = "DEFAULT";
+
 /**
  * The types for atom parameters.
  */
@@ -58,6 +74,38 @@ typedef enum {
   JAVA_TYPE_BYTE_ARRAY = -2,
 } java_type_t;
 
+enum AnnotationType {
+    ANNOTATION_TYPE_UNKNOWN = 0,
+    ANNOTATION_TYPE_INT = 1,
+    ANNOTATION_TYPE_BOOL = 2,
+};
+
+union AnnotationValue {
+    int intValue;
+    bool boolValue;
+
+    AnnotationValue(const int value): intValue(value) {}
+    AnnotationValue(const bool value): boolValue(value) {}
+};
+
+struct Annotation {
+    const unsigned char annotationId;
+    const int atomId;
+    AnnotationType type;
+    AnnotationValue value;
+
+    inline Annotation(unsigned char annotationId, int atomId, AnnotationType type,
+            AnnotationValue value):
+            annotationId(annotationId), atomId(atomId), type(type), value(value) {}
+    inline ~Annotation() {}
+
+    inline bool operator<(const Annotation& that) const {
+        return atomId == that.atomId ? annotationId < that.annotationId : atomId < that.atomId;
+    }
+};
+
+using FieldNumberToAnnotations =  map<int, set<shared_ptr<Annotation>>>;
+
 /**
  * The name and type for an atom field.
  */
@@ -72,6 +120,7 @@ struct AtomField {
     inline AtomField(const AtomField& that) :name(that.name),
                                              javaType(that.javaType),
                                              enumValues(that.enumValues) {}
+
     inline AtomField(string n, java_type_t jt) :name(n), javaType(jt) {}
     inline ~AtomField() {}
 };
@@ -86,6 +135,8 @@ struct AtomDecl {
     string message;
     vector<AtomField> fields;
 
+    FieldNumberToAnnotations fieldNumberToAnnotations;
+
     vector<int> primaryFields;
     int exclusiveField = 0;
     int defaultState = INT_MAX;
@@ -98,8 +149,6 @@ struct AtomDecl {
 
     vector<int> binaryFields;
 
-    set<string> moduleNames;
-
     AtomDecl();
     AtomDecl(const AtomDecl& that);
     AtomDecl(int code, const string& name, const string& message);
@@ -111,17 +160,17 @@ struct AtomDecl {
 };
 
 struct Atoms {
-    map<vector<java_type_t>, set<string>> signatures_to_modules;
+    map<vector<java_type_t>, FieldNumberToAnnotations> signatureInfoMap;
     set<AtomDecl> decls;
     set<AtomDecl> non_chained_decls;
-    map<vector<java_type_t>, set<string>> non_chained_signatures_to_modules;
+    map<vector<java_type_t>, FieldNumberToAnnotations> nonChainedSignatureInfoMap;
     int maxPushedAtomId;
 };
 
 /**
  * Gather the information about the atoms.  Returns the number of errors.
  */
-int collate_atoms(const Descriptor* descriptor, Atoms* atoms);
+int collate_atoms(const Descriptor* descriptor, const string& moduleName, Atoms* atoms);
 int collate_atom(const Descriptor *atom, AtomDecl *atomDecl, vector<java_type_t> *signature);
 
 }  // namespace stats_log_api_gen
diff --git a/tools/stats_log_api_gen/java_writer.cpp b/tools/stats_log_api_gen/java_writer.cpp
index 209d511c8a053..8da0801567745 100644
--- a/tools/stats_log_api_gen/java_writer.cpp
+++ b/tools/stats_log_api_gen/java_writer.cpp
@@ -23,10 +23,8 @@ namespace stats_log_api_gen {
 
 static int write_java_q_logger_class(
         FILE* out,
-        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
-        const AtomDecl &attributionDecl,
-        const string& moduleName
-        ) {
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& 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");
@@ -37,7 +35,7 @@ static int write_java_q_logger_class(
     fprintf(out, "\n");
     fprintf(out, "        // Write methods.\n");
     write_java_methods_q_schema(
-            out, signatures_to_modules, attributionDecl, moduleName, "        ");
+            out, signatureInfoMap, attributionDecl, "        ");
 
     fprintf(out, "    }\n");
     return 0;
@@ -46,21 +44,15 @@ static int write_java_q_logger_class(
 
 static int write_java_methods(
         FILE* out,
-        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
         const AtomDecl &attributionDecl,
-        const string& moduleName,
         const bool supportQ
         ) {
-    for (auto signature_to_modules_it = signatures_to_modules.begin();
-            signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
-        // Skip if this signature is not needed for the module.
-        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
-            continue;
-        }
-
+    for (auto signatureInfoMapIt = signatureInfoMap.begin();
+            signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
         // Print method signature.
         fprintf(out, "    public static void write(int code");
-        vector<java_type_t> signature = signature_to_modules_it->first;
+        vector<java_type_t> signature = signatureInfoMapIt->first;
         int argIndex = 1;
         for (vector<java_type_t>::const_iterator arg = signature.begin();
                 arg != signature.end(); arg++) {
@@ -249,7 +241,7 @@ static int write_java_methods(
 }
 
 int write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
-                                    const string& moduleName, const string& javaClass,
+                                    const string& javaClass,
                                     const string& javaPackage, const bool supportQ,
                                     const bool supportWorkSource) {
     // Print prelude
@@ -273,24 +265,24 @@ int write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attribut
     fprintf(out, " */\n");
     fprintf(out, "public class %s {\n", javaClass.c_str());
 
-    write_java_atom_codes(out, atoms, moduleName);
-    write_java_enum_values(out, atoms, moduleName);
+    write_java_atom_codes(out, atoms);
+    write_java_enum_values(out, atoms);
 
     int errors = 0;
 
     // Print write methods.
     fprintf(out, "    // Write methods\n");
     errors += write_java_methods(
-            out, atoms.signatures_to_modules, attributionDecl, moduleName, supportQ);
+            out, atoms.signatureInfoMap, attributionDecl, supportQ);
     errors += write_java_non_chained_methods(
-            out, atoms.non_chained_signatures_to_modules, moduleName);
+            out, atoms.nonChainedSignatureInfoMap);
     if (supportWorkSource) {
-        errors += write_java_work_source_methods(out, atoms.signatures_to_modules, moduleName);
+        errors += write_java_work_source_methods(out, atoms.signatureInfoMap);
     }
 
     if (supportQ) {
         errors += write_java_q_logger_class(
-                out, atoms.signatures_to_modules, attributionDecl, moduleName);
+                out, atoms.signatureInfoMap, attributionDecl);
     }
 
     fprintf(out, "}\n");
diff --git a/tools/stats_log_api_gen/java_writer.h b/tools/stats_log_api_gen/java_writer.h
index 5b78f059c5b95..4e1365e8dccef 100644
--- a/tools/stats_log_api_gen/java_writer.h
+++ b/tools/stats_log_api_gen/java_writer.h
@@ -31,7 +31,7 @@ namespace stats_log_api_gen {
 using namespace std;
 
 int write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
-                         const string& moduleName, const string& javaClass,
+                         const string& javaClass,
                          const string& javaPackage, const bool supportQ,
                          const bool supportWorkSource);
 
diff --git a/tools/stats_log_api_gen/java_writer_q.cpp b/tools/stats_log_api_gen/java_writer_q.cpp
index 8f2112a03e466..329c25def441c 100644
--- a/tools/stats_log_api_gen/java_writer_q.cpp
+++ b/tools/stats_log_api_gen/java_writer_q.cpp
@@ -51,20 +51,14 @@ 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>, set<string>>& signatures_to_modules,
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
         const AtomDecl &attributionDecl,
-        const string& moduleName,
         const string& indent) {
     int requiredHelpers = 0;
-    for (auto signature_to_modules_it = signatures_to_modules.begin();
-            signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
-        // Skip if this signature is not needed for the module.
-        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
-            continue;
-        }
-
+    for (auto signatureInfoMapIt = signatureInfoMap.begin();
+            signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
         // Print method signature.
-        vector<java_type_t> signature = signature_to_modules_it->first;
+        vector<java_type_t> signature = signatureInfoMapIt->first;
         fprintf(out, "%spublic static void write(int code", indent.c_str());
         int argIndex = 1;
         for (vector<java_type_t>::const_iterator arg = signature.begin();
@@ -568,7 +562,7 @@ void write_java_helpers_for_q_schema_methods(
 // This method is called in main.cpp to generate StatsLog for modules that's compatible with
 // Q at compile-time.
 int write_stats_log_java_q_for_module(FILE* out, const Atoms& atoms,
-                                      const AtomDecl &attributionDecl, const string& moduleName,
+                                      const AtomDecl &attributionDecl,
                                       const string& javaClass, const string& javaPackage,
                                       const bool supportWorkSource) {
     // Print prelude
@@ -589,19 +583,18 @@ int write_stats_log_java_q_for_module(FILE* out, const Atoms& atoms,
 
     write_java_q_logging_constants(out, "    ");
 
-    write_java_atom_codes(out, atoms, moduleName);
+    write_java_atom_codes(out, atoms);
 
-    write_java_enum_values(out, atoms, moduleName);
+    write_java_enum_values(out, atoms);
 
     int errors = 0;
     // Print write methods
     fprintf(out, "    // Write methods\n");
-    errors += write_java_methods_q_schema(out, atoms.signatures_to_modules, attributionDecl,
-            moduleName, "    ");
-    errors += write_java_non_chained_methods(out, atoms.non_chained_signatures_to_modules,
-            moduleName);
+    errors += write_java_methods_q_schema(out, atoms.signatureInfoMap, attributionDecl,
+            "    ");
+    errors += write_java_non_chained_methods(out, atoms.nonChainedSignatureInfoMap);
     if (supportWorkSource) {
-        errors += write_java_work_source_methods(out, atoms.signatures_to_modules, moduleName);
+        errors += write_java_work_source_methods(out, atoms.signatureInfoMap);
     }
 
     fprintf(out, "}\n");
diff --git a/tools/stats_log_api_gen/java_writer_q.h b/tools/stats_log_api_gen/java_writer_q.h
index 6ccb225ce9739..0f33b6cf94a1a 100644
--- a/tools/stats_log_api_gen/java_writer_q.h
+++ b/tools/stats_log_api_gen/java_writer_q.h
@@ -34,9 +34,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>, set<string>>& signatures_to_modules,
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
         const AtomDecl &attributionDecl,
-        const string& moduleName,
         const string& indent);
 
 void write_java_helpers_for_q_schema_methods(
@@ -46,7 +45,7 @@ void write_java_helpers_for_q_schema_methods(
         const string& indent);
 
 int write_stats_log_java_q_for_module(FILE* out, const Atoms& atoms,
-        const AtomDecl &attributionDecl, const string& moduleName, const string& javaClass,
+        const AtomDecl &attributionDecl, const string& javaClass,
         const string& javaPackage, const bool supportWorkSource);
 
 }  // namespace stats_log_api_gen
diff --git a/tools/stats_log_api_gen/main.cpp b/tools/stats_log_api_gen/main.cpp
index e9723a114e88b..4f791a35be7b6 100644
--- a/tools/stats_log_api_gen/main.cpp
+++ b/tools/stats_log_api_gen/main.cpp
@@ -194,7 +194,7 @@ run(int argc, char const*const* argv)
 
     // Collate the parameters
     Atoms atoms;
-    int errorCount = collate_atoms(Atom::descriptor(), &atoms);
+    int errorCount = collate_atoms(Atom::descriptor(), moduleName, &atoms);
     if (errorCount != 0) {
         return 1;
     }
@@ -246,7 +246,7 @@ run(int argc, char const*const* argv)
             return 1;
         }
         errorCount = android::stats_log_api_gen::write_stats_log_cpp(
-            out, atoms, attributionDecl, moduleName, cppNamespace, cppHeaderImport, supportQ);
+            out, atoms, attributionDecl, cppNamespace, cppHeaderImport, supportQ);
         fclose(out);
     }
 
@@ -262,7 +262,7 @@ run(int argc, char const*const* argv)
             fprintf(stderr, "Must supply --namespace if supplying a specific module\n");
         }
         errorCount = android::stats_log_api_gen::write_stats_log_header(
-            out, atoms, attributionDecl, moduleName, cppNamespace);
+            out, atoms, attributionDecl, cppNamespace);
         fclose(out);
     }
 
@@ -291,11 +291,11 @@ run(int argc, char const*const* argv)
 
         if (compileQ) {
             errorCount = android::stats_log_api_gen::write_stats_log_java_q_for_module(
-                    out, atoms, attributionDecl, moduleName, javaClass, javaPackage,
+                    out, atoms, attributionDecl, javaClass, javaPackage,
                     supportWorkSource);
         } else {
             errorCount = android::stats_log_api_gen::write_stats_log_java(
-                    out, atoms, attributionDecl, moduleName, javaClass, javaPackage, supportQ,
+                    out, atoms, attributionDecl, javaClass, javaPackage, supportQ,
                     supportWorkSource);
         }
 
diff --git a/tools/stats_log_api_gen/native_writer.cpp b/tools/stats_log_api_gen/native_writer.cpp
index da207d665e205..c06356abec1c6 100644
--- a/tools/stats_log_api_gen/native_writer.cpp
+++ b/tools/stats_log_api_gen/native_writer.cpp
@@ -21,14 +21,11 @@ namespace android {
 namespace stats_log_api_gen {
 
 static int write_native_stats_write_methods(FILE* out, const Atoms& atoms,
-        const AtomDecl& attributionDecl, const string& moduleName, const bool supportQ) {
+        const AtomDecl& attributionDecl, const bool supportQ) {
     fprintf(out, "\n");
-    for (auto signature_to_modules_it = atoms.signatures_to_modules.begin();
-        signature_to_modules_it != atoms.signatures_to_modules.end(); signature_to_modules_it++) {
-        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
-            continue;
-        }
-        vector<java_type_t> signature = signature_to_modules_it->first;
+    for (auto signatureInfoMapIt = atoms.signatureInfoMap.begin();
+            signatureInfoMapIt != atoms.signatureInfoMap.end(); signatureInfoMapIt++) {
+        vector<java_type_t> signature = signatureInfoMapIt->first;
         // Key value pairs not supported in native.
         if (find(signature.begin(), signature.end(), JAVA_TYPE_KEY_VALUE_PAIR) != signature.end()) {
             continue;
@@ -133,13 +130,10 @@ static int write_native_stats_write_methods(FILE* out, const Atoms& atoms,
 }
 
 static void write_native_stats_write_non_chained_methods(FILE* out, const Atoms& atoms,
-        const AtomDecl& attributionDecl, const string& moduleName) {
+        const AtomDecl& attributionDecl) {
     fprintf(out, "\n");
-    for (auto signature_it = atoms.non_chained_signatures_to_modules.begin();
-            signature_it != atoms.non_chained_signatures_to_modules.end(); signature_it++) {
-        if (!signature_needed_for_module(signature_it->second, moduleName)) {
-            continue;
-        }
+    for (auto signature_it = atoms.nonChainedSignatureInfoMap.begin();
+            signature_it != atoms.nonChainedSignatureInfoMap.end(); signature_it++) {
         vector<java_type_t> signature = signature_it->first;
         // Key value pairs not supported in native.
         if (find(signature.begin(), signature.end(), JAVA_TYPE_KEY_VALUE_PAIR) != signature.end()) {
@@ -176,16 +170,12 @@ static void write_native_stats_write_non_chained_methods(FILE* out, const Atoms&
 static void write_native_method_header(
         FILE* out,
         const string& methodName,
-        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
-        const AtomDecl &attributionDecl, const string& moduleName) {
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap,
+        const AtomDecl &attributionDecl) {
 
-    for (auto signature_to_modules_it = signatures_to_modules.begin();
-            signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
-        // Skip if this signature is not needed for the module.
-        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
-            continue;
-        }
-        vector<java_type_t> signature = signature_to_modules_it->first;
+    for (auto signatureInfoMapIt = signatureInfoMap.begin();
+            signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
+        vector<java_type_t> signature = signatureInfoMapIt->first;
 
         // Key value pairs not supported in native.
         if (find(signature.begin(), signature.end(), JAVA_TYPE_KEY_VALUE_PAIR) != signature.end()) {
@@ -196,7 +186,7 @@ static void write_native_method_header(
 }
 
 int write_stats_log_cpp(FILE *out, const Atoms &atoms, const AtomDecl &attributionDecl,
-                        const string& moduleName, const string& cppNamespace,
+                        const string& cppNamespace,
                         const string& importHeader, const bool supportQ) {
     // Print prelude
     fprintf(out, "// This file is autogenerated\n");
@@ -212,8 +202,8 @@ int write_stats_log_cpp(FILE *out, const Atoms &atoms, const AtomDecl &attributi
     fprintf(out, "\n");
     write_namespace(out, cppNamespace);
 
-    write_native_stats_write_methods(out, atoms, attributionDecl, moduleName, supportQ);
-    write_native_stats_write_non_chained_methods(out, atoms, attributionDecl, moduleName);
+    write_native_stats_write_methods(out, atoms, attributionDecl, supportQ);
+    write_native_stats_write_non_chained_methods(out, atoms, attributionDecl);
 
     // Print footer
     fprintf(out, "\n");
@@ -223,7 +213,7 @@ int write_stats_log_cpp(FILE *out, const Atoms &atoms, const AtomDecl &attributi
 }
 
 int write_stats_log_header(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
-        const string& moduleName, const string& cppNamespace) {
+        const string& cppNamespace) {
     // Print prelude
     fprintf(out, "// This file is autogenerated\n");
     fprintf(out, "\n");
@@ -242,7 +232,7 @@ int write_stats_log_header(FILE* out, const Atoms& atoms, const AtomDecl &attrib
     fprintf(out, " */\n");
     fprintf(out, "\n");
 
-    write_native_atom_constants(out, atoms, attributionDecl, moduleName);
+    write_native_atom_constants(out, atoms, attributionDecl);
 
     // Print constants for the enum values.
     fprintf(out, "//\n");
@@ -250,10 +240,6 @@ int write_stats_log_header(FILE* out, const Atoms& atoms, const AtomDecl &attrib
     fprintf(out, "//\n\n");
     for (set<AtomDecl>::const_iterator atom = atoms.decls.begin();
         atom != atoms.decls.end(); atom++) {
-        // Skip if the atom is not needed for the module.
-        if (!atom_needed_for_module(*atom, moduleName)) {
-            continue;
-        }
 
         for (vector<AtomField>::const_iterator field = atom->fields.begin();
             field != atom->fields.end(); field++) {
@@ -286,14 +272,13 @@ int write_stats_log_header(FILE* out, const Atoms& atoms, const AtomDecl &attrib
     fprintf(out, "//\n");
     fprintf(out, "// Write methods\n");
     fprintf(out, "//\n");
-    write_native_method_header(out, "int stats_write", atoms.signatures_to_modules, attributionDecl,
-            moduleName);
+    write_native_method_header(out, "int stats_write", atoms.signatureInfoMap, attributionDecl);
 
     fprintf(out, "//\n");
     fprintf(out, "// Write flattened methods\n");
     fprintf(out, "//\n");
     write_native_method_header(out, "int stats_write_non_chained",
-            atoms.non_chained_signatures_to_modules, attributionDecl, moduleName);
+            atoms.nonChainedSignatureInfoMap, attributionDecl);
 
     fprintf(out, "\n");
     write_closing_namespace(out, cppNamespace);
diff --git a/tools/stats_log_api_gen/native_writer.h b/tools/stats_log_api_gen/native_writer.h
index aafa96ece2d0c..6e603259600d1 100644
--- a/tools/stats_log_api_gen/native_writer.h
+++ b/tools/stats_log_api_gen/native_writer.h
@@ -27,11 +27,11 @@ namespace stats_log_api_gen {
 using namespace std;
 
 int write_stats_log_cpp(FILE *out, const Atoms &atoms, const AtomDecl &attributionDecl,
-        const string& moduleName, const string& cppNamespace, const string& importHeader,
+        const string& cppNamespace, const string& importHeader,
         const bool supportQ);
 
 int write_stats_log_header(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
-        const string& moduleName, const string& cppNamespace);
+        const string& cppNamespace);
 
 }  // namespace stats_log_api_gen
 }  // namespace android
diff --git a/tools/stats_log_api_gen/test.proto b/tools/stats_log_api_gen/test.proto
index f6c89c22ee65c..480442fe9d3dd 100644
--- a/tools/stats_log_api_gen/test.proto
+++ b/tools/stats_log_api_gen/test.proto
@@ -206,7 +206,7 @@ message ListedAtoms {
 }
 
 message ModuleOneAtom {
-    optional int32 field = 1;
+    optional int32 field = 1 [(android.os.statsd.is_uid) = true];
 }
 
 message ModuleTwoAtom {
@@ -214,7 +214,7 @@ message ModuleTwoAtom {
 }
 
 message ModuleOneAndTwoAtom {
-    optional int32 field = 1;
+    optional int32 field = 1 [(android.os.statsd.state_field_option).option = EXCLUSIVE_STATE];
 }
 
 message NoModuleAtom {
diff --git a/tools/stats_log_api_gen/test_collation.cpp b/tools/stats_log_api_gen/test_collation.cpp
index 73abaef1d91bc..5032ac088f4f3 100644
--- a/tools/stats_log_api_gen/test_collation.cpp
+++ b/tools/stats_log_api_gen/test_collation.cpp
@@ -29,10 +29,10 @@ using std::set;
 using std::vector;
 
 /**
- * Return whether the set contains a vector of the elements provided.
+ * Return whether the map contains a vector of the elements provided.
  */
 static bool
-set_contains_vector(const map<vector<java_type_t>, set<string>>& s, int count, ...)
+map_contains_vector(const map<vector<java_type_t>, FieldNumberToAnnotations>& s, int count, ...)
 {
     va_list args;
     vector<java_type_t> v;
@@ -47,12 +47,12 @@ set_contains_vector(const map<vector<java_type_t>, set<string>>& s, int count, .
 }
 
 /**
- * Expect that the provided set contains the elements provided.
+ * Expect that the provided map contains the elements provided.
  */
-#define EXPECT_SET_CONTAINS_SIGNATURE(s, ...) \
+#define EXPECT_MAP_CONTAINS_SIGNATURE(s, ...) \
     do { \
         int count = sizeof((int[]){__VA_ARGS__})/sizeof(int); \
-        EXPECT_TRUE(set_contains_vector(s, count, __VA_ARGS__)); \
+        EXPECT_TRUE(map_contains_vector(s, count, __VA_ARGS__)); \
     } while(0)
 
 /** Expects that the provided atom has no enum values for any field. */
@@ -83,20 +83,20 @@ set_contains_vector(const map<vector<java_type_t>, set<string>>& s, int count, .
  */
 TEST(CollationTest, CollateStats) {
     Atoms atoms;
-    int errorCount = collate_atoms(Event::descriptor(), &atoms);
+    int errorCount = collate_atoms(Event::descriptor(), DEFAULT_MODULE_NAME, &atoms);
 
     EXPECT_EQ(0, errorCount);
-    EXPECT_EQ(3ul, atoms.signatures_to_modules.size());
+    EXPECT_EQ(3ul, atoms.signatureInfoMap.size());
 
     // IntAtom, AnotherIntAtom
-    EXPECT_SET_CONTAINS_SIGNATURE(atoms.signatures_to_modules, JAVA_TYPE_INT);
+    EXPECT_MAP_CONTAINS_SIGNATURE(atoms.signatureInfoMap, JAVA_TYPE_INT);
 
     // OutOfOrderAtom
-    EXPECT_SET_CONTAINS_SIGNATURE(atoms.signatures_to_modules, JAVA_TYPE_INT, JAVA_TYPE_INT);
+    EXPECT_MAP_CONTAINS_SIGNATURE(atoms.signatureInfoMap, JAVA_TYPE_INT, JAVA_TYPE_INT);
 
     // AllTypesAtom
-    EXPECT_SET_CONTAINS_SIGNATURE(
-        atoms.signatures_to_modules,
+    EXPECT_MAP_CONTAINS_SIGNATURE(
+        atoms.signatureInfoMap,
         JAVA_TYPE_ATTRIBUTION_CHAIN, // AttributionChain
         JAVA_TYPE_FLOAT,             // float
         JAVA_TYPE_LONG,              // int64
@@ -150,7 +150,7 @@ TEST(CollationTest, CollateStats) {
  */
 TEST(CollationTest, NonMessageTypeFails) {
     Atoms atoms;
-    int errorCount = collate_atoms(IntAtom::descriptor(), &atoms);
+    int errorCount = collate_atoms(IntAtom::descriptor(), DEFAULT_MODULE_NAME, &atoms);
 
     EXPECT_EQ(1, errorCount);
 }
@@ -160,7 +160,7 @@ TEST(CollationTest, NonMessageTypeFails) {
  */
 TEST(CollationTest, FailOnBadTypes) {
     Atoms atoms;
-    int errorCount = collate_atoms(BadTypesEvent::descriptor(), &atoms);
+    int errorCount = collate_atoms(BadTypesEvent::descriptor(), DEFAULT_MODULE_NAME, &atoms);
 
     EXPECT_EQ(4, errorCount);
 }
@@ -170,7 +170,7 @@ TEST(CollationTest, FailOnBadTypes) {
  */
 TEST(CollationTest, FailOnSkippedFieldsSingle) {
     Atoms atoms;
-    int errorCount = collate_atoms(BadSkippedFieldSingle::descriptor(), &atoms);
+    int errorCount = collate_atoms(BadSkippedFieldSingle::descriptor(), DEFAULT_MODULE_NAME, &atoms);
 
     EXPECT_EQ(1, errorCount);
 }
@@ -181,7 +181,7 @@ TEST(CollationTest, FailOnSkippedFieldsSingle) {
  */
 TEST(CollationTest, FailOnSkippedFieldsMultiple) {
     Atoms atoms;
-    int errorCount = collate_atoms(BadSkippedFieldMultiple::descriptor(), &atoms);
+    int errorCount = collate_atoms(BadSkippedFieldMultiple::descriptor(), DEFAULT_MODULE_NAME, &atoms);
 
     EXPECT_EQ(2, errorCount);
 }
@@ -193,48 +193,48 @@ TEST(CollationTest, FailOnSkippedFieldsMultiple) {
 TEST(CollationTest, FailBadAttributionNodePosition) {
   Atoms atoms;
   int errorCount =
-      collate_atoms(BadAttributionNodePosition::descriptor(), &atoms);
+      collate_atoms(BadAttributionNodePosition::descriptor(), DEFAULT_MODULE_NAME, &atoms);
 
   EXPECT_EQ(1, errorCount);
 }
 
 TEST(CollationTest, FailOnBadStateAtomOptions) {
     Atoms atoms;
-    int errorCount = collate_atoms(BadStateAtoms::descriptor(), &atoms);
+    int errorCount = collate_atoms(BadStateAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
 
     EXPECT_EQ(3, errorCount);
 }
 
 TEST(CollationTest, PassOnGoodStateAtomOptions) {
     Atoms atoms;
-    int errorCount = collate_atoms(GoodStateAtoms::descriptor(), &atoms);
+    int errorCount = collate_atoms(GoodStateAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     EXPECT_EQ(0, errorCount);
 }
 
 TEST(CollationTest, PassOnGoodBinaryFieldAtom) {
     Atoms atoms;
     int errorCount =
-            collate_atoms(GoodEventWithBinaryFieldAtom::descriptor(), &atoms);
+            collate_atoms(GoodEventWithBinaryFieldAtom::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     EXPECT_EQ(0, errorCount);
 }
 
 TEST(CollationTest, FailOnBadBinaryFieldAtom) {
     Atoms atoms;
     int errorCount =
-            collate_atoms(BadEventWithBinaryFieldAtom::descriptor(), &atoms);
+            collate_atoms(BadEventWithBinaryFieldAtom::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     EXPECT_TRUE(errorCount > 0);
 }
 
 TEST(CollationTest, PassOnWhitelistedAtom) {
     Atoms atoms;
-    int errorCount = collate_atoms(ListedAtoms::descriptor(), &atoms);
+    int errorCount = collate_atoms(ListedAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     EXPECT_EQ(errorCount, 0);
     EXPECT_EQ(atoms.decls.size(), 2ul);
 }
 
 TEST(CollationTest, RecogniseWhitelistedAtom) {
     Atoms atoms;
-    collate_atoms(ListedAtoms::descriptor(), &atoms);
+    collate_atoms(ListedAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     for (const auto& atomDecl : atoms.decls) {
         if (atomDecl.code == 1) {
             EXPECT_TRUE(atomDecl.whitelisted);
@@ -246,45 +246,80 @@ TEST(CollationTest, RecogniseWhitelistedAtom) {
 
 TEST(CollationTest, PassOnLogFromModuleAtom) {
     Atoms atoms;
-    int errorCount = collate_atoms(ModuleAtoms::descriptor(), &atoms);
+    int errorCount = collate_atoms(ModuleAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     EXPECT_EQ(errorCount, 0);
     EXPECT_EQ(atoms.decls.size(), 4ul);
 }
 
 TEST(CollationTest, RecognizeModuleAtom) {
     Atoms atoms;
-    int errorCount = collate_atoms(ModuleAtoms::descriptor(), &atoms);
+    int errorCount = collate_atoms(ModuleAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
     EXPECT_EQ(errorCount, 0);
     EXPECT_EQ(atoms.decls.size(), 4ul);
-    for (const auto& atomDecl: atoms.decls) {
-        if (atomDecl.code == 1) {
-            EXPECT_EQ(1ul, atomDecl.moduleNames.size());
-            EXPECT_NE(atomDecl.moduleNames.end(), atomDecl.moduleNames.find("module1"));
-        } else if (atomDecl.code == 2) {
-            EXPECT_EQ(1ul, atomDecl.moduleNames.size());
-            EXPECT_NE(atomDecl.moduleNames.end(), atomDecl.moduleNames.find("module2"));
-        } else if (atomDecl.code == 3) {
-            EXPECT_EQ(2ul, atomDecl.moduleNames.size());
-            EXPECT_NE(atomDecl.moduleNames.end(), atomDecl.moduleNames.find("module1"));
-            EXPECT_NE(atomDecl.moduleNames.end(), atomDecl.moduleNames.find("module2"));
-        } else {
-            EXPECT_TRUE(atomDecl.moduleNames.empty());
+    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);
+                }
+            }
         }
     }
+}
 
-    EXPECT_EQ(atoms.signatures_to_modules.size(), 2u);
-    EXPECT_SET_CONTAINS_SIGNATURE(atoms.signatures_to_modules, JAVA_TYPE_INT);
-    EXPECT_SET_CONTAINS_SIGNATURE(atoms.signatures_to_modules, JAVA_TYPE_STRING);
-    for (auto signature_to_modules_it : atoms.signatures_to_modules) {
-        vector<java_type_t> signature = signature_to_modules_it.first;
-        if (signature[0] == JAVA_TYPE_STRING) {
-            EXPECT_EQ(signature_to_modules_it.second.size(), 0u);
-        } else if (signature[0] == JAVA_TYPE_INT) {
-            set<string> modules = signature_to_modules_it.second;
-            EXPECT_EQ(modules.size(), 2u);
-            // Assert that the set contains "module1" and "module2".
-            EXPECT_NE(modules.find("module1"), modules.end());
-            EXPECT_NE(modules.find("module2"), modules.end());
+TEST(CollationTest, RecognizeModule1Atom) {
+    Atoms atoms;
+    const string moduleName = "module1";
+    int errorCount = collate_atoms(ModuleAtoms::descriptor(), moduleName, &atoms);
+    EXPECT_EQ(errorCount, 0);
+    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);
+            }
         }
     }
 }
diff --git a/tools/stats_log_api_gen/utils.cpp b/tools/stats_log_api_gen/utils.cpp
index 9dc4ff8273806..cc9117c468af0 100644
--- a/tools/stats_log_api_gen/utils.cpp
+++ b/tools/stats_log_api_gen/utils.cpp
@@ -98,20 +98,6 @@ const char* java_type_name(java_type_t type) {
     }
 }
 
-bool atom_needed_for_module(const AtomDecl& atomDecl, const string& moduleName) {
-    if (moduleName == DEFAULT_MODULE_NAME) {
-        return true;
-    }
-    return atomDecl.moduleNames.find(moduleName) != atomDecl.moduleNames.end();
-}
-
-bool signature_needed_for_module(const set<string>& modules, const string& moduleName) {
-    if (moduleName == DEFAULT_MODULE_NAME) {
-        return true;
-    }
-    return modules.find(moduleName) != modules.end();
-}
-
 // Native
 // Writes namespaces for the cpp and header files, returning the number of namespaces written.
 void write_namespace(FILE* out, const string& cppNamespaces) {
@@ -165,8 +151,7 @@ static void write_cpp_usage(
     fprintf(out, ");\n");
 }
 
-void write_native_atom_constants(FILE* out, const Atoms& atoms, const AtomDecl& attributionDecl,
-        const string& moduleName) {
+void write_native_atom_constants(FILE* out, const Atoms& atoms, const AtomDecl& attributionDecl) {
     fprintf(out, "/**\n");
     fprintf(out, " * Constants for atom codes.\n");
     fprintf(out, " */\n");
@@ -179,10 +164,6 @@ void write_native_atom_constants(FILE* out, const Atoms& atoms, const AtomDecl&
     // Print atom constants
     for (set<AtomDecl>::const_iterator atom = atoms.decls.begin();
         atom != atoms.decls.end(); atom++) {
-        // Skip if the atom is not needed for the module.
-        if (!atom_needed_for_module(*atom, moduleName)) {
-            continue;
-        }
         string constant = make_constant_name(atom->name);
         fprintf(out, "\n");
         fprintf(out, "    /**\n");
@@ -264,7 +245,7 @@ void write_native_method_call(FILE* out, const string& methodName,
 }
 
 // Java
-void write_java_atom_codes(FILE* out, const Atoms& atoms, const string& moduleName) {
+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;
@@ -273,10 +254,6 @@ void write_java_atom_codes(FILE* out, const Atoms& atoms, const string& moduleNa
     // Print constants for the atom codes.
     for (set<AtomDecl>::const_iterator atom = atoms.decls.begin();
             atom != atoms.decls.end(); atom++) {
-        // Skip if the atom is not needed for the module.
-        if (!atom_needed_for_module(*atom, moduleName)) {
-            continue;
-        }
         string constant = make_constant_name(atom->name);
         fprintf(out, "\n");
         fprintf(out, "    /**\n");
@@ -292,14 +269,10 @@ void write_java_atom_codes(FILE* out, const Atoms& atoms, const string& moduleNa
     fprintf(out, "\n");
 }
 
-void write_java_enum_values(FILE* out, const Atoms& atoms, const string& moduleName) {
+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++) {
-        // Skip if the atom is not needed for the module.
-        if (!atom_needed_for_module(*atom, moduleName)) {
-            continue;
-        }
         for (vector<AtomField>::const_iterator field = atom->fields.begin();
             field != atom->fields.end(); field++) {
             if (field->javaType == JAVA_TYPE_ENUM) {
@@ -340,19 +313,12 @@ void write_java_usage(FILE* out, const string& method_name, const string& atom_c
 
 int write_java_non_chained_methods(
         FILE* out,
-        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
-        const string& moduleName
-        ) {
-    for (auto signature_to_modules_it = signatures_to_modules.begin();
-            signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
-        // Skip if this signature is not needed for the module.
-        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
-            continue;
-        }
-
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap) {
+    for (auto signatureInfoMapIt = signatureInfoMap.begin();
+            signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
         // Print method signature.
         fprintf(out, "    public static void write_non_chained(int code");
-        vector<java_type_t> signature = signature_to_modules_it->first;
+        vector<java_type_t> signature = signatureInfoMapIt->first;
         int argIndex = 1;
         for (vector<java_type_t>::const_iterator arg = signature.begin();
                 arg != signature.end(); arg++) {
@@ -392,17 +358,11 @@ int write_java_non_chained_methods(
 
 int write_java_work_source_methods(
         FILE* out,
-        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
-        const string& moduleName
-        ) {
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap) {
     fprintf(out, "    // WorkSource methods.\n");
-    for (auto signature_to_modules_it = signatures_to_modules.begin();
-            signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
-        // Skip if this signature is not needed for the module.
-        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
-            continue;
-        }
-        vector<java_type_t> signature = signature_to_modules_it->first;
+    for (auto signatureInfoMapIt = signatureInfoMap.begin();
+            signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
+        vector<java_type_t> signature = signatureInfoMapIt->first;
         // Determine if there is Attribution in this signature.
         int attributionArg = -1;
         int argIndexMax = 0;
diff --git a/tools/stats_log_api_gen/utils.h b/tools/stats_log_api_gen/utils.h
index 715d42bc160a8..a6b3ef9fe99e0 100644
--- a/tools/stats_log_api_gen/utils.h
+++ b/tools/stats_log_api_gen/utils.h
@@ -30,7 +30,6 @@ namespace stats_log_api_gen {
 
 using namespace std;
 
-const string DEFAULT_MODULE_NAME = "DEFAULT";
 const string DEFAULT_CPP_NAMESPACE = "android,util";
 const string DEFAULT_CPP_HEADER_IMPORT = "statslog.h";
 const string DEFAULT_ATOMS_INFO_CPP_HEADER_IMPORT = "atoms_info.h";
@@ -45,17 +44,12 @@ const char* cpp_type_name(java_type_t type);
 
 const char* java_type_name(java_type_t type);
 
-bool atom_needed_for_module(const AtomDecl& atomDecl, const string& moduleName);
-
-bool signature_needed_for_module(const set<string>& modules, const string& moduleName);
-
 // Common Native helpers
 void write_namespace(FILE* out, const string& cppNamespaces);
 
 void write_closing_namespace(FILE* out, const string& cppNamespaces);
 
-void write_native_atom_constants(FILE* out, const Atoms& atoms, const AtomDecl& attributionDecl,
-        const string& moduleName);
+void write_native_atom_constants(FILE* out, const Atoms& atoms, const AtomDecl& attributionDecl);
 
 void write_native_method_signature(FILE* out, const string& methodName,
         const vector<java_type_t>& signature, const AtomDecl& attributionDecl,
@@ -65,21 +59,19 @@ void write_native_method_call(FILE* out, const string& methodName,
         const vector<java_type_t>& signature, const AtomDecl& attributionDecl, int argIndex = 1);
 
 // Common Java helpers.
-void write_java_atom_codes(FILE* out, const Atoms& atoms, const string& moduleName);
+void write_java_atom_codes(FILE* out, const Atoms& atoms);
 
-void write_java_enum_values(FILE* out, const Atoms& atoms, const string& moduleName);
+void write_java_enum_values(FILE* out, const Atoms& atoms);
 
 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>,
-        set<string>>& signatures_to_modules,
-        const string& moduleName);
+        FieldNumberToAnnotations>& signatureInfoMap);
 
 int write_java_work_source_methods(
         FILE* out,
-        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
-        const string& moduleName);
+        const map<vector<java_type_t>, FieldNumberToAnnotations>& signatureInfoMap);
 
 }  // namespace stats_log_api_gen
 }  // namespace android