Merge Android R (rvc-dev-plus-aosp-without-vendor@6692709)

Bug: 166295507
Merged-In: I3d92a6de21a938f6b352ec26dc23420c0fe02b27
Change-Id: Ifdb80563ef042738778ebb8a7581a97c4e3d96e2

1 files changed, 373 insertions, 275 deletions

diff --git a/tools/stats_log_api_gen/Collation.cpp b/tools/stats_log_api_gen/Collation.cpp
index 75deb017e41b..a230de46dcf3 100644
--- a/tools/stats_log_api_gen/Collation.cpp
+++ b/tools/stats_log_api_gen/Collation.cpp
@@ -15,11 +15,13 @@
  */
 
 #include "Collation.h"
-#include "frameworks/base/cmds/statsd/src/atoms.pb.h"
 
 #include <stdio.h>
+
 #include <map>
 
+#include "frameworks/base/cmds/statsd/src/atoms.pb.h"
+
 namespace android {
 namespace stats_log_api_gen {
 
@@ -27,17 +29,16 @@ 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
 //
 
-AtomDecl::AtomDecl()
-    :code(0),
-     name()
-{
+AtomDecl::AtomDecl() : code(0), name() {
 }
 
 AtomDecl::AtomDecl(const AtomDecl& that)
@@ -45,38 +46,33 @@ AtomDecl::AtomDecl(const AtomDecl& that)
       name(that.name),
       message(that.message),
       fields(that.fields),
+      fieldNumberToAnnotations(that.fieldNumberToAnnotations),
       primaryFields(that.primaryFields),
       exclusiveField(that.exclusiveField),
-      uidField(that.uidField),
-      whitelisted(that.whitelisted),
-      binaryFields(that.binaryFields),
-      hasModule(that.hasModule),
-      moduleName(that.moduleName) {}
-
-AtomDecl::AtomDecl(int c, const string& n, const string& m)
-    :code(c),
-     name(n),
-     message(m)
-{
+      defaultState(that.defaultState),
+      triggerStateReset(that.triggerStateReset),
+      nested(that.nested),
+      uidField(that.uidField) {
 }
 
-AtomDecl::~AtomDecl()
-{
+AtomDecl::AtomDecl(int c, const string& n, const string& m) : code(c), name(n), message(m) {
 }
 
+AtomDecl::~AtomDecl() {
+}
 
 /**
- * Print an error message for a FieldDescriptor, including the file name and line number.
+ * Print an error message for a FieldDescriptor, including the file name and
+ * line number.
  */
-static void
-print_error(const FieldDescriptor* field, const char* format, ...)
-{
+static void print_error(const FieldDescriptor* field, const char* format, ...) {
     const Descriptor* message = field->containing_type();
     const FileDescriptor* file = message->file();
 
     SourceLocation loc;
     if (field->GetSourceLocation(&loc)) {
-        // TODO: this will work if we can figure out how to pass --include_source_info to protoc
+        // TODO: this will work if we can figure out how to pass
+        // --include_source_info to protoc
         fprintf(stderr, "%s:%d: ", file->name().c_str(), loc.start_line);
     } else {
         fprintf(stderr, "%s: ", file->name().c_str());
@@ -84,15 +80,13 @@ print_error(const FieldDescriptor* field, const char* format, ...)
     va_list args;
     va_start(args, format);
     vfprintf(stderr, format, args);
-    va_end (args);
+    va_end(args);
 }
 
 /**
  * Convert a protobuf type into a java type.
  */
-static java_type_t
-java_type(const FieldDescriptor* field)
-{
+static java_type_t java_type(const FieldDescriptor* field) {
     int protoType = field->type();
     switch (protoType) {
         case FieldDescriptor::TYPE_DOUBLE:
@@ -117,12 +111,10 @@ java_type(const FieldDescriptor* field)
             return JAVA_TYPE_UNKNOWN;
         case FieldDescriptor::TYPE_MESSAGE:
             // TODO: not the final package name
-            if (field->message_type()->full_name() ==
-                "android.os.statsd.AttributionNode") {
-              return JAVA_TYPE_ATTRIBUTION_CHAIN;
-            } else if (field->message_type()->full_name() ==
-                       "android.os.statsd.KeyValuePair") {
-              return JAVA_TYPE_KEY_VALUE_PAIR;
+            if (field->message_type()->full_name() == "android.os.statsd.AttributionNode") {
+                return JAVA_TYPE_ATTRIBUTION_CHAIN;
+            } else if (field->message_type()->full_name() == "android.os.statsd.KeyValuePair") {
+                return JAVA_TYPE_KEY_VALUE_PAIR;
             } else if (field->options().GetExtension(os::statsd::log_mode) ==
                        os::statsd::LogMode::MODE_BYTES) {
                 return JAVA_TYPE_BYTE_ARRAY;
@@ -151,210 +143,307 @@ java_type(const FieldDescriptor* field)
 /**
  * Gather the enums info.
  */
-void collate_enums(const EnumDescriptor &enumDescriptor, AtomField *atomField) {
+void collate_enums(const EnumDescriptor& enumDescriptor, AtomField* atomField) {
     for (int i = 0; i < enumDescriptor.value_count(); i++) {
         atomField->enumValues[enumDescriptor.value(i)->number()] =
-            enumDescriptor.value(i)->name().c_str();
+                enumDescriptor.value(i)->name().c_str();
     }
 }
 
-/**
- * Gather the info about an atom proto.
- */
-int collate_atom(const Descriptor *atom, AtomDecl *atomDecl,
-                 vector<java_type_t> *signature) {
-
-  int errorCount = 0;
-
-  // Build a sorted list of the fields. Descriptor has them in source file
-  // order.
-  map<int, const FieldDescriptor *> fields;
-  for (int j = 0; j < atom->field_count(); j++) {
-    const FieldDescriptor *field = atom->field(j);
-    fields[field->number()] = field;
-  }
-
-  // Check that the parameters start at 1 and go up sequentially.
-  int expectedNumber = 1;
-  for (map<int, const FieldDescriptor *>::const_iterator it = fields.begin();
-       it != fields.end(); it++) {
-    const int number = it->first;
-    const FieldDescriptor *field = it->second;
-    if (number != expectedNumber) {
-      print_error(field,
-                  "Fields must be numbered consecutively starting at 1:"
-                  " '%s' is %d but should be %d\n",
-                  field->name().c_str(), number, expectedNumber);
-      errorCount++;
-      expectedNumber = number;
-      continue;
-    }
-    expectedNumber++;
-  }
-
-  // Check that only allowed types are present. Remove any invalid ones.
-  for (map<int, const FieldDescriptor *>::const_iterator it = fields.begin();
-       it != fields.end(); it++) {
-    const FieldDescriptor *field = it->second;
-    bool isBinaryField = field->options().GetExtension(os::statsd::log_mode) ==
-                         os::statsd::LogMode::MODE_BYTES;
-
-    java_type_t javaType = java_type(field);
-
-    if (javaType == JAVA_TYPE_UNKNOWN) {
-      print_error(field, "Unkown type for field: %s\n", field->name().c_str());
-      errorCount++;
-      continue;
-    } else if (javaType == JAVA_TYPE_OBJECT &&
-               atomDecl->code < PULL_ATOM_START_ID) {
-        // Allow attribution chain, but only at position 1.
-        print_error(field,
-                    "Message type not allowed for field in pushed atoms: %s\n",
-                    field->name().c_str());
-        errorCount++;
-        continue;
-    } else if (javaType == JAVA_TYPE_BYTE_ARRAY && !isBinaryField) {
-        print_error(field, "Raw bytes type not allowed for field: %s\n",
-                    field->name().c_str());
-        errorCount++;
-        continue;
+static void addAnnotationToAtomDecl(AtomDecl* atomDecl, const int fieldNumber,
+                                    const AnnotationId 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));
+}
 
-    if (isBinaryField && javaType != JAVA_TYPE_BYTE_ARRAY) {
-        print_error(field, "Cannot mark field %s as bytes.\n",
-                    field->name().c_str());
-        errorCount++;
-        continue;
-    }
+static int collate_field_annotations(AtomDecl* atomDecl, const FieldDescriptor* field,
+                                     const int fieldNumber, const java_type_t& javaType) {
+    int errorCount = 0;
+
+    if (field->options().HasExtension(os::statsd::state_field_option)) {
+        const os::statsd::StateAtomFieldOption& stateFieldOption =
+                field->options().GetExtension(os::statsd::state_field_option);
+        const bool primaryField = stateFieldOption.primary_field();
+        const bool exclusiveState = stateFieldOption.exclusive_state();
+        const bool primaryFieldFirstUid = stateFieldOption.primary_field_first_uid();
+
+        // Check the field is only one of primaryField, exclusiveState, or primaryFieldFirstUid.
+        if (primaryField + primaryFieldFirstUid + exclusiveState > 1) {
+            print_error(field,
+                        "Field can be max 1 of primary_field, exclusive_state, "
+                        "or primary_field_first_uid: '%s'\n",
+                        atomDecl->message.c_str());
+            errorCount++;
+        }
 
-    // Doubles are not supported yet.
-    if (javaType == JAVA_TYPE_DOUBLE) {
-        print_error(field, "Doubles are not supported in atoms. Please change field %s to float\n",
-                    field->name().c_str());
-        errorCount++;
-        continue;
-    }
-  }
-
-  // Check that if there's an attribution chain, it's at position 1.
-  for (map<int, const FieldDescriptor *>::const_iterator it = fields.begin();
-       it != fields.end(); it++) {
-    int number = it->first;
-    if (number != 1) {
-      const FieldDescriptor *field = it->second;
-      java_type_t javaType = java_type(field);
-      if (javaType == JAVA_TYPE_ATTRIBUTION_CHAIN) {
-        print_error(
-            field,
-            "AttributionChain fields must have field id 1, in message: '%s'\n",
-            atom->name().c_str());
-        errorCount++;
-      }
+        if (primaryField) {
+            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",
+                            atomDecl->message.c_str());
+                errorCount++;
+            } else {
+                atomDecl->primaryFields.push_back(fieldNumber);
+                addAnnotationToAtomDecl(atomDecl, fieldNumber, ANNOTATION_ID_PRIMARY_FIELD,
+                                        ANNOTATION_TYPE_BOOL, AnnotationValue(true));
+            }
+        }
+
+        if (primaryFieldFirstUid) {
+            if (javaType != JAVA_TYPE_ATTRIBUTION_CHAIN) {
+                print_error(field,
+                            "PRIMARY_FIELD_FIRST_UID annotation is only for AttributionChains: "
+                            "'%s'\n",
+                            atomDecl->message.c_str());
+                errorCount++;
+            } else {
+                atomDecl->primaryFields.push_back(FIRST_UID_IN_CHAIN_ID);
+                addAnnotationToAtomDecl(atomDecl, fieldNumber,
+                                        ANNOTATION_ID_PRIMARY_FIELD_FIRST_UID, ANNOTATION_TYPE_BOOL,
+                                        AnnotationValue(true));
+            }
+        }
+
+        if (exclusiveState) {
+            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",
+                            atomDecl->message.c_str());
+                errorCount++;
+            }
+
+            if (atomDecl->exclusiveField != 0) {
+                print_error(field,
+                            "Cannot have more than one exclusive state field in an "
+                            "atom: '%s'\n",
+                            atomDecl->message.c_str());
+                errorCount++;
+            } else {
+                atomDecl->exclusiveField = fieldNumber;
+                addAnnotationToAtomDecl(atomDecl, fieldNumber, ANNOTATION_ID_EXCLUSIVE_STATE,
+                                        ANNOTATION_TYPE_BOOL, AnnotationValue(true));
+            }
+
+            if (stateFieldOption.has_default_state_value()) {
+                const int defaultState = stateFieldOption.default_state_value();
+                atomDecl->defaultState = defaultState;
+
+                addAnnotationToAtomDecl(atomDecl, fieldNumber, ANNOTATION_ID_DEFAULT_STATE,
+                                        ANNOTATION_TYPE_INT, AnnotationValue(defaultState));
+            }
+
+            if (stateFieldOption.has_trigger_state_reset_value()) {
+                const int triggerStateReset = stateFieldOption.trigger_state_reset_value();
+
+                atomDecl->triggerStateReset = triggerStateReset;
+                addAnnotationToAtomDecl(atomDecl, fieldNumber, ANNOTATION_ID_TRIGGER_STATE_RESET,
+                                        ANNOTATION_TYPE_INT, AnnotationValue(triggerStateReset));
+            }
+
+            if (stateFieldOption.has_nested()) {
+                const bool nested = stateFieldOption.nested();
+                atomDecl->nested = nested;
+
+                addAnnotationToAtomDecl(atomDecl, fieldNumber, ANNOTATION_ID_STATE_NESTED,
+                                        ANNOTATION_TYPE_BOOL, AnnotationValue(nested));
+            }
+        }
     }
-  }
-
-  // Build the type signature and the atom data.
-  for (map<int, const FieldDescriptor *>::const_iterator it = fields.begin();
-       it != fields.end(); it++) {
-    const FieldDescriptor *field = it->second;
-    java_type_t javaType = java_type(field);
-    bool isBinaryField = field->options().GetExtension(os::statsd::log_mode) ==
-                         os::statsd::LogMode::MODE_BYTES;
-
-    AtomField atField(field->name(), javaType);
-    // Generate signature for pushed atoms
-    if (atomDecl->code < PULL_ATOM_START_ID) {
-      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 {
-          signature->push_back(javaType);
-      }
+
+    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",
+                        atomDecl->message.c_str());
+            errorCount++;
+        }
+
+        if (atomDecl->uidField == 0) {
+            atomDecl->uidField = fieldNumber;
+
+            addAnnotationToAtomDecl(atomDecl, fieldNumber, 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",
+                        atomDecl->message.c_str());
+            errorCount++;
+        }
     }
-    if (javaType == JAVA_TYPE_ENUM) {
-      // All enums are treated as ints when it comes to function signatures.
-      collate_enums(*field->enum_type(), &atField);
+
+    return errorCount;
+}
+
+/**
+ * Gather the info about an atom proto.
+ */
+int collate_atom(const Descriptor* atom, AtomDecl* atomDecl, vector<java_type_t>* signature) {
+    int errorCount = 0;
+
+    // Build a sorted list of the fields. Descriptor has them in source file
+    // order.
+    map<int, const FieldDescriptor*> fields;
+    for (int j = 0; j < atom->field_count(); j++) {
+        const FieldDescriptor* field = atom->field(j);
+        fields[field->number()] = field;
     }
-    atomDecl->fields.push_back(atField);
 
-    if (field->options().GetExtension(os::statsd::state_field_option).option() ==
-        os::statsd::StateField::PRIMARY) {
-        if (javaType == JAVA_TYPE_UNKNOWN ||
-            javaType == JAVA_TYPE_ATTRIBUTION_CHAIN ||
-            javaType == JAVA_TYPE_OBJECT || javaType == JAVA_TYPE_BYTE_ARRAY) {
+    // Check that the parameters start at 1 and go up sequentially.
+    int expectedNumber = 1;
+    for (map<int, const FieldDescriptor*>::const_iterator it = fields.begin(); it != fields.end();
+         it++) {
+        const int number = it->first;
+        const FieldDescriptor* field = it->second;
+        if (number != expectedNumber) {
+            print_error(field,
+                        "Fields must be numbered consecutively starting at 1:"
+                        " '%s' is %d but should be %d\n",
+                        field->name().c_str(), number, expectedNumber);
             errorCount++;
+            expectedNumber = number;
+            continue;
         }
-        atomDecl->primaryFields.push_back(it->first);
+        expectedNumber++;
     }
 
-    if (field->options().GetExtension(os::statsd::state_field_option).option() ==
-        os::statsd::StateField::EXCLUSIVE) {
-        if (javaType == JAVA_TYPE_UNKNOWN ||
-            javaType == JAVA_TYPE_ATTRIBUTION_CHAIN ||
-            javaType == JAVA_TYPE_OBJECT || javaType == JAVA_TYPE_BYTE_ARRAY) {
+    // Check that only allowed types are present. Remove any invalid ones.
+    for (map<int, const FieldDescriptor*>::const_iterator it = fields.begin(); it != fields.end();
+         it++) {
+        const FieldDescriptor* field = it->second;
+        bool isBinaryField = field->options().GetExtension(os::statsd::log_mode) ==
+                             os::statsd::LogMode::MODE_BYTES;
+
+        java_type_t javaType = java_type(field);
+
+        if (javaType == JAVA_TYPE_UNKNOWN) {
+            print_error(field, "Unknown type for field: %s\n", field->name().c_str());
+            errorCount++;
+            continue;
+        } else if (javaType == JAVA_TYPE_OBJECT && atomDecl->code < PULL_ATOM_START_ID) {
+            // Allow attribution chain, but only at position 1.
+            print_error(field, "Message type not allowed for field in pushed atoms: %s\n",
+                        field->name().c_str());
+            errorCount++;
+            continue;
+        } else if (javaType == JAVA_TYPE_BYTE_ARRAY && !isBinaryField) {
+            print_error(field, "Raw bytes type not allowed for field: %s\n", field->name().c_str());
             errorCount++;
+            continue;
         }
 
-        if (atomDecl->exclusiveField == 0) {
-            atomDecl->exclusiveField = it->first;
-        } else {
+        if (isBinaryField && javaType != JAVA_TYPE_BYTE_ARRAY) {
+            print_error(field, "Cannot mark field %s as bytes.\n", field->name().c_str());
             errorCount++;
+            continue;
         }
-    }
 
-    if (field->options().GetExtension(os::statsd::is_uid) == true) {
-        if (javaType != JAVA_TYPE_INT) {
+        // Doubles are not supported yet.
+        if (javaType == JAVA_TYPE_DOUBLE) {
+            print_error(field,
+                        "Doubles are not supported in atoms. Please change field %s "
+                        "to float\n",
+                        field->name().c_str());
             errorCount++;
+            continue;
         }
 
-        if (atomDecl->uidField == 0) {
-            atomDecl->uidField = it->first;
-        } else {
+        if (field->is_repeated() &&
+            !(javaType == JAVA_TYPE_ATTRIBUTION_CHAIN || javaType == JAVA_TYPE_KEY_VALUE_PAIR)) {
+            print_error(field,
+                        "Repeated fields are not supported in atoms. Please make "
+                        "field %s not "
+                        "repeated.\n",
+                        field->name().c_str());
             errorCount++;
+            continue;
         }
     }
-    // Binary field validity is already checked above.
-    if (isBinaryField) {
-        atomDecl->binaryFields.push_back(it->first);
+
+    // Check that if there's an attribution chain, it's at position 1.
+    for (map<int, const FieldDescriptor*>::const_iterator it = fields.begin(); it != fields.end();
+         it++) {
+        int number = it->first;
+        if (number != 1) {
+            const FieldDescriptor* field = it->second;
+            java_type_t javaType = java_type(field);
+            if (javaType == JAVA_TYPE_ATTRIBUTION_CHAIN) {
+                print_error(field,
+                            "AttributionChain fields must have field id 1, in message: '%s'\n",
+                            atom->name().c_str());
+                errorCount++;
+            }
+        }
+    }
+
+    // Build the type signature and the atom data.
+    for (map<int, const FieldDescriptor*>::const_iterator it = fields.begin(); it != fields.end();
+         it++) {
+        const FieldDescriptor* field = it->second;
+        java_type_t javaType = java_type(field);
+        bool isBinaryField = field->options().GetExtension(os::statsd::log_mode) ==
+                             os::statsd::LogMode::MODE_BYTES;
+
+        AtomField atField(field->name(), javaType);
+
+        if (javaType == JAVA_TYPE_ENUM) {
+            // All enums are treated as ints when it comes to function signatures.
+            collate_enums(*field->enum_type(), &atField);
+        }
+
+        // Generate signature for pushed atoms
+        if (atomDecl->code < PULL_ATOM_START_ID) {
+            if (javaType == JAVA_TYPE_ENUM) {
+                // All enums are treated as ints when it comes to function signatures.
+                signature->push_back(JAVA_TYPE_INT);
+            } else if (javaType == JAVA_TYPE_OBJECT && isBinaryField) {
+                signature->push_back(JAVA_TYPE_BYTE_ARRAY);
+            } else {
+                signature->push_back(javaType);
+            }
+        }
+
+        atomDecl->fields.push_back(atField);
+
+        errorCount += collate_field_annotations(atomDecl, field, it->first, javaType);
     }
-  }
 
-  return errorCount;
+    return errorCount;
 }
 
-// This function flattens the fields of the AttributionNode proto in an Atom proto and generates
-// the corresponding atom decl and signature.
-bool get_non_chained_node(const Descriptor *atom, AtomDecl *atomDecl,
-                          vector<java_type_t> *signature) {
+// This function flattens the fields of the AttributionNode proto in an Atom
+// proto and generates the corresponding atom decl and signature.
+bool get_non_chained_node(const Descriptor* atom, AtomDecl* atomDecl,
+                          vector<java_type_t>* signature) {
     // Build a sorted list of the fields. Descriptor has them in source file
     // order.
-    map<int, const FieldDescriptor *> fields;
+    map<int, const FieldDescriptor*> fields;
     for (int j = 0; j < atom->field_count(); j++) {
-        const FieldDescriptor *field = atom->field(j);
+        const FieldDescriptor* field = atom->field(j);
         fields[field->number()] = field;
     }
 
     AtomDecl attributionDecl;
     vector<java_type_t> attributionSignature;
-    collate_atom(android::os::statsd::AttributionNode::descriptor(),
-                 &attributionDecl, &attributionSignature);
+    collate_atom(android::os::statsd::AttributionNode::descriptor(), &attributionDecl,
+                 &attributionSignature);
 
     // Build the type signature and the atom data.
     bool has_attribution_node = false;
-    for (map<int, const FieldDescriptor *>::const_iterator it = fields.begin();
-        it != fields.end(); it++) {
-        const FieldDescriptor *field = it->second;
+    for (map<int, const FieldDescriptor*>::const_iterator it = fields.begin(); it != fields.end();
+         it++) {
+        const FieldDescriptor* field = it->second;
         java_type_t javaType = java_type(field);
         if (javaType == JAVA_TYPE_ATTRIBUTION_CHAIN) {
-            atomDecl->fields.insert(
-                atomDecl->fields.end(),
-                attributionDecl.fields.begin(), attributionDecl.fields.end());
-            signature->insert(
-                signature->end(),
-                attributionSignature.begin(), attributionSignature.end());
+            atomDecl->fields.insert(atomDecl->fields.end(), attributionDecl.fields.begin(),
+                                    attributionDecl.fields.end());
+            signature->insert(signature->end(), attributionSignature.begin(),
+                              attributionSignature.end());
             has_attribution_node = true;
 
         } else {
@@ -372,106 +461,115 @@ bool get_non_chained_node(const Descriptor *atom, AtomDecl *atomDecl,
     return has_attribution_node;
 }
 
+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;
+        (*fieldNumberToAtomDeclSet)[fieldNumber].insert(atomDecl);
+    }
+}
+
 /**
  * Gather the info about the atoms.
  */
-int collate_atoms(const Descriptor *descriptor, Atoms *atoms) {
-  int errorCount = 0;
-  const bool dbg = false;
+int collate_atoms(const Descriptor* descriptor, const string& moduleName, Atoms* atoms) {
+    int errorCount = 0;
+
+    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;
+                }
+            }
 
-  int maxPushedAtomId = 2;
-  for (int i = 0; i < descriptor->field_count(); i++) {
-    const FieldDescriptor *atomField = descriptor->field(i);
+            // 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());
-    }
+        if (dbg) {
+            printf("   %s (%d)\n", atomField->name().c_str(), atomField->number());
+        }
 
-    // StatsEvent only has one oneof, which contains only messages. Don't allow
-    // other types.
-    if (atomField->type() != FieldDescriptor::TYPE_MESSAGE) {
-      print_error(atomField,
-                  "Bad type for atom. StatsEvent can only have message type "
-                  "fields: %s\n",
-                  atomField->name().c_str());
-      errorCount++;
-      continue;
-    }
+        // StatsEvent only has one oneof, which contains only messages. Don't allow
+        // other types.
+        if (atomField->type() != FieldDescriptor::TYPE_MESSAGE) {
+            print_error(atomField,
+                        "Bad type for atom. StatsEvent can only have message type "
+                        "fields: %s\n",
+                        atomField->name().c_str());
+            errorCount++;
+            continue;
+        }
 
-    const Descriptor *atom = atomField->message_type();
-    AtomDecl atomDecl(atomField->number(), atomField->name(), atom->name());
+        const Descriptor* atom = atomField->message_type();
+        shared_ptr<AtomDecl> atomDecl =
+                make_shared<AtomDecl>(atomField->number(), atomField->name(), atom->name());
+
+        if (atomDecl->code < PULL_ATOM_START_ID &&
+            atomField->options().GetExtension(os::statsd::truncate_timestamp)) {
+            addAnnotationToAtomDecl(atomDecl.get(), ATOM_ID_FIELD_NUMBER,
+                                    ANNOTATION_ID_TRUNCATE_TIMESTAMP, ANNOTATION_TYPE_BOOL,
+                                    AnnotationValue(true));
+            if (dbg) {
+                printf("%s can have timestamp truncated\n", atomField->name().c_str());
+            }
+        }
 
-    if (atomField->options().GetExtension(os::statsd::allow_from_any_uid) == true) {
-        atomDecl.whitelisted = true;
-    }
+        vector<java_type_t> signature;
+        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;
+        }
 
-    if (atomField->options().HasExtension(os::statsd::log_from_module)) {
-        atomDecl.hasModule = true;
-        atomDecl.moduleName = atomField->options().GetExtension(os::statsd::log_from_module);
-    }
+        FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet = atoms->signatureInfoMap[signature];
+        populateFieldNumberToAtomDeclSet(atomDecl, &fieldNumberToAtomDeclSet);
 
-    vector<java_type_t> signature;
-    errorCount += collate_atom(atom, &atomDecl, &signature);
-    if (atomDecl.primaryFields.size() != 0 && atomDecl.exclusiveField == 0) {
-        errorCount++;
-    }
+        atoms->decls.insert(atomDecl);
 
-    // Add the signature if does not already exist.
-    auto signature_to_modules_it = atoms->signatures_to_modules.find(signature);
-    if (signature_to_modules_it == atoms->signatures_to_modules.end()) {
-        set<string> modules;
-        if (atomDecl.hasModule) {
-            modules.insert(atomDecl.moduleName);
-        }
-        atoms->signatures_to_modules[signature] = modules;
-    } else {
-        if (atomDecl.hasModule) {
-            signature_to_modules_it->second.insert(atomDecl.moduleName);
+        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.get(), &nonChainedSignature)) {
+            FieldNumberToAtomDeclSet& nonChainedFieldNumberToAtomDeclSet =
+                    atoms->nonChainedSignatureInfoMap[nonChainedSignature];
+            populateFieldNumberToAtomDeclSet(nonChainedAtomDecl,
+                                             &nonChainedFieldNumberToAtomDeclSet);
+
+            atoms->non_chained_decls.insert(nonChainedAtomDecl);
         }
     }
-    atoms->decls.insert(atomDecl);
-
-    AtomDecl nonChainedAtomDecl(atomField->number(), atomField->name(), atom->name());
-    vector<java_type_t> nonChainedSignature;
-    if (get_non_chained_node(atom, &nonChainedAtomDecl, &nonChainedSignature)) {
-        auto it = atoms->non_chained_signatures_to_modules.find(nonChainedSignature);
-        if (it == atoms->non_chained_signatures_to_modules.end()) {
-            set<string> modules_non_chained;
-            if (atomDecl.hasModule) {
-                modules_non_chained.insert(atomDecl.moduleName);
-            }
-            atoms->non_chained_signatures_to_modules[nonChainedSignature] = modules_non_chained;
-        } else {
-            if (atomDecl.hasModule) {
-                it->second.insert(atomDecl.moduleName);
+
+    if (dbg) {
+        printf("signatures = [\n");
+        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();
+                 jt++) {
+                printf(" %d", (int)*jt);
             }
+            printf("\n");
         }
-        atoms->non_chained_decls.insert(nonChainedAtomDecl);
-    }
-
-    if (atomDecl.code < PULL_ATOM_START_ID && atomDecl.code > maxPushedAtomId) {
-        maxPushedAtomId = atomDecl.code;
-    }
-  }
-
-  atoms->maxPushedAtomId = maxPushedAtomId;
-
-  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++) {
-      printf("   ");
-      for (vector<java_type_t>::const_iterator jt = it->first.begin();
-           jt != it->first.end(); jt++) {
-        printf(" %d", (int)*jt);
-      }
-      printf("\n");
+        printf("]\n");
     }
-    printf("]\n");
-  }
 
-  return errorCount;
+    return errorCount;
 }
 
 }  // namespace stats_log_api_gen