Wean statsd off batterystats: workchain (cp)

Modified cherry pick from master:
(cherry picked from commit d888df2ba99e117b50a5553e18426ad370294a20)
covering ag/4210168 and ag/4230692.

The tool that generates StatsLogInternal.java now includes the following
type of code, to accomodate WorkSources automatically.

    public static void write(int code, WorkSource ws, int arg2) {
        for (int i = 0; i < ws.size(); ++i) {
            write_non_chained(code, ws.get(i), ws.getName(i), arg2);
        }
        ArrayList<WorkSource.WorkChain> workChains = ws.getWorkChains();
        if (workChains != null) {
            for (WorkSource.WorkChain wc : workChains) {
                write(code, wc.getUids(), wc.getTags(), arg2);
            }
        }
    }

This will reduce a lot of code repitition for the StatsLog.write calls
in BatteryStats. Consequently, doing so will make it easier to remove
these StatsLog.write calls out of BatteryStats.

Bug: 80308558
Bug: 80314737
Test: manual verification
Test: com.android.internal.os.BatteryStatsTests
Change-Id: Ied537a6195121f16ef49b6bccf58c8c31964909c
Merged-In: Ibd28297345f4ab926dec01a89de0f1bfff9f8871
diff --git a/tools/stats_log_api_gen/main.cpp b/tools/stats_log_api_gen/main.cpp
index 638549d..1f13cc2 100644
--- a/tools/stats_log_api_gen/main.cpp
+++ b/tools/stats_log_api_gen/main.cpp
@@ -614,23 +614,19 @@
     return 0;
 }
 
-static void write_java_usage(
-    FILE* out, const string& method_name, const string& atom_code_name,
-    const AtomDecl& atom, const AtomDecl &attributionDecl) {
+static void write_java_usage(FILE* out, const string& method_name, const string& atom_code_name,
+        const AtomDecl& atom) {
     fprintf(out, "     * Usage: StatsLog.%s(StatsLog.%s",
         method_name.c_str(), atom_code_name.c_str());
     for (vector<AtomField>::const_iterator field = atom.fields.begin();
         field != atom.fields.end(); field++) {
         if (field->javaType == JAVA_TYPE_ATTRIBUTION_CHAIN) {
-            for (auto chainField : attributionDecl.fields) {
-                fprintf(out, ", %s[] %s",
-                    java_type_name(chainField.javaType), chainField.name.c_str());
-            }
+            fprintf(out, ", android.os.WorkSource workSource");
         } else {
             fprintf(out, ", %s %s", java_type_name(field->javaType), field->name.c_str());
         }
     }
-    fprintf(out, ");\n");
+    fprintf(out, ");<br>\n");
 }
 
 static void write_java_method(
@@ -656,6 +652,75 @@
     }
 }
 
+static void write_java_work_source_method(FILE* out, const set<vector<java_type_t>>& signatures) {
+    fprintf(out, "\n    // WorkSource methods.\n");
+    for (set<vector<java_type_t>>::const_iterator signature = signatures.begin();
+            signature != signatures.end(); signature++) {
+        // Determine if there is Attribution in this signature.
+        int attributionArg = -1;
+        int argIndexMax = 0;
+        for (vector<java_type_t>::const_iterator arg = signature->begin();
+                arg != signature->end(); arg++) {
+            argIndexMax++;
+            if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
+                if (attributionArg > -1) {
+                    fprintf(stderr, "An atom contains multiple AttributionNode fields.\n");
+                    fprintf(stderr, "This is not supported. Aborting WorkSource method writing.\n");
+                    fprintf(out, "\n// Invalid for WorkSource: more than one attribution chain.\n");
+                    return;
+                }
+                attributionArg = argIndexMax;
+            }
+        }
+        if (attributionArg < 0) {
+            continue;
+        }
+
+        // Method header (signature)
+        fprintf(out, "    public static void write(int code");
+        int argIndex = 1;
+        for (vector<java_type_t>::const_iterator arg = signature->begin();
+                arg != signature->end(); arg++) {
+            if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
+                fprintf(out, ", WorkSource ws");
+            } else {
+                fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
+            }
+            argIndex++;
+        }
+        fprintf(out, ") {\n");
+
+        // write_non_chained() component. TODO: Remove when flat uids are no longer needed.
+        fprintf(out, "        for (int i = 0; i < ws.size(); ++i) {\n");
+        fprintf(out, "            write_non_chained(code");
+        for (int argIndex = 1; argIndex <= argIndexMax; argIndex++) {
+            if (argIndex == attributionArg) {
+                fprintf(out, ", ws.get(i), ws.getName(i)");
+            } else {
+               fprintf(out, ", arg%d", argIndex);
+            }
+        }
+        fprintf(out, ");\n");
+        fprintf(out, "        }\n"); // close flor-loop
+
+        // write() component.
+        fprintf(out, "        ArrayList<WorkSource.WorkChain> workChains = ws.getWorkChains();\n");
+        fprintf(out, "        if (workChains != null) {\n");
+        fprintf(out, "            for (WorkSource.WorkChain wc : workChains) {\n");
+        fprintf(out, "                write(code");
+        for (int argIndex = 1; argIndex <= argIndexMax; argIndex++) {
+            if (argIndex == attributionArg) {
+                fprintf(out, ", wc.getUids(), wc.getTags()");
+            } else {
+               fprintf(out, ", arg%d", argIndex);
+            }
+        }
+        fprintf(out, ");\n");
+        fprintf(out, "            }\n"); // close for-loop
+        fprintf(out, "        }\n"); // close if
+        fprintf(out, "    }\n"); // close method
+    }
+}
 
 static int
 write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl)
@@ -665,6 +730,9 @@
     fprintf(out, "\n");
     fprintf(out, "package android.util;\n");
     fprintf(out, "\n");
+    fprintf(out, "import android.os.WorkSource;\n");
+    fprintf(out, "import java.util.ArrayList;\n");
+    fprintf(out, "\n");
     fprintf(out, "\n");
     fprintf(out, "/**\n");
     fprintf(out, " * API For logging statistics events.\n");
@@ -682,12 +750,11 @@
         string constant = make_constant_name(atom->name);
         fprintf(out, "\n");
         fprintf(out, "    /**\n");
-        fprintf(out, "     * %s %s\n", atom->message.c_str(), atom->name.c_str());
-        write_java_usage(out, "write", constant, *atom, attributionDecl);
+        fprintf(out, "     * %s %s<br>\n", atom->message.c_str(), atom->name.c_str());
+        write_java_usage(out, "write", constant, *atom);
         auto non_chained_decl = atom_code_to_non_chained_decl_map.find(atom->code);
         if (non_chained_decl != atom_code_to_non_chained_decl_map.end()) {
-            write_java_usage(out, "write_non_chained", constant, *non_chained_decl->second,
-             attributionDecl);
+            write_java_usage(out, "write_non_chained", constant, *non_chained_decl->second);
         }
         fprintf(out, "     */\n");
         fprintf(out, "    public static final int %s = %d;\n", constant.c_str(), atom->code);
@@ -720,6 +787,7 @@
     fprintf(out, "    // Write methods\n");
     write_java_method(out, "write", atoms.signatures, attributionDecl);
     write_java_method(out, "write_non_chained", atoms.non_chained_signatures, attributionDecl);
+    write_java_work_source_method(out, atoms.signatures);
 
     fprintf(out, "}\n");