Add flag to stats-log-api-gen for WorkSource.

As part of getting rid of StatsLogInternal and moving every atom to be
logged from a module, this CL addresses 2 issues:
1. Add WorkSource logging to modules.
2. Gate WorkSource logging via a CLI flag to prevent WorkSource methods
from being included in modules that are Mainline modules and can't have
@hide calls.

Bug: 147541828
Test: m
Test: atest android.cts.statsd.atom.UidAtomTests

Change-Id: Icce2f6ed7298be5e5ceebfb05d87e6fa7cd6bdbb
Merged-In: Icce2f6ed7298be5e5ceebfb05d87e6fa7cd6bdbb
diff --git a/Android.bp b/Android.bp
index 3046b20..3bd97a4 100644
--- a/Android.bp
+++ b/Android.bp
@@ -552,7 +552,7 @@
 genrule {
     name: "framework-statslog-gen",
     tools: ["stats-log-api-gen"],
-    cmd: "$(location stats-log-api-gen) --java $(out)",
+    cmd: "$(location stats-log-api-gen) --java $(out) --worksource",
     out: ["android/util/StatsLogInternal.java"],
 }
 
diff --git a/tools/stats_log_api_gen/Collation.cpp b/tools/stats_log_api_gen/Collation.cpp
index fa55601..75deb01 100644
--- a/tools/stats_log_api_gen/Collation.cpp
+++ b/tools/stats_log_api_gen/Collation.cpp
@@ -434,7 +434,7 @@
     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(signature);
+        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) {
diff --git a/tools/stats_log_api_gen/java_writer.cpp b/tools/stats_log_api_gen/java_writer.cpp
index dd3f0bb..c29936b 100644
--- a/tools/stats_log_api_gen/java_writer.cpp
+++ b/tools/stats_log_api_gen/java_writer.cpp
@@ -73,7 +73,7 @@
                         java_type_name(chainField.javaType), chainField.name.c_str());
                 }
             } else if (*arg == JAVA_TYPE_KEY_VALUE_PAIR) {
-                fprintf(out, ", SparseArray<Object> valueMap");
+                fprintf(out, ", android.util.SparseArray<Object> valueMap");
             } else {
                 fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
             }
@@ -142,16 +142,16 @@
                 fprintf(out,
                         "%s        final int count = valueMap.size();\n", indent.c_str());
                 fprintf(out,
-                        "%s        SparseIntArray intMap = null;\n",
+                        "%s        android.util.SparseIntArray intMap = null;\n",
                         indent.c_str());
                 fprintf(out,
-                        "%s        SparseLongArray longMap = null;\n",
+                        "%s        android.util.SparseLongArray longMap = null;\n",
                         indent.c_str());
                 fprintf(out,
-                        "%s        SparseArray<String> stringMap = null;\n",
+                        "%s        android.util.SparseArray<String> stringMap = null;\n",
                         indent.c_str());
                 fprintf(out,
-                        "%s        SparseArray<Float> floatMap = null;\n",
+                        "%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());
@@ -165,7 +165,7 @@
                 fprintf(out,
                         "%s                if (null == intMap) {\n", indent.c_str());
                 fprintf(out,
-                        "%s                    intMap = new SparseIntArray();\n", indent.c_str());
+                        "%s                    intMap = new android.util.SparseIntArray();\n", indent.c_str());
                 fprintf(out,
                         "%s                }\n", indent.c_str());
                 fprintf(out,
@@ -175,7 +175,7 @@
                 fprintf(out,
                         "%s                if (null == longMap) {\n", indent.c_str());
                 fprintf(out,
-                        "%s                    longMap = new SparseLongArray();\n", indent.c_str());
+                        "%s                    longMap = new android.util.SparseLongArray();\n", indent.c_str());
                 fprintf(out,
                         "%s                }\n", indent.c_str());
                 fprintf(out,
@@ -185,7 +185,7 @@
                 fprintf(out,
                         "%s                if (null == stringMap) {\n", indent.c_str());
                 fprintf(out,
-                        "%s                    stringMap = new SparseArray<>();\n", indent.c_str());
+                        "%s                    stringMap = new android.util.SparseArray<>();\n", indent.c_str());
                 fprintf(out,
                         "%s                }\n", indent.c_str());
                 fprintf(out,
@@ -195,7 +195,7 @@
                 fprintf(out,
                         "%s                if (null == floatMap) {\n", indent.c_str());
                 fprintf(out,
-                        "%s                    floatMap = new SparseArray<>();\n", indent.c_str());
+                        "%s                    floatMap = new android.util.SparseArray<>();\n", indent.c_str());
                 fprintf(out,
                         "%s                }\n", indent.c_str());
                 fprintf(out,
@@ -253,7 +253,8 @@
 
 int write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
                                     const string& moduleName, const string& javaClass,
-                                    const string& javaPackage, const bool supportQ) {
+                                    const string& javaPackage, const bool supportQ,
+                                    const bool supportWorkSource) {
     // Print prelude
     fprintf(out, "// This file is autogenerated\n");
     fprintf(out, "\n");
@@ -265,25 +266,9 @@
         fprintf(out, "import android.os.SystemClock;\n");
     }
 
