Merge "Added support for pulled atoms for stats-log-api-gen"

6 files changed, 338 insertions, 174 deletions

diff --git a/cmds/statsd/src/atoms.proto b/cmds/statsd/src/atoms.proto
index 8c54c242e04a..02c0d933b3cd 100644
--- a/cmds/statsd/src/atoms.proto
+++ b/cmds/statsd/src/atoms.proto
@@ -4590,7 +4590,7 @@ message GeneralExternalStorageAccessStats {
     // Includes file path and ContentResolver accesses
     optional uint32 secondary_storage_accesses = 4;
     // Comma-separated list of mime types that were accessed.
-    optional MimeTypes mime_types_accessed = 5;
+    optional MimeTypes mime_types_accessed = 5 [(log_mode) = MODE_BYTES];
 }
 
 /**
@@ -6174,7 +6174,7 @@ message ProcessStatsAvailablePagesProto {
  * Pulled from ProcessStatsService.java
  */
 message ProcStats {
-    optional ProcessStatsSectionProto proc_stats_section = 1;
+    optional ProcessStatsSectionProto proc_stats_section = 1 [(log_mode) = MODE_BYTES];
     // Data pulled from device into this is sometimes sharded across multiple atoms to work around
     // a size limit. When this happens, this shard ID will contain an increasing 1-indexed integer
     // with the number of this shard.
@@ -6185,7 +6185,7 @@ message ProcStats {
  * Pulled from ProcessStatsService.java
  */
 message ProcStatsPkgProc {
-    optional ProcessStatsSectionProto proc_stats_section = 1;
+    optional ProcessStatsSectionProto proc_stats_section = 1 [(log_mode) = MODE_BYTES];
 }
 
 // Next Tag: 2
@@ -6203,7 +6203,7 @@ message NotificationRemoteViewsProto {
  * Pulled from NotificationManagerService.java
  */
 message NotificationRemoteViews {
-    optional NotificationRemoteViewsProto notification_remote_views = 1;
+    optional NotificationRemoteViewsProto notification_remote_views = 1 [(log_mode) = MODE_BYTES];
 }
 
 /**
@@ -6271,7 +6271,7 @@ message DNDModeProto {
     // May also be "MANUAL_RULE" to indicate app-activation of the manual rule.
     optional string id = 5;
     optional int32 uid = 6 [(is_uid) = true]; // currently only SYSTEM_UID or 0 for other
-    optional DNDPolicyProto policy = 7;
+    optional DNDPolicyProto policy = 7 [(log_mode) = MODE_BYTES];
 }
 
 /**
@@ -6436,7 +6436,7 @@ message PowerProfileProto {
  * Pulled from PowerProfile.java
  */
 message PowerProfile {
-    optional PowerProfileProto power_profile = 1;
+    optional PowerProfileProto power_profile = 1 [(log_mode) = MODE_BYTES];
 }
 
 /**
@@ -8131,7 +8131,7 @@ message DeviceIdentifierAccessDenied {
 message TrainInfo {
     optional int64 train_version_code = 1;
 
-    optional TrainExperimentIds train_experiment_id = 2;
+    optional TrainExperimentIds train_experiment_id = 2 [(log_mode) = MODE_BYTES];
 
     optional string train_name = 3;
 
@@ -11179,8 +11179,8 @@ message BlobInfo {
     optional int64 expiry_timestamp_millis = 3;
 
     // List of committers of this Blob
-    optional BlobCommitterListProto committers = 4;
+    optional BlobCommitterListProto committers = 4 [(log_mode) = MODE_BYTES];
 
     // List of leasees of this Blob
-    optional BlobLeaseeListProto leasees = 5;
+    optional BlobLeaseeListProto leasees = 5 [(log_mode) = MODE_BYTES];
 }
diff --git a/tools/stats_log_api_gen/Collation.cpp b/tools/stats_log_api_gen/Collation.cpp
index a230de46dcf3..4c741c49cfdb 100644
--- a/tools/stats_log_api_gen/Collation.cpp
+++ b/tools/stats_log_api_gen/Collation.cpp
@@ -25,6 +25,7 @@
 namespace android {
 namespace stats_log_api_gen {
 
+using google::protobuf::OneofDescriptor;
 using google::protobuf::EnumDescriptor;
 using google::protobuf::FieldDescriptor;
 using google::protobuf::FileDescriptor;
@@ -396,16 +397,14 @@ int collate_atom(const Descriptor* atom, AtomDecl* atomDecl, vector<java_type_t>
             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);
-            }
+        // Generate signature for atom.
+        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);
@@ -518,8 +517,7 @@ int collate_atoms(const Descriptor* descriptor, const string& moduleName, Atoms*
         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)) {
+        if (atomField->options().GetExtension(os::statsd::truncate_timestamp)) {
             addAnnotationToAtomDecl(atomDecl.get(), ATOM_ID_FIELD_NUMBER,
                                     ANNOTATION_ID_TRUNCATE_TIMESTAMP, ANNOTATION_TYPE_BOOL,
                                     AnnotationValue(true));
@@ -537,7 +535,24 @@ int collate_atoms(const Descriptor* descriptor, const string& moduleName, Atoms*
             continue;
         }
 
-        FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet = atoms->signatureInfoMap[signature];
+        const OneofDescriptor* oneofAtom = atomField->containing_oneof();
+        if (oneofAtom == nullptr) {
+            print_error(atomField, "Atom is not declared in a `oneof` field: %s\n",
+                        atomField->name().c_str());
+            errorCount++;
+            continue;
+        }
+        else if ((oneofAtom->name() != ONEOF_PUSHED_ATOM_NAME) &&
+                 (oneofAtom->name() != ONEOF_PULLED_ATOM_NAME)) {
+            print_error(atomField, "Atom is neither a pushed nor pulled atom: %s\n",
+                        atomField->name().c_str());
+            errorCount++;
+            continue;
+        }
+
+        FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet = oneofAtom->name() ==
+            ONEOF_PUSHED_ATOM_NAME ? atoms->signatureInfoMap[signature] :
+            atoms->pulledAtomsSignatureInfoMap[signature];
         populateFieldNumberToAtomDeclSet(atomDecl, &fieldNumberToAtomDeclSet);
 
         atoms->decls.insert(atomDecl);
@@ -556,6 +571,7 @@ int collate_atoms(const Descriptor* descriptor, const string& moduleName, Atoms*
     }
 
     if (dbg) {
+        // Signatures for pushed atoms.
         printf("signatures = [\n");
         for (SignatureInfoMap::const_iterator it = atoms->signatureInfoMap.begin();
              it != atoms->signatureInfoMap.end(); it++) {
@@ -566,6 +582,17 @@ int collate_atoms(const Descriptor* descriptor, const string& moduleName, Atoms*
             }
             printf("\n");
         }
+
+        // Signatures for pull atoms.
+        for (SignatureInfoMap::const_iterator it = atoms->pulledAtomsSignatureInfoMap.begin();
+             it != atoms->pulledAtomsSignatureInfoMap.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");
     }
 
diff --git a/tools/stats_log_api_gen/Collation.h b/tools/stats_log_api_gen/Collation.h
index 10b34ecf5f54..e637ed945a08 100644
--- a/tools/stats_log_api_gen/Collation.h
+++ b/tools/stats_log_api_gen/Collation.h
@@ -29,6 +29,7 @@
 namespace android {
 namespace stats_log_api_gen {
 
+using google::protobuf::OneofDescriptor;
 using google::protobuf::Descriptor;
 using google::protobuf::FieldDescriptor;
 using std::map;
@@ -41,6 +42,14 @@ const int PULL_ATOM_START_ID = 10000;
 
 const int FIRST_UID_IN_CHAIN_ID = 0;
 
+/**
+ * The types of oneof atoms.
+ *
+ * `OneofDescriptor::name()` returns the name of the oneof.
+ */
+const string ONEOF_PUSHED_ATOM_NAME = "pushed";
+const string ONEOF_PULLED_ATOM_NAME = "pulled";
+
 enum AnnotationId : uint8_t {
     ANNOTATION_ID_IS_UID = 1,
     ANNOTATION_ID_TRUNCATE_TIMESTAMP = 2,
@@ -184,6 +193,7 @@ using SignatureInfoMap = map<vector<java_type_t>, FieldNumberToAtomDeclSet>;
 
 struct Atoms {
     SignatureInfoMap signatureInfoMap;
+    SignatureInfoMap pulledAtomsSignatureInfoMap;
     AtomDeclSet decls;
     AtomDeclSet non_chained_decls;
     SignatureInfoMap nonChainedSignatureInfoMap;
diff --git a/tools/stats_log_api_gen/java_writer.cpp b/tools/stats_log_api_gen/java_writer.cpp
index f4c937c3f599..ffbe9f800736 100644
--- a/tools/stats_log_api_gen/java_writer.cpp
+++ b/tools/stats_log_api_gen/java_writer.cpp
@@ -96,29 +96,160 @@ static void write_annotations(FILE* out, int argIndex,
     }
 }
 
+static void write_method_signature(FILE* out, const vector<java_type_t>& signature,
+                                   const AtomDecl& attributionDecl) {
+    int argIndex = 1;
+    for (vector<java_type_t>::const_iterator arg = signature.begin(); arg != signature.end();
+         arg++) {
+        if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
+            for (auto chainField : attributionDecl.fields) {
+                fprintf(out, ", %s[] %s", java_type_name(chainField.javaType),
+                        chainField.name.c_str());
+            }
+        } else if (*arg == JAVA_TYPE_KEY_VALUE_PAIR) {
+            fprintf(out, ", android.util.SparseArray<Object> valueMap");
+        } else {
+            fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
+        }
+        argIndex++;
+    }
+}
+
+static int write_method_body(FILE* out, const vector<java_type_t>& signature,
+                             const FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet,
+                             const AtomDecl& attributionDecl, const string& indent) {
+    // Start StatsEvent.Builder.
+    fprintf(out,
+            "%s        final StatsEvent.Builder builder = "
+            "StatsEvent.newBuilder();\n",
+            indent.c_str());
+
+    // Write atom code.
+    fprintf(out, "%s        builder.setAtomId(code);\n", indent.c_str());
+    write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAtomDeclSet);
+
+    // Write the args.
+    int argIndex = 1;
+    for (vector<java_type_t>::const_iterator arg = signature.begin(); arg != signature.end();
+         arg++) {
+        switch (*arg) {
+            case JAVA_TYPE_BOOLEAN:
+                fprintf(out, "%s        builder.writeBoolean(arg%d);\n", indent.c_str(),
+                        argIndex);
+                break;
+            case JAVA_TYPE_INT:
+            case JAVA_TYPE_ENUM:
+                fprintf(out, "%s        builder.writeInt(arg%d);\n", indent.c_str(), argIndex);
+                break;
+            case JAVA_TYPE_FLOAT:
+                fprintf(out, "%s        builder.writeFloat(arg%d);\n", indent.c_str(),
+                        argIndex);
+                break;
+            case JAVA_TYPE_LONG:
+                fprintf(out, "%s        builder.writeLong(arg%d);\n", indent.c_str(), argIndex);
+                break;
+            case JAVA_TYPE_STRING:
+                fprintf(out, "%s        builder.writeString(arg%d);\n", indent.c_str(),
+                        argIndex);
+                break;
+            case JAVA_TYPE_BYTE_ARRAY:
+                fprintf(out,
+                        "%s        builder.writeByteArray(null == arg%d ? new byte[0] : "
+                        "arg%d);\n",
+                        indent.c_str(), argIndex, argIndex);
+                break;
+            case JAVA_TYPE_ATTRIBUTION_CHAIN: {
+                const char* uidName = attributionDecl.fields.front().name.c_str();
+                const char* tagName = attributionDecl.fields.back().name.c_str();
+
+                fprintf(out, "%s        builder.writeAttributionChain(\n", indent.c_str());
+                fprintf(out, "%s                null == %s ? new int[0] : %s,\n",
+                        indent.c_str(), uidName, uidName);
+                fprintf(out, "%s                null == %s ? new String[0] : %s);\n",
+                        indent.c_str(), tagName, tagName);
+                break;
+            }
+            case JAVA_TYPE_KEY_VALUE_PAIR:
+                fprintf(out, "\n");
+                fprintf(out, "%s        // Write KeyValuePairs.\n", indent.c_str());
+                fprintf(out, "%s        final int count = valueMap.size();\n", indent.c_str());
+                fprintf(out, "%s        android.util.SparseIntArray intMap = null;\n",
+                        indent.c_str());
+                fprintf(out, "%s        android.util.SparseLongArray longMap = null;\n",
+                        indent.c_str());
+                fprintf(out, "%s        android.util.SparseArray<String> stringMap = null;\n",
+                        indent.c_str());
+                fprintf(out, "%s        android.util.SparseArray<Float> floatMap = null;\n",
+                        indent.c_str());
+                fprintf(out, "%s        for (int i = 0; i < count; i++) {\n", indent.c_str());
+                fprintf(out, "%s            final int key = valueMap.keyAt(i);\n",
+                        indent.c_str());
+                fprintf(out, "%s            final Object value = valueMap.valueAt(i);\n",
+                        indent.c_str());
+                fprintf(out, "%s            if (value instanceof Integer) {\n", indent.c_str());
+                fprintf(out, "%s                if (null == intMap) {\n", indent.c_str());
+                fprintf(out,
+                        "%s                    intMap = new "
+                        "android.util.SparseIntArray();\n",
+                        indent.c_str());
+                fprintf(out, "%s                }\n", indent.c_str());
+                fprintf(out, "%s                intMap.put(key, (Integer) value);\n",
+                        indent.c_str());
+                fprintf(out, "%s            } else if (value instanceof Long) {\n",
+                        indent.c_str());
+                fprintf(out, "%s                if (null == longMap) {\n", indent.c_str());
+                fprintf(out,
+                        "%s                    longMap = new "
+                        "android.util.SparseLongArray();\n",
+                        indent.c_str());
+                fprintf(out, "%s                }\n", indent.c_str());
+                fprintf(out, "%s                longMap.put(key, (Long) value);\n",
+                        indent.c_str());
+                fprintf(out, "%s            } else if (value instanceof String) {\n",
+                        indent.c_str());
+                fprintf(out, "%s                if (null == stringMap) {\n", indent.c_str());
+                fprintf(out,
+                        "%s                    stringMap = new "
+                        "android.util.SparseArray<>();\n",
+                        indent.c_str());
+                fprintf(out, "%s                }\n", indent.c_str());
+                fprintf(out, "%s                stringMap.put(key, (String) value);\n",
+                        indent.c_str());
+                fprintf(out, "%s            } else if (value instanceof Float) {\n",
+                        indent.c_str());
+                fprintf(out, "%s                if (null == floatMap) {\n", indent.c_str());
+                fprintf(out,
+                        "%s                    floatMap = new "
+                        "android.util.SparseArray<>();\n",
+                        indent.c_str());
+                fprintf(out, "%s                }\n", indent.c_str());
+                fprintf(out, "%s                floatMap.put(key, (Float) value);\n",
+                        indent.c_str());
+                fprintf(out, "%s            }\n", indent.c_str());
+                fprintf(out, "%s        }\n", indent.c_str());
+                fprintf(out,
+                        "%s        builder.writeKeyValuePairs("
+                        "intMap, longMap, stringMap, floatMap);\n",
+                        indent.c_str());
+                break;
+            default:
+                // Unsupported types: OBJECT, DOUBLE.
+                fprintf(stderr, "Encountered unsupported type.");
+                return 1;
+        }
+        write_annotations(out, argIndex, fieldNumberToAtomDeclSet);
+        argIndex++;
+    }
+    return 0;
+}
+
 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 FieldNumberToAtomDeclSet& fieldNumberToAtomDeclSet = signatureInfoMapIt->second;
-        int argIndex = 1;
-        for (vector<java_type_t>::const_iterator arg = signature.begin(); arg != signature.end();
-             arg++) {
-            if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
-                for (auto chainField : attributionDecl.fields) {
-                    fprintf(out, ", %s[] %s", java_type_name(chainField.javaType),
-                            chainField.name.c_str());
-                }
-            } else if (*arg == JAVA_TYPE_KEY_VALUE_PAIR) {
-                fprintf(out, ", android.util.SparseArray<Object> valueMap");
-            } else {
-                fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
-            }
-            argIndex++;
-        }
+        write_method_signature(out, signatureInfoMapIt->first, attributionDecl);
         fprintf(out, ") {\n");
 
         // Print method body.
@@ -128,130 +259,13 @@ static int write_java_methods(FILE* out, const SignatureInfoMap& signatureInfoMa
             indent = "    ";
         }
 
-        // Start StatsEvent.Builder.
-        fprintf(out,
-                "%s        final StatsEvent.Builder builder = "
-                "StatsEvent.newBuilder();\n",
-                indent.c_str());
-
-        // Write atom code.
-        fprintf(out, "%s        builder.setAtomId(code);\n", indent.c_str());
-        write_annotations(out, ATOM_ID_FIELD_NUMBER, fieldNumberToAtomDeclSet);
-
-        // Write the args.
-        argIndex = 1;
-        for (vector<java_type_t>::const_iterator arg = signature.begin(); arg != signature.end();
-             arg++) {
-            switch (*arg) {
-                case JAVA_TYPE_BOOLEAN:
-                    fprintf(out, "%s        builder.writeBoolean(arg%d);\n", indent.c_str(),
-                            argIndex);
-                    break;
-                case JAVA_TYPE_INT:
-                case JAVA_TYPE_ENUM:
-                    fprintf(out, "%s        builder.writeInt(arg%d);\n", indent.c_str(), argIndex);
-                    break;
-                case JAVA_TYPE_FLOAT:
-                    fprintf(out, "%s        builder.writeFloat(arg%d);\n", indent.c_str(),
-                            argIndex);
-                    break;
-                case JAVA_TYPE_LONG:
-                    fprintf(out, "%s        builder.writeLong(arg%d);\n", indent.c_str(), argIndex);
-                    break;
-                case JAVA_TYPE_STRING:
-                    fprintf(out, "%s        builder.writeString(arg%d);\n", indent.c_str(),
-                            argIndex);
-                    break;
-                case JAVA_TYPE_BYTE_ARRAY:
-                    fprintf(out,
-                            "%s        builder.writeByteArray(null == arg%d ? new byte[0] : "
-                            "arg%d);\n",
-                            indent.c_str(), argIndex, argIndex);
-                    break;
-                case JAVA_TYPE_ATTRIBUTION_CHAIN: {
-                    const char* uidName = attributionDecl.fields.front().name.c_str();
-                    const char* tagName = attributionDecl.fields.back().name.c_str();
-
-                    fprintf(out, "%s        builder.writeAttributionChain(\n", indent.c_str());
-                    fprintf(out, "%s                null == %s ? new int[0] : %s,\n",
-                            indent.c_str(), uidName, uidName);
-                    fprintf(out, "%s                null == %s ? new String[0] : %s);\n",
-                            indent.c_str(), tagName, tagName);
-                    break;
-                }
-                case JAVA_TYPE_KEY_VALUE_PAIR:
-                    fprintf(out, "\n");
-                    fprintf(out, "%s        // Write KeyValuePairs.\n", indent.c_str());
-                    fprintf(out, "%s        final int count = valueMap.size();\n", indent.c_str());
-                    fprintf(out, "%s        android.util.SparseIntArray intMap = null;\n",
-                            indent.c_str());
-                    fprintf(out, "%s        android.util.SparseLongArray longMap = null;\n",
-                            indent.c_str());
-                    fprintf(out, "%s        android.util.SparseArray<String> stringMap = null;\n",
-                            indent.c_str());
-                    fprintf(out, "%s        android.util.SparseArray<Float> floatMap = null;\n",
-                            indent.c_str());
-                    fprintf(out, "%s        for (int i = 0; i < count; i++) {\n", indent.c_str());
-                    fprintf(out, "%s            final int key = valueMap.keyAt(i);\n",
-                            indent.c_str());
-                    fprintf(out, "%s            final Object value = valueMap.valueAt(i);\n",
-                            indent.c_str());
-                    fprintf(out, "%s            if (value instanceof Integer) {\n", indent.c_str());
-                    fprintf(out, "%s                if (null == intMap) {\n", indent.c_str());
-                    fprintf(out,
-                            "%s                    intMap = new "
-                            "android.util.SparseIntArray();\n",
-                            indent.c_str());
-                    fprintf(out, "%s                }\n", indent.c_str());
-                    fprintf(out, "%s                intMap.put(key, (Integer) value);\n",
-                            indent.c_str());
-                    fprintf(out, "%s            } else if (value instanceof Long) {\n",
-                            indent.c_str());
-                    fprintf(out, "%s                if (null == longMap) {\n", indent.c_str());
-                    fprintf(out,
-                            "%s                    longMap = new "
-                            "android.util.SparseLongArray();\n",
-                            indent.c_str());
-                    fprintf(out, "%s                }\n", indent.c_str());
-                    fprintf(out, "%s                longMap.put(key, (Long) value);\n",
-                            indent.c_str());
-                    fprintf(out, "%s            } else if (value instanceof String) {\n",
-                            indent.c_str());
-                    fprintf(out, "%s                if (null == stringMap) {\n", indent.c_str());
-                    fprintf(out,
-                            "%s                    stringMap = new "
-                            "android.util.SparseArray<>();\n",
-                            indent.c_str());
-                    fprintf(out, "%s                }\n", indent.c_str());
-                    fprintf(out, "%s                stringMap.put(key, (String) value);\n",
-                            indent.c_str());
-                    fprintf(out, "%s            } else if (value instanceof Float) {\n",
-                            indent.c_str());
-                    fprintf(out, "%s                if (null == floatMap) {\n", indent.c_str());
-                    fprintf(out,
-                            "%s                    floatMap = new "
-                            "android.util.SparseArray<>();\n",
-                            indent.c_str());
-                    fprintf(out, "%s                }\n", indent.c_str());
-                    fprintf(out, "%s                floatMap.put(key, (Float) value);\n",
-                            indent.c_str());
-                    fprintf(out, "%s            }\n", indent.c_str());
-                    fprintf(out, "%s        }\n", indent.c_str());
-                    fprintf(out,
-                            "%s        builder.writeKeyValuePairs("
-                            "intMap, longMap, stringMap, floatMap);\n",
-                            indent.c_str());
-                    break;
-                default:
-                    // Unsupported types: OBJECT, DOUBLE.
-                    fprintf(stderr, "Encountered unsupported type.");
-                    return 1;
-            }
-            write_annotations(out, argIndex, fieldNumberToAtomDeclSet);
-            argIndex++;
+        int ret = write_method_body(out, signatureInfoMapIt->first, signatureInfoMapIt->second,
+                                    attributionDecl, indent);
+        if (ret != 0) {
+            return ret;
         }
-
         fprintf(out, "\n");
+
         fprintf(out, "%s        builder.usePooledBuffer();\n", indent.c_str());
         fprintf(out, "%s        StatsLog.write(builder.build());\n", indent.c_str());
 
@@ -259,9 +273,9 @@ static int write_java_methods(FILE* out, const SignatureInfoMap& signatureInfoMa
         if (supportQ) {
             fprintf(out, "        } else {\n");
             fprintf(out, "            QLogger.write(code");
-            argIndex = 1;
-            for (vector<java_type_t>::const_iterator arg = signature.begin();
-                 arg != signature.end(); arg++) {
+            int argIndex = 1;
+            for (vector<java_type_t>::const_iterator arg = signatureInfoMapIt->first.begin();
+                 arg != signatureInfoMapIt->first.end(); arg++) {
                 if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
                     const char* uidName = attributionDecl.fields.front().name.c_str();
                     const char* tagName = attributionDecl.fields.back().name.c_str();
@@ -285,6 +299,32 @@ static int write_java_methods(FILE* out, const SignatureInfoMap& signatureInfoMa
     return 0;
 }
 
+static int write_java_build_stats_event_methods(FILE* out, const SignatureInfoMap& signatureInfoMap,
+                              const AtomDecl& attributionDecl) {
+    for (auto signatureInfoMapIt = signatureInfoMap.begin();
+         signatureInfoMapIt != signatureInfoMap.end(); signatureInfoMapIt++) {
+        // Print method signature.
+        fprintf(out, "    public static StatsEvent buildStatsEvent(int code");
+        write_method_signature(out, signatureInfoMapIt->first, attributionDecl);
+        fprintf(out, ") {\n");
+
+        // Print method body.
+        string indent("");
+        int ret = write_method_body(out, signatureInfoMapIt->first, signatureInfoMapIt->second,
+                                    attributionDecl, indent);
+        if (ret != 0) {
+            return ret;
+        }
+        fprintf(out, "\n");
+
+        fprintf(out, "%s        return builder.build();\n", indent.c_str());
+
+        fprintf(out, "    }\n");  // method
+        fprintf(out, "\n");
+    }
+    return 0;
+}
+
 int write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl& attributionDecl,
                          const string& javaClass, const string& javaPackage, const bool supportQ,
                          const bool supportWorkSource) {
@@ -319,6 +359,8 @@ int write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl& attribut
     fprintf(out, "    // Write methods\n");
     errors += write_java_methods(out, atoms.signatureInfoMap, attributionDecl, supportQ);
     errors += write_java_non_chained_methods(out, atoms.nonChainedSignatureInfoMap);
+    errors += write_java_build_stats_event_methods(out, atoms.pulledAtomsSignatureInfoMap,
+                                                   attributionDecl);
     if (supportWorkSource) {
         errors += write_java_work_source_methods(out, atoms.signatureInfoMap);
     }
diff --git a/tools/stats_log_api_gen/test.proto b/tools/stats_log_api_gen/test.proto
index aaa488e44fee..e658b62b8daa 100644
--- a/tools/stats_log_api_gen/test.proto
+++ b/tools/stats_log_api_gen/test.proto
@@ -58,7 +58,7 @@ message AllTypesAtom {
 }
 
 message Event {
-    oneof event {
+    oneof pushed {
         OutOfOrderAtom out_of_order_atom = 2;
         IntAtom int_atom = 1;
         AnotherIntAtom another_int_atom = 3;
@@ -74,7 +74,7 @@ message BadTypesAtom {
 }
 
 message BadTypesEvent {
-    oneof event {
+    oneof pushed {
         BadTypesAtom bad_types_atom = 1;
     }
 }
@@ -84,7 +84,7 @@ message BadSkippedFieldSingleAtom {
 }
 
 message BadSkippedFieldSingle {
-    oneof event {
+    oneof pushed {
         BadSkippedFieldSingleAtom bad = 1;
     }
 }
@@ -96,7 +96,7 @@ message BadSkippedFieldMultipleAtom {
 }
 
 message BadSkippedFieldMultiple {
-    oneof event {
+    oneof pushed {
         BadSkippedFieldMultipleAtom bad = 1;
     }
 }
@@ -107,11 +107,11 @@ message BadAttributionNodePositionAtom {
 }
 
 message BadAttributionNodePosition {
-  oneof event { BadAttributionNodePositionAtom bad = 1; }
+  oneof pushed { BadAttributionNodePositionAtom bad = 1; }
 }
 
 message GoodEventWithBinaryFieldAtom {
-    oneof event { GoodBinaryFieldAtom field1 = 1; }
+    oneof pushed { GoodBinaryFieldAtom field1 = 1; }
 }
 
 message ComplexField {
@@ -124,7 +124,7 @@ message GoodBinaryFieldAtom {
 }
 
 message BadEventWithBinaryFieldAtom {
-    oneof event { BadBinaryFieldAtom field1 = 1; }
+    oneof pushed { BadBinaryFieldAtom field1 = 1; }
 }
 
 message BadBinaryFieldAtom {
@@ -133,7 +133,7 @@ message BadBinaryFieldAtom {
 }
 
 message BadStateAtoms {
-    oneof event {
+    oneof pushed {
         BadStateAtom1 bad1 = 1;
         BadStateAtom2 bad2 = 2;
         BadStateAtom3 bad3 = 3;
@@ -141,7 +141,7 @@ message BadStateAtoms {
 }
 
 message GoodStateAtoms {
-    oneof event {
+    oneof pushed {
         GoodStateAtom1 good1 = 1;
         GoodStateAtom2 good2 = 2;
     }
@@ -204,7 +204,7 @@ message NoModuleAtom {
 }
 
 message ModuleAtoms {
-    oneof event {
+    oneof pushed {
         ModuleOneAtom module_one_atom = 1 [(android.os.statsd.module) = "module1"];
         ModuleTwoAtom module_two_atom = 2 [(android.os.statsd.module) = "module2"];
         ModuleOneAndTwoAtom module_one_and_two_atom = 3 [
@@ -213,3 +213,24 @@ message ModuleAtoms {
         NoModuleAtom no_module_atom = 4;
     }
 }
+
+message NotAPushNorPullAtom {
+    oneof event {
+        IntAtom int_atom = 1;
+    }
+}
+
+message AtomNotInAOneof {
+    optional IntAtom int_atom = 1;
+}
+
+message PushedAndPulledAtoms {
+    oneof pushed {
+        IntAtom int_atom_1 = 1;
+    }
+
+    oneof pulled {
+        OutOfOrderAtom out_of_order_atom = 11;
+        AnotherIntAtom another_int_atom = 10;
+    }
+}
diff --git a/tools/stats_log_api_gen/test_collation.cpp b/tools/stats_log_api_gen/test_collation.cpp
index dbae58889333..5fd728a29c07 100644
--- a/tools/stats_log_api_gen/test_collation.cpp
+++ b/tools/stats_log_api_gen/test_collation.cpp
@@ -365,5 +365,69 @@ TEST(CollationTest, RecognizeModule1Atom) {
     EXPECT_TRUE(annotation->value.boolValue);
 }
 
+/**
+ * Test that an atom is not a pushed nor pulled atom.
+ */
+TEST(CollationTest, InvalidAtomType) {
+    Atoms atoms;
+    int errorCount = collate_atoms(NotAPushNorPullAtom::descriptor(), DEFAULT_MODULE_NAME, &atoms);
+
+    EXPECT_EQ(1, errorCount);
+}
+
+/**
+ * Test that an atom was not declared in a `oneof` field.
+ */
+TEST(CollationTest, AtomNotDeclaredInAOneof) {
+    Atoms atoms;
+    int errorCount = collate_atoms(AtomNotInAOneof::descriptor(), DEFAULT_MODULE_NAME, &atoms);
+
+    EXPECT_EQ(1, errorCount);
+}
+
+/**
+ * Test a correct collation with pushed and pulled atoms.
+ */
+TEST(CollationTest, CollatePushedAndPulledAtoms) {
+    Atoms atoms;
+    int errorCount = collate_atoms(PushedAndPulledAtoms::descriptor(), DEFAULT_MODULE_NAME, &atoms);
+
+    EXPECT_EQ(0, errorCount);
+    EXPECT_EQ(1ul, atoms.signatureInfoMap.size());
+    EXPECT_EQ(2ul, atoms.pulledAtomsSignatureInfoMap.size());
+
+    // IntAtom
+    EXPECT_MAP_CONTAINS_SIGNATURE(atoms.signatureInfoMap, JAVA_TYPE_INT);
+
+    // AnotherIntAtom
+    EXPECT_MAP_CONTAINS_SIGNATURE(atoms.pulledAtomsSignatureInfoMap, JAVA_TYPE_INT);
+
+    // OutOfOrderAtom
+    EXPECT_MAP_CONTAINS_SIGNATURE(atoms.pulledAtomsSignatureInfoMap, JAVA_TYPE_INT, JAVA_TYPE_INT);
+
+    EXPECT_EQ(3ul, atoms.decls.size());
+
+    AtomDeclSet::const_iterator atomIt = atoms.decls.begin();
+    EXPECT_EQ(1, (*atomIt)->code);
+    EXPECT_EQ("int_atom_1", (*atomIt)->name);
+    EXPECT_EQ("IntAtom", (*atomIt)->message);
+    EXPECT_NO_ENUM_FIELD((*atomIt));
+    atomIt++;
+
+    EXPECT_EQ(10, (*atomIt)->code);
+    EXPECT_EQ("another_int_atom", (*atomIt)->name);
+    EXPECT_EQ("AnotherIntAtom", (*atomIt)->message);
+    EXPECT_NO_ENUM_FIELD((*atomIt));
+    atomIt++;
+
+    EXPECT_EQ(11, (*atomIt)->code);
+    EXPECT_EQ("out_of_order_atom", (*atomIt)->name);
+    EXPECT_EQ("OutOfOrderAtom", (*atomIt)->message);
+    EXPECT_NO_ENUM_FIELD((*atomIt));
+    atomIt++;
+
+    EXPECT_EQ(atoms.decls.end(), atomIt);
+}
+
 }  // namespace stats_log_api_gen
 }  // namespace android