Generate TEST_MAPPING for test modules in .bp file

* Share srcs in rust_defaults if len(srcs) is 1.

Test: regenerate .bp and TEST_MAPPING files in external/rust/crates
Test: make
Test: atest -c  --include-subdirs external/rust/crates
Bug: 161259631
Change-Id: I22a1d567ccabcd945e4cac870049ac3ce7a3049a
diff --git a/scripts/cargo2android.py b/scripts/cargo2android.py
index df841d9..4a05eeb 100755
--- a/scripts/cargo2android.py
+++ b/scripts/cargo2android.py
@@ -48,6 +48,9 @@
       --cargo "build --target x86_64-unknown-linux-gnu"
       --cargo "build --tests --target x86_64-unknown-linux-gnu"
 
+    Note that when there are test modules generated into Android.bp,
+    corresponding test entries will also be added into the TEST_MAPPING file.
+
 If there are rustc warning messages, this script will add
 a warning comment to the owner crate module in Android.bp.
 """
@@ -180,6 +183,38 @@
   return s.replace('"', '\\"')
 
 
+class TestMapping(object):
+  """Entries for a TEST_MAPPING file."""
+
+  def __init__(self):
+    self.entries = []
+
+  def add_test(self, name, host):
+    self.entries.append((name, host))
+
+  def is_empty(self):
+    return not self.entries
+
+  def dump(self, outf_name):
+    """Append all entries into the output file."""
+    if self.is_empty():
+      return
+    with open(outf_name, 'w') as outf:
+      outf.write('// Generated by cargo2android.py for tests in Android.bp\n')
+      outf.write('{\n  "presubmit": [\n')
+      is_first = True
+      for (name, host) in self.entries:
+        if not is_first:  # add comma and '\n' after the previous entry
+          outf.write(',\n')
+        is_first = False
+        outf.write('    {\n      "name": "' + name + '"')
+        if host:
+          outf.write(',\n      "host": true\n    }')
+        else:
+          outf.write('\n    }')
+      outf.write('\n  ]\n}\n')
+
+
 class Crate(object):
   """Information of a Rust crate to collect/emit for an Android.bp module."""
 
@@ -198,6 +233,7 @@
     self.module_name = ''  # unique in Android build system
     self.module_type = ''  # rust_{binary,library,test}[_host] etc.
     self.defaults = ''  # rust_defaults used by rust_test* modules
+    self.default_srcs = False  # use 'srcs' defined in self.defaults
     self.root_pkg = ''  # parent package name of a sub/test packge, from -L
     self.srcs = list()  # main_src or merged multiple source files
     self.stem = ''  # real base name of output file
@@ -560,6 +596,11 @@
     self.write('\nrust_defaults {')
     self.write('    name: "' + name + '",')
     self.write('    crate_name: "' + self.crate_name + '",')
+    if len(self.srcs) == 1:  # only one source file; share it in defaults
+      self.default_srcs = True
+      if self.has_warning and not self.cap_lints:
+        self.write('    // has rustc warnings')
+      self.write('    srcs: ["' + self.main_src + '"],')
     if 'test' in self.crate_types:
       self.write('    test_suites: ["general-tests"],')
       self.write('    auto_gen_config: true,')
@@ -591,12 +632,14 @@
         self.module_name = self.test_module_name()
         self.decide_one_module_type(crate_type)
         self.dump_one_android_module(crate_type)
+        self.runner.add_test(self.outf_name, self.module_name, True)
       if saved_device_supported:
         self.device_supported = True
         self.host_supported = False
         self.module_name = self.test_module_name()
         self.decide_one_module_type(crate_type)
         self.dump_one_android_module(crate_type)
+        self.runner.add_test(self.outf_name, self.module_name, False)
       self.host_supported = saved_host_supported
       self.device_supported = saved_device_supported
       self.main_src = saved_main_src
@@ -734,7 +777,7 @@
       self.write('    defaults: ["' + self.defaults + '"],')
     if self.stem != self.module_name:
       self.write('    stem: "' + self.stem + '",')
-    if self.has_warning and not self.cap_lints:
+    if self.has_warning and not self.cap_lints and not self.default_srcs:
       self.write('    // has rustc warnings')
     if self.host_supported and self.device_supported:
       self.write('    host_supported: true,')
@@ -743,7 +786,7 @@
     if len(self.srcs) > 1:
       self.srcs = sorted(set(self.srcs))
       self.dump_android_property_list('srcs', '"%s"', self.srcs)
-    else:
+    elif not self.default_srcs:
       self.write('    srcs: ["' + self.main_src + '"],')
     if 'test' in self.crate_types and not self.defaults:
       # self.root_pkg can have multiple test modules, with different *_tests[n]
@@ -964,6 +1007,7 @@
 
   def __init__(self, args):
     self.bp_files = set()  # Remember all output Android.bp files.
+    self.test_mappings = {}  # Map from Android.bp file path to TestMapping.
     self.root_pkg = ''  # name of package in ./Cargo.toml
     # Saved flags, modes, and data.
     self.args = args
@@ -1004,6 +1048,21 @@
       with open(name, 'w') as outf:
         outf.write(ANDROID_BP_HEADER.format(args=' '.join(sys.argv[1:])))
 
+  def dump_test_mapping_files(self):
+    if self.dry_run:
+      print('Dry-run skip dump of TEST_MAPPING')
+    else:
+      for bp_file_name in self.test_mappings:
+        name = os.path.join(os.path.dirname(bp_file_name), 'TEST_MAPPING')
+        self.test_mappings[bp_file_name].dump(name)
+    return self
+
+  def add_test(self, bp_file_name, test_name, host):
+    if bp_file_name not in self.test_mappings:
+      self.test_mappings[bp_file_name] = TestMapping()
+    mapping = self.test_mappings[bp_file_name]
+    mapping.add_test(test_name, host)
+
   def try_claim_module_name(self, name, owner):
     """Reserve and return True if it has not been reserved yet."""
     if name not in self.name_owners or owner == self.name_owners[name]:
@@ -1316,7 +1375,7 @@
   args = parse_args()
   if not args.run:  # default is dry-run
     print(DRY_RUN_NOTE)
-  Runner(args).run_cargo().gen_bp()
+  Runner(args).run_cargo().gen_bp().dump_test_mapping_files()
 
 
 if __name__ == '__main__':