-    if (DEFAULT_MODULE_NAME == moduleName) {
-        // Mainline modules don't use WorkSource logging.
-        fprintf(out, "import android.os.WorkSource;\n");
-
-        // SparseArray is used for writing KeyValuePairs; not supported for Mainline modules.
-        fprintf(out, "import android.util.SparseArray;\n");
-        fprintf(out, "import android.util.SparseIntArray;\n");
-        fprintf(out, "import android.util.SparseLongArray;\n");
-    }
-
     fprintf(out, "import android.util.StatsEvent;\n");
     fprintf(out, "import android.util.StatsLog;\n");
 
-    if (DEFAULT_MODULE_NAME == moduleName) {
-        // List is used for WorkSource writing. Only needed for default module.
-        fprintf(out, "\n");
-        fprintf(out, "import java.util.ArrayList;\n");
-    }
-
     fprintf(out, "\n");
     fprintf(out, "\n");
     fprintf(out, "/**\n");
@@ -305,7 +290,7 @@
             out, atoms.signatures_to_modules, attributionDecl, moduleName, supportQ);
     errors += write_java_non_chained_methods(
             out, atoms.non_chained_signatures_to_modules, moduleName);
-    if (DEFAULT_MODULE_NAME == moduleName) {
+    if (supportWorkSource) {
         errors += write_java_work_source_methods(out, atoms.signatures_to_modules, moduleName);
     }
 
diff --git a/tools/stats_log_api_gen/java_writer.h b/tools/stats_log_api_gen/java_writer.h
index 9324b23..5b78f05 100644
--- a/tools/stats_log_api_gen/java_writer.h
+++ b/tools/stats_log_api_gen/java_writer.h
@@ -31,8 +31,9 @@
 using namespace std;
 
 int write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
-                                    const string& moduleName, const string& javaClass,
-                                    const string& javaPackage, const bool supportQ);
+                         const string& moduleName, const string& javaClass,
+                         const string& javaPackage, const bool supportQ,
+                         const bool supportWorkSource);
 
 }  // namespace stats_log_api_gen
 }  // namespace android
diff --git a/tools/stats_log_api_gen/java_writer_q.cpp b/tools/stats_log_api_gen/java_writer_q.cpp
index d6899f6..f866129 100644
--- a/tools/stats_log_api_gen/java_writer_q.cpp
+++ b/tools/stats_log_api_gen/java_writer_q.cpp
@@ -382,7 +382,7 @@
                         java_type_name(chainField.javaType), chainField.name.c_str());
                 }
             } else if (*arg == JAVA_TYPE_KEY_VALUE_PAIR) {
-                fprintf(out, ", SparseArray<Object> value_map");
+                fprintf(out, ", android.util.SparseArray<Object> value_map");
             } else {
                 fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
             }
@@ -393,16 +393,13 @@
     }
 }
 
-int write_stats_log_java_q(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl) {
+int write_stats_log_java_q(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
+                           const bool supportWorkSource) {
     // Print prelude
     fprintf(out, "// This file is autogenerated\n");
     fprintf(out, "\n");
     fprintf(out, "package android.util;\n");
     fprintf(out, "\n");
-    fprintf(out, "import android.os.WorkSource;\n");
-    fprintf(out, "import android.util.SparseArray;\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");
@@ -418,16 +415,19 @@
     write_java_method(out, "write", atoms.signatures_to_modules, attributionDecl);
     write_java_method(out, "write_non_chained", atoms.non_chained_signatures_to_modules,
             attributionDecl);
-    write_java_work_source_methods(out, atoms.signatures_to_modules, DEFAULT_MODULE_NAME);
+    if (supportWorkSource) {
+        write_java_work_source_methods(out, atoms.signatures_to_modules, DEFAULT_MODULE_NAME);
+    }
 
     fprintf(out, "}\n");
 
     return 0;
 }
 
-int write_stats_log_java_q_for_module(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
-                                    const string& moduleName, const string& javaClass,
-                                    const string& javaPackage) {
+int write_stats_log_java_q_for_module(FILE* out, const Atoms& atoms,
+                                      const AtomDecl &attributionDecl, const string& moduleName,
+                                      const string& javaClass, const string& javaPackage,
+                                      const bool supportWorkSource) {
     // Print prelude
     fprintf(out, "// This file is autogenerated\n");
     fprintf(out, "\n");
@@ -438,8 +438,6 @@
     fprintf(out, "import android.util.StatsLog;\n");
     fprintf(out, "import android.os.SystemClock;\n");
     fprintf(out, "\n");
-    fprintf(out, "import java.util.ArrayList;\n");
-    fprintf(out, "\n");
     fprintf(out, "\n");
     fprintf(out, "/**\n");
     fprintf(out, " * Utility class for logging statistics events.\n");
@@ -459,6 +457,9 @@
             moduleName, "    ");
     errors += write_java_non_chained_methods(out, atoms.non_chained_signatures_to_modules,
             moduleName);
+    if (supportWorkSource) {
+        errors += write_java_work_source_methods(out, atoms.signatures_to_modules, moduleName);
+    }
 
     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 96ac745..36df1d8 100644
--- a/tools/stats_log_api_gen/java_writer_q.h
+++ b/tools/stats_log_api_gen/java_writer_q.h
@@ -46,11 +46,12 @@
         const string& indent);
 
 #if defined(STATS_SCHEMA_LEGACY)
-int write_stats_log_java_q(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl);
+int write_stats_log_java_q(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
+                           const bool supportWorkSource);
 
 int write_stats_log_java_q_for_module(FILE* out, const Atoms& atoms,
         const AtomDecl &attributionDecl, const string& moduleName, const string& javaClass,
-        const string& javaPackage);
+        const string& javaPackage, const bool supportWorkSource);
 #endif
 }  // namespace stats_log_api_gen
 }  // namespace android
diff --git a/tools/stats_log_api_gen/main.cpp b/tools/stats_log_api_gen/main.cpp
index 00a3704..6089532 100644
--- a/tools/stats_log_api_gen/main.cpp
+++ b/tools/stats_log_api_gen/main.cpp
@@ -512,6 +512,7 @@
     fprintf(stderr, "                       Optional for Java with module.\n");
     fprintf(stderr, "                       Default is \"StatsLogInternal\"\n");
     fprintf(stderr, "  --supportQ           Include support for Android Q.\n");
+    fprintf(stderr, "  --worksource         Include support for logging WorkSource objects.\n");
 }
 
 /**
@@ -534,6 +535,7 @@
     string javaPackage = DEFAULT_JAVA_PACKAGE;
     string javaClass = DEFAULT_JAVA_CLASS;
     bool supportQ = false;
+    bool supportWorkSource = false;
 
     int index = 1;
     while (index < argc) {
@@ -626,6 +628,8 @@
             atomsInfoCppHeaderImport = argv[index];
         } else if (0 == strcmp("--supportQ", argv[index])) {
             supportQ = true;
+        } else if (0 == strcmp("--worksource", argv[index])) {
+            supportWorkSource = true;
         }
 
         index++;
@@ -728,19 +732,15 @@
             fprintf(stderr, "Unable to open file for write: %s\n", javaFilename.c_str());
             return 1;
         }
-        // If this is for a specific module, the java package must also be provided.
-        if (moduleName != DEFAULT_MODULE_NAME && javaPackage== DEFAULT_JAVA_PACKAGE) {
-            fprintf(stderr, "Must supply --javaPackage if supplying a specific module\n");
-            return 1;
-        }
 
 #if defined(STATS_SCHEMA_LEGACY)
         if (moduleName == DEFAULT_MODULE_NAME) {
             errorCount = android::stats_log_api_gen::write_stats_log_java_q(
-                    out, atoms, attributionDecl);
+                    out, atoms, attributionDecl, supportWorkSource);
         } else {
             errorCount = android::stats_log_api_gen::write_stats_log_java_q_for_module(
-                    out, atoms, attributionDecl, moduleName, javaClass, javaPackage);
+                    out, atoms, attributionDecl, moduleName, javaClass, javaPackage,
+                    supportWorkSource);
 
         }
 #else
@@ -749,7 +749,8 @@
             javaPackage = "android.util";
         }
         errorCount = android::stats_log_api_gen::write_stats_log_java(
-                out, atoms, attributionDecl, moduleName, javaClass, javaPackage, supportQ);
+                out, atoms, attributionDecl, moduleName, javaClass, javaPackage, supportQ,
+                supportWorkSource);
 #endif
 
         fclose(out);
diff --git a/tools/stats_log_api_gen/utils.cpp b/tools/stats_log_api_gen/utils.cpp
index 5b830ee..6414042 100644
--- a/tools/stats_log_api_gen/utils.cpp
+++ b/tools/stats_log_api_gen/utils.cpp
@@ -334,7 +334,7 @@
         if (field->javaType == JAVA_TYPE_ATTRIBUTION_CHAIN) {
             fprintf(out, ", android.os.WorkSource workSource");
         } else if (field->javaType == JAVA_TYPE_KEY_VALUE_PAIR) {
-            fprintf(out, ", SparseArray<Object> value_map");
+            fprintf(out, ", android.util.SparseArray<Object> value_map");
         } else if (field->javaType == JAVA_TYPE_BYTE_ARRAY) {
             fprintf(out, ", byte[] %s", field->name.c_str());
         } else {
@@ -442,7 +442,7 @@
         for (vector<java_type_t>::const_iterator arg = signature.begin();
                 arg != signature.end(); arg++) {
             if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
-                fprintf(out, ", WorkSource ws");
+                fprintf(out, ", android.os.WorkSource ws");
             } else {
                 fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
             }
@@ -464,9 +464,10 @@
         fprintf(out, "        }\n"); // close for-loop
 
         // write() component.
-        fprintf(out, "        ArrayList<WorkSource.WorkChain> workChains = ws.getWorkChains();\n");
+        fprintf(out, "        java.util.ArrayList<android.os.WorkSource.WorkChain> workChains = "
+                "ws.getWorkChains();\n");
         fprintf(out, "        if (workChains != null) {\n");
-        fprintf(out, "            for (WorkSource.WorkChain wc : workChains) {\n");
+        fprintf(out, "            for (android.os.WorkSource.WorkChain wc : workChains) {\n");
         fprintf(out, "                write(code");
         for (int argIndex = 1; argIndex <= argIndexMax; argIndex++) {
             if (argIndex == attributionArg) {