#! /usr/bin/env python3
#
# Copyright 2020 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Regenerate some ART test related files.

# This script handles only a subset of ART run-tests at the moment; additional
# cases will be added later.

import argparse
import copy
import collections
import json
import logging
import os
import re
import sys
import textwrap
import xml.dom.minidom

logging.basicConfig(format='%(levelname)s: %(message)s')

ME = os.path.basename(sys.argv[0])

# Common advisory placed at the top of all generated files.
ADVISORY = f"Generated by `{ME}`. Do not edit manually."

# Default indentation unit.
INDENT = "  "

# Indentation unit for XML files.
XML_INDENT = "    "

def reindent(str, indent = ""):
  """Reindent literal string while removing common leading spaces."""
  return textwrap.indent(textwrap.dedent(str), indent)

def copyright_header_text(year):
  """Return the copyright header text used in XML files."""
  return reindent(f"""\
    Copyright (C) {year} The Android Open Source Project

        Licensed under the Apache License, Version 2.0 (the "License");
        you may not use this file except in compliance with the License.
        You may obtain a copy of the License at

             http://www.apache.org/licenses/LICENSE-2.0

        Unless required by applicable law or agreed to in writing, software
        distributed under the License is distributed on an "AS IS" BASIS,
        WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
        See the License for the specific language governing permissions and
        limitations under the License.
    """, " ")

def split_list(l, n):
  """Return a list of `n` sublists of (contiguous) elements of list `l`."""
  assert n > 0
  (d, m) = divmod(len(l), n)
  # If the length of `l` is divisible by `n`, use that that divisor (`d`) as size of each sublist;
  # otherwise, the next integer value (`d + 1`).
  s = d if m == 0 else d + 1
  result = [l[i:i + s] for i in range(0, len(l), s)]
  assert len(result) == n
  return result

# The prefix used in the Soong module name of all ART run-tests.
ART_RUN_TEST_MODULE_NAME_PREFIX = "art-run-test-"

# Number of shards used to declare ART run-tests in the sharded ART MTS test plan.
NUM_MTS_ART_RUN_TEST_SHARDS = 1

# Known failing ART run-tests.
# TODO(rpl): Investigate and address the causes of failures.
known_failing_tests = frozenset([
  "004-SignalTest",
  # 004-ThreadStress: Custom `run` script + dependency on `libarttest` + custom `check` script.
  "004-ThreadStress",
  "004-UnsafeTest",
  "030-bad-finalizer",
  "034-call-null",
  "038-inner-null",
  "044-proxy",
  "051-thread",
  "054-uncaught",
  "086-null-super",
  "087-gc-after-link",
  # 1002-notify-startup: Dependency on `libarttest` + custom `check` script.
  "1002-notify-startup",
  "115-native-bridge",
  "116-nodex2oat",
  "1336-short-finalizer-timeout",
  "1337-gc-coverage",
  "1339-dead-reference-safe",
  "134-nodex2oat-nofallback",
  "136-daemon-jni-shutdown",
  "139-register-natives",
  "148-multithread-gc-annotations",
  "149-suspend-all-stress",
  "150-loadlibrary",
  "154-gc-loop",
  # 157-void-class: Custom `run` script + dependency on `libarttest`.
  "157-void-class",
  "158-app-image-class-table",
  "169-threadgroup-jni",
  "172-app-image-twice",
  "177-visibly-initialized-deadlock",
  "178-app-image-native-method",
  "179-nonvirtual-jni",
  "674-HelloWorld-Dm",
  "1900-track-alloc",
  "1901-get-bytecodes",
  "1902-suspend",
  "1903-suspend-self",
  "1904-double-suspend",
  "1905-suspend-native",
  "1906-suspend-list-me-first",
  "1907-suspend-list-self-twice",
  "1908-suspend-native-resume-self",
  "1909-per-agent-tls",
  "1910-transform-with-default",
  "1911-get-local-var-table",
  "1912-get-set-local-primitive",
  "1913-get-set-local-objects",
  "1914-get-local-instance",
  "1915-get-set-local-current-thread",
  "1916-get-set-current-frame",
  "1917-get-stack-frame",
  "1919-vminit-thread-start-timing",
  "1920-suspend-native-monitor",
  "1921-suspend-native-recursive-monitor",
  "1922-owned-monitors-info",
  "1923-frame-pop",
  "1924-frame-pop-toggle",
  "1925-self-frame-pop",
  "1926-missed-frame-pop",
  "1927-exception-event",
  "1928-exception-event-exception",
  "1930-monitor-info",
  "1931-monitor-events",
  "1932-monitor-events-misc",
  "1933-monitor-current-contended",
  "1934-jvmti-signal-thread",
  "1935-get-set-current-frame-jit",
  "1936-thread-end-events",
  "1937-transform-soft-fail",
  "1938-transform-abstract-single-impl",
  "1939-proxy-frames",
  "1941-dispose-stress",
  "1942-suspend-raw-monitor-exit",
  "1943-suspend-raw-monitor-wait",
  "1945-proxy-method-arguments",
  "1947-breakpoint-redefine-deopt",
  "1949-short-dex-file",
  "1951-monitor-enter-no-suspend",
  "1953-pop-frame",
  "1954-pop-frame-jit",
  "1955-pop-frame-jit-called",
  "1956-pop-frame-jit-calling",
  "1957-error-ext",
  "1958-transform-try-jit",
  "1959-redefine-object-instrument",
  "1960-obsolete-jit-multithread-native",
  "1961-obsolete-jit-multithread",
  "1962-multi-thread-events",
  "1963-add-to-dex-classloader-in-memory",
  "1967-get-set-local-bad-slot",
  "1968-force-early-return",
  "1969-force-early-return-void",
  "1970-force-early-return-long",
  "1971-multi-force-early-return",
  "1972-jni-id-swap-indices",
  "1973-jni-id-swap-pointer",
  "1974-resize-array",
  "1975-hello-structural-transformation",
  "1976-hello-structural-static-methods",
  "1977-hello-structural-obsolescence",
  "1978-regular-obsolete-then-structural-obsolescence",
  "1979-threaded-structural-transformation",
  "1980-obsolete-object-cleared",
  # 1981-structural-redef-private-method-handles: Custom `run` script.
  "1981-structural-redef-private-method-handles",
  "1982-no-virtuals-structural-redefinition",
  # 1983-structural-redefinition-failures: Custom `run` script.
  "1983-structural-redefinition-failures",
  "1984-structural-redefine-field-trace",
  "1985-structural-redefine-stack-scope",
  "1986-structural-redefine-multi-thread-stack-scope",
  "1987-structural-redefine-recursive-stack-scope",
  "1988-multi-structural-redefine",
  "1989-transform-bad-monitor",
  "1990-structural-bad-verify",
  "1991-hello-structural-retransform",
  "1992-retransform-no-such-field",
  "1993-fallback-non-structural",
  "1994-final-virtual-structural",
  "1995-final-virtual-structural-multithread",
  "1996-final-override-virtual-structural",
  "1997-structural-shadow-method",
  "1998-structural-shadow-field",
  "1999-virtual-structural",
  "2003-double-virtual-structural",
  "2004-double-virtual-structural-abstract",
  "2005-pause-all-redefine-multithreaded",
  "2008-redefine-then-old-reflect-field",
  "2011-stack-walk-concurrent-instrument",
  # 2012-structural-redefinition-failures-jni-id: Custom `run` script.
  "2012-structural-redefinition-failures-jni-id",
  # 2040-huge-native-alloc: Fails with:
  #
  #   Test command execution failed with status FAILED: CommandResult: exit code=1, out=, err=Exception in thread "main" java.lang.ArrayIndexOutOfBoundsException: length=0; index=0
  #           at Main.main(Main.java:56)
  #
  "2040-huge-native-alloc",
  "2041-bad-cleaner",
  "203-multi-checkpoint",
  "2031-zygote-compiled-frame-deopt",
  "2033-shutdown-mechanics",
  "2036-jni-filechannel",
  "2037-thread-name-inherit",
  # 2230-profile-save-hotness: Custom `run` script + dependency on `libarttest`.
  "2230-profile-save-hotness",
  "2235-JdkUnsafeTest",
  "305-other-fault-handler",
  # 449-checker-bce: Dependency on `libarttest`.
  "449-checker-bce",
  "454-get-vreg",
  "461-get-reference-vreg",
  "466-get-live-vreg",
  "497-inlining-and-class-loader",
  "530-regression-lse",
  "555-UnsafeGetLong-regression",
  "566-polymorphic-inlining",
  "595-profile-saving",
  # 596-monitor-inflation: Dependency on `libarttest`.
  "596-monitor-inflation",
  "597-deopt-busy-loop",
  "597-deopt-invoke-stub",
  "597-deopt-new-string",
  "602-deoptimizeable",
  "604-hot-static-interface",
  "616-cha-abstract",
  "616-cha-interface",
  "616-cha-miranda",
  "616-cha-native",
  "616-cha-regression-proxy-method",
  "616-cha",
  # 623-checker-loop-regressions: Dependency on `libarttest`.
  "623-checker-loop-regressions",
  "626-set-resolved-string",
  "629-vdex-speed",
  # 638-checker-inline-cache-intrinsic: Custom `run` script + dependency on `libarttest`.
  "638-checker-inline-cache-intrinsic",
  "642-fp-callees",
  "647-jni-get-field-id",
  "652-deopt-intrinsic",
  "655-jit-clinit",
  "656-loop-deopt",
  "660-clinit",
  "661-oat-writer-layout",
  "664-aget-verifier",
  "667-jit-jni-stub",
  "674-hotness-compiled",
  "679-locks",
  # 680-checker-deopt-dex-pc-0: Dependency on `libarttest`.
  "680-checker-deopt-dex-pc-0",
  "685-deoptimizeable",
  "687-deopt",
  "689-zygote-jit-deopt",
  "693-vdex-inmem-loader-evict",
  # 707-checker-invalid-profile: Custom `run` script + custom `check` script.
  "707-checker-invalid-profile",
  "708-jit-cache-churn",
  # 716-jli-jit-samples: Dependency on `libarttest`.
  "716-jli-jit-samples",
  "717-integer-value-of",
  "720-thread-priority",
  # 728-imt-conflict-zygote: Custom `run` script + dependency on `libarttest`.
  "728-imt-conflict-zygote",
  # 729-checker-polymorphic-intrinsic: Custom `run` script.
  "729-checker-polymorphic-intrinsic",
  # 730-cha-deopt: Fails with:
  #
  #   Test command execution failed with status FAILED: CommandResult: exit code=1, out=, err=Exception in thread "main" java.lang.ArrayIndexOutOfBoundsException: length=0; index=0
  #           at Main.main(Main.java:24)
  #
  "730-cha-deopt",
  # 813-fp-args: Dependency on `libarttest`.
  "813-fp-args",
  # 821-many-args: Fails with:
  #
  #   Test command execution failed with status FAILED: CommandResult: exit code=1, out=, err=Exception in thread "main" java.lang.ArrayIndexOutOfBoundsException: length=0; index=0
  #           at Main.main(Main.java:20)
  #
  "821-many-args",
  # 826-infinite-loop: The test expects an argument passed to `Main.main` (the test library,
  # usually `arttestd` or `arttest)`, but the ART run-test TradeFed test runner
  # (`com.android.tradefed.testtype.ArtRunTest`) does not implement this yet.
  "826-infinite-loop",
  # 832-cha-recursive: Dependency on `libarttest`.
  "832-cha-recursive",
  "900-hello-plugin",
  "901-hello-ti-agent",
  "902-hello-transformation",
  "903-hello-tagging",
  "904-object-allocation",
  "905-object-free",
  "906-iterate-heap",
  "907-get-loaded-classes",
  # 909-attach-agent: Custom `run` script + dependency on `libarttest`.
  "909-attach-agent",
  "908-gc-start-finish",
  "910-methods",
  "911-get-stack-trace",
  # 912-classes: Custom `run` script + dependency on `libtiagent`.
  "912-classes",
  "913-heaps",
  "914-hello-obsolescence",
  "915-obsolete-2",
  "916-obsolete-jit",
  "917-fields-transformation",
  "918-fields",
  "919-obsolete-fields",
  "920-objects",
  "921-hello-failure",
  "922-properties",
  "923-monitors",
  "924-threads",
  "925-threadgroups",
  "926-multi-obsolescence",
  "927-timers",
  "928-jni-table",
  "930-hello-retransform",
  "931-agent-thread",
  "932-transform-saves",
  "933-misc-events",
  "937-hello-retransform-package",
  "939-hello-transformation-bcp",
  "940-recursive-obsolete",
  "941-recursive-obsolete-jit",
  "942-private-recursive",
  "943-private-recursive-jit",
  "944-transform-classloaders",
  "945-obsolete-native",
  "946-obsolete-throw",
  "947-reflect-method",
  "949-in-memory-transform",
  "950-redefine-intrinsic",
  "951-threaded-obsolete",
  # 981-dedup-original-dex: Custom `run` script + dependency on `libtiagent`.
  "981-dedup-original-dex",
  "982-ok-no-retransform",
  "983-source-transform-verify",
  "984-obsolete-invoke",
  "985-re-obsolete",
  "986-native-method-bind",
  "987-agent-bind",
  # 988-method-trace: Custom `run` script.
  "988-method-trace",
  "989-method-trace-throw",
  "990-field-trace",
  "991-field-trace-2",
  "992-source-data",
  "993-breakpoints",
  "994-breakpoint-line",
  "995-breakpoints-throw",
  "996-breakpoint-obsolete",
  "997-single-step",
])

# Percentage of ART run-tests (among the ones expected to succeed) to include in
# the `presubmit` test group in `TEST_MAPPING` file -- the rest will be included
# in `postsubmit` test group.
# This value has to be a number between 0 and 100.
presubmit_tests_percentage = 100

# Percentage of ART run-tests (among the ones expected to succeed) to include in
# the `mainline-presubmit` test group in `TEST_MAPPING` file.
# This value has to be a number between 0 and 100.
mainline_presubmit_tests_percentage = 100

# ART gtests that do not need root access to the device.
art_gtest_user_module_names = [
    "art_standalone_cmdline_tests",
    "art_standalone_compiler_tests",
    "art_standalone_dex2oat_tests",
    "art_standalone_dexdump_tests",
    "art_standalone_dexlist_tests",
    "art_standalone_libartbase_tests",
    "art_standalone_libartpalette_tests",
    "art_standalone_libartservice_tests",
    "art_standalone_libarttools_tests",
    "art_standalone_libdexfile_support_tests",
    "art_standalone_libdexfile_tests",
    "art_standalone_libprofile_tests",
    "art_standalone_oatdump_tests",
    "art_standalone_odrefresh_tests",
    "art_standalone_runtime_compiler_tests",
    "art_standalone_runtime_tests",
    "art_standalone_sigchain_tests",
]

# ART gtests that need root access to the device.
art_gtest_eng_only_module_names = [
    "art_standalone_dexoptanalyzer_tests",
    "art_standalone_profman_tests",
]

# All supported ART gtests.
art_gtest_module_names = sorted(art_gtest_user_module_names + art_gtest_eng_only_module_names)

# ART gtests supported in MTS that do not need root access to the device.
art_gtest_mts_user_module_names = copy.copy(art_gtest_user_module_names)
# Temporarily disable `art_standalone_odrefresh_tests` in MTS,
# as it is currently failing in Mainline testing
# (b/206335809); a fix is in the works but may take some time
# to land.
#
# TODO(b/206335809): Re-enable this test when the fix has landed.
art_gtest_mts_user_module_names.remove("art_standalone_odrefresh_tests")

# ART gtests supported in Mainline presubmits.
art_gtests_mainline_presubmit_module_names = copy.copy(art_gtest_module_names)
# Temporarily disable `art_standalone_odrefresh_tests` in Mainline
# presubmits, as it is currently failing in Mainline testing
# (b/206335809); a fix is in the works but may take some time to
# land.
#
# TODO(b/206335809): Re-enable this test when the fix has landed.
art_gtests_mainline_presubmit_module_names.remove("art_standalone_odrefresh_tests")


# Is `run_test` a Checker test (i.e. a test containing Checker
# assertions)?
def is_checker_test(run_test):
  return re.match("^[0-9]+-checker-", run_test)

# Is `run_test` expected to succeed?
def is_expected_succeeding(run_test):
  return run_test not in known_failing_tests


class Generator:
  def __init__(self, top_dir):
    """Generator of ART test files for an Android source tree anchored at `top_dir`."""
    # Path to the Android top source tree.
    self.top_dir = top_dir
    # Path to the ART directory
    self.art_dir = os.path.join(top_dir, "art")
    # Path to the ART tests directory.
    self.art_test_dir = os.path.join(self.art_dir, "test")
    # Path to the MTS configuration directory.
    self.mts_config_dir = os.path.join(
        top_dir, "test", "mts", "tools", "mts-tradefed", "res", "config")

  def enumerate_run_tests(self):
    return sorted([run_test
                   for run_test in os.listdir(self.art_test_dir)
                   if re.match("^[0-9]{3,}-", run_test)])

  # Read build file (Bash script) and return a canonized version of it
  # (without comments, blank lines, "debugging" statements, etc.).
  def canonize_build_script(self, build_file):

    def is_comment(line):
      return re.match("^\\s*#", line)

    def is_blank(line):
      return re.match("^\\s*$", line)

    # Is `line` a `set -e` statement?
    def is_set_e(line):
      return re.match("^\\s*set -e\\s*", line)

    # Should `line` be kept in the canonized build script?
    def keep_line(line):
      return not (is_comment(line) or is_blank(line) or is_set_e(line))

    with open(build_file, "r") as f:
      lines = f.readlines()
    return list(filter(keep_line, lines))

  # Can the build script in `build_file` be safely ignored?
  def can_ignore_build_script(self, build_file):
    build_script = self.canonize_build_script(build_file)
    if len(build_script) == 1:
      if build_script[0] == "./default-build \"$@\" --experimental var-handles\n":
        # Soong builds JARs with VarHandle support by default (i.e. by
        # using an API level greater or equal to 28), so we can ignore
        # build scripts that just request support for this feature.
        return True
    return False

  # Is building `run_test` supported?
  # TODO(b/147814778): Add build support for more tests.
  def is_buildable(self, run_test):
    run_test_path = os.path.join(self.art_test_dir, run_test)

    # Skip tests with non-default build rules, unless these build
    # rules can be safely ignored.
    if os.path.isfile(os.path.join(run_test_path, "build")):
      if not self.can_ignore_build_script(os.path.join(run_test_path, "build")):
        return False
    # Skip tests with sources outside the `src` directory.
    for subdir in ["jasmin",
                   "jasmin-multidex",
                   "smali",
                   "smali-ex",
                   "smali-multidex",
                   "src-dex2oat-unresolved",
                   "src-ex",
                   "src-ex2",
                   "src-multidex",
                   "src2"]:
      if os.path.isdir(os.path.join(run_test_path, subdir)):
        return False
    # Skip tests that have both an `src` directory and an `src-art` directory.
    if os.path.isdir(os.path.join(run_test_path, "src")) and \
       os.path.isdir(os.path.join(run_test_path, "src-art")):
        return False
    # Skip tests that have neither an `src` directory nor an `src-art` directory.
    if not os.path.isdir(os.path.join(run_test_path, "src")) and \
       not os.path.isdir(os.path.join(run_test_path, "src-art")):
      return False
    # Skip test with a copy of `sun.misc.Unsafe`.
    if os.path.isfile(os.path.join(run_test_path, "src", "sun", "misc", "Unsafe.java")):
      return False
    # Skip tests with Hidden API specs.
    if os.path.isfile(os.path.join(run_test_path, "hiddenapi-flags.csv")):
      return False
    # All other tests are considered buildable.
    return True

  def regen_bp_files(self, run_tests, buildable_tests):
    for run_test in run_tests:
      # Remove any previously generated file.
      bp_file = os.path.join(self.art_test_dir, run_test, "Android.bp")
      if os.path.exists(bp_file):
        logging.debug(f"Removing `{bp_file}`.")
        os.remove(bp_file)

    for run_test in buildable_tests:
      self.regen_bp_file(run_test)

  def regen_bp_file(self, run_test):
    """Regenerate Blueprint file for an ART run-test."""

    run_test_path = os.path.join(self.art_test_dir, run_test)
    bp_file = os.path.join(run_test_path, "Android.bp")

    run_test_module_name = ART_RUN_TEST_MODULE_NAME_PREFIX + run_test

    if is_expected_succeeding(run_test):
      test_config_template = "art-run-test-target-template"
    else:
      test_config_template = "art-run-test-target-no-test-suite-tag-template"

    if is_checker_test(run_test):
      include_src = """\

          // Include the Java source files in the test's artifacts, to make Checker assertions
          // available to the TradeFed test runner.
          include_srcs: true,"""
    else:
      include_src = ""

    # The default source directory is `src`, except if `src-art` exists.
    if os.path.isdir(os.path.join(run_test_path, "src-art")):
      source_dir = "src-art"
    else:
      source_dir = "src"

    with open(bp_file, "w") as f:
      logging.debug(f"Writing `{bp_file}`.")
      f.write(textwrap.dedent(f"""\
      // {ADVISORY}

      // Build rules for ART run-test `{run_test}`.

      package {{
          // See: http://go/android-license-faq
          // A large-scale-change added 'default_applicable_licenses' to import
          // all of the 'license_kinds' from "art_license"
          // to get the below license kinds:
          //   SPDX-license-identifier-Apache-2.0
          default_applicable_licenses: ["art_license"],
      }}

      // Test's Dex code.
      java_test {{
          name: "{run_test_module_name}",
          defaults: ["art-run-test-defaults"],
          test_config_template: ":{test_config_template}",
          srcs: ["{source_dir}/**/*.java"],
          data: [
              ":{run_test_module_name}-expected-stdout",
              ":{run_test_module_name}-expected-stderr",
          ],{include_src}
      }}

      // Test's expected standard output.
      genrule {{
          name: "{run_test_module_name}-expected-stdout",
          out: ["{run_test_module_name}-expected-stdout.txt"],
          srcs: ["expected-stdout.txt"],
          cmd: "cp -f $(in) $(out)",
      }}

      // Test's expected standard error.
      genrule {{
          name: "{run_test_module_name}-expected-stderr",
          out: ["{run_test_module_name}-expected-stderr.txt"],
          srcs: ["expected-stderr.txt"],
          cmd: "cp -f $(in) $(out)",
      }}
      """))

  def regen_test_mapping_file(self, art_run_tests, num_presubmit_run_tests,
                              num_mainline_presubmit_run_tests):
    """Regenerate ART's `TEST_MAPPING`."""

    run_test_module_names = [ART_RUN_TEST_MODULE_NAME_PREFIX + t for t in art_run_tests]

    # Mainline presubmits.
    mainline_other_presubmit_tests = [
        "ComposHostTestCases",
    ]
    mainline_presubmit_run_tests = run_test_module_names[0:num_mainline_presubmit_run_tests]
    mainline_presubmit_tests = (mainline_other_presubmit_tests + mainline_presubmit_run_tests +
                                art_gtests_mainline_presubmit_module_names)
    mainline_presubmit_tests_with_apex = [t + "[com.google.android.art.apex]"
                                          for t
                                          in mainline_presubmit_tests]
    mainline_presubmit_tests_dict = [{"name": t} for t in mainline_presubmit_tests_with_apex]

    # Presubmits.
    other_presubmit_tests = [
        "CtsJdwpTestCases",
        "BootImageProfileTest",
        "ArtServiceTests",
        "ComposHostTestCases",
    ]
    presubmit_run_tests = run_test_module_names[0:num_presubmit_run_tests]
    presubmit_tests = other_presubmit_tests + presubmit_run_tests + art_gtest_module_names
    presubmit_tests_dict = [{"name": t} for t in presubmit_tests]

    # Postsubmits.
    postsubmit_run_tests = run_test_module_names[num_presubmit_run_tests:]
    postsubmit_tests_dict = [{"name": t} for t in postsubmit_run_tests]

    # Use an `OrderedDict` container to preserve the order in which items are inserted.
    # Do not produce an entry for a test group if it is empty.
    test_mapping_dict = collections.OrderedDict([
        (test_group_name, test_group_dict)
        for (test_group_name, test_group_dict)
        in [
            ("mainline-presubmit", mainline_presubmit_tests_dict),
            ("presubmit", presubmit_tests_dict),
            ("postsubmit", postsubmit_tests_dict),
        ]
        if test_group_dict
    ])
    test_mapping_contents = json.dumps(test_mapping_dict, indent = INDENT)

    test_mapping_file = os.path.join(self.art_dir, "TEST_MAPPING")
    with open(test_mapping_file, "w") as f:
      logging.debug(f"Writing `{test_mapping_file}`.")
      f.write(f"// {ADVISORY}\n")
      f.write(test_mapping_contents)
      f.write("\n")

  def create_mts_test_shard(self, description, tests, shard_num, copyright_year, comments = []):
    """Factory method instantiating an `MtsTestShard`."""
    return self.MtsTestShard(self.mts_config_dir,
                             description, tests, shard_num, copyright_year, comments)

  class MtsTestShard:
    """Class encapsulating data and generation logic for an ART MTS test shard."""

    def __init__(self, mts_config_dir, description, tests, shard_num, copyright_year, comments):
      self.mts_config_dir = mts_config_dir
      self.description = description
      self.tests = tests
      self.shard_num = shard_num
      self.copyright_year = copyright_year
      self.comments = comments

    def shard_id(self):
      return f"{self.shard_num:02}"

    def test_plan_name(self):
      return "mts-art-shard-" + self.shard_id()

    def test_list_name(self):
      return "mts-art-tests-list-user-shard-" + self.shard_id()

    def regen_test_plan_file(self):
      """Regenerate ART MTS test plan file shard (`mts-art-shard-<shard_num>.xml`)."""
      root = xml.dom.minidom.Document()

      advisory_header = root.createComment(f" {ADVISORY} ")
      root.appendChild(advisory_header)
      copyright_header = root.createComment(copyright_header_text(self.copyright_year))
      root.appendChild(copyright_header)

      configuration = root.createElement("configuration")
      root.appendChild(configuration)
      configuration.setAttribute(
          "description",
          f"Run mts-art-shard-{self.shard_id()} from a preexisting MTS installation.")

      # Included XML files.
      included_xml_files = ["mts", self.test_list_name()]
      for xml_file in included_xml_files:
        include = root.createElement("include")
        include.setAttribute("name", xml_file)
        configuration.appendChild(include)

      # Test plan name.
      option = root.createElement("option")
      option.setAttribute("name", "plan")
      option.setAttribute("value", self.test_plan_name())
      configuration.appendChild(option)

      xml_str = root.toprettyxml(indent = XML_INDENT, encoding = "utf-8")

      test_plan_file = os.path.join(self.mts_config_dir, self.test_plan_name() + ".xml")
      with open(test_plan_file, "wb") as f:
        logging.debug(f"Writing `{test_plan_file}`.")
        f.write(xml_str)

    def regen_test_list_file(self):
      """Regenerate ART MTS test list file (`mts-art-tests-list-user-shard-<shard_num>.xml`)."""
      root = xml.dom.minidom.Document()

      advisory_header = root.createComment(f" {ADVISORY} ")
      root.appendChild(advisory_header)
      copyright_header = root.createComment(copyright_header_text(self.copyright_year))
      root.appendChild(copyright_header)

      configuration = root.createElement("configuration")
      root.appendChild(configuration)
      configuration.setAttribute(
          "description",
          f"List of ART MTS tests that do not need root access (shard {self.shard_id()})"
      )

      # Test declarations.
      # ------------------

      def append_test_declaration(test):
        option = root.createElement("option")
        option.setAttribute("name", "compatibility:include-filter")
        option.setAttribute("value", test)
        configuration.appendChild(option)

      test_declarations_comments = [self.description + "."]
      test_declarations_comments.extend(self.comments)
      for c in test_declarations_comments:
        xml_comment = root.createComment(f" {c} ")
        configuration.appendChild(xml_comment)
      for t in self.tests:
        append_test_declaration(t)

      # `MainlineTestModuleController` configurations.
      # ----------------------------------------------

      def append_module_controller_configuration(test):
        option = root.createElement("option")
        option.setAttribute("name", "compatibility:module-arg")
        option.setAttribute("value", f"{test}:enable:true")
        configuration.appendChild(option)

      module_controller_configuration_comments = [
          f"Enable MainlineTestModuleController for {self.description}."]
      module_controller_configuration_comments.extend(self.comments)
      for c in module_controller_configuration_comments:
        xml_comment = root.createComment(f" {c} ")
        configuration.appendChild(xml_comment)
      for t in self.tests:
        append_module_controller_configuration(t)

      xml_str = root.toprettyxml(indent = XML_INDENT, encoding = "utf-8")

      test_list_file = os.path.join(self.mts_config_dir, self.test_list_name() + ".xml")
      with open(test_list_file, "wb") as f:
        logging.debug(f"Writing `{test_list_file}`.")
        f.write(xml_str)

  def regen_mts_art_tests_list_user_file(self, num_mts_art_run_test_shards):
    """Regenerate ART MTS test list file (`mts-art-tests-list-user.xml`)."""
    root = xml.dom.minidom.Document()

    advisory_header = root.createComment(f" {ADVISORY} ")
    root.appendChild(advisory_header)
    copyright_header = root.createComment(copyright_header_text(2020))
    root.appendChild(copyright_header)

    configuration = root.createElement("configuration")
    root.appendChild(configuration)
    configuration.setAttribute("description", "List of ART MTS tests that do not need root access.")

    # Included XML files.
    for s in range(num_mts_art_run_test_shards):
      include = root.createElement("include")
      include.setAttribute("name", f"mts-art-tests-list-user-shard-{s:02}")
      configuration.appendChild(include)

    xml_str = root.toprettyxml(indent = XML_INDENT, encoding = "utf-8")

    mts_art_tests_list_user_file = os.path.join(self.mts_config_dir, "mts-art-tests-list-user.xml")
    with open(mts_art_tests_list_user_file, "wb") as f:
      logging.debug(f"Writing `{mts_art_tests_list_user_file}`.")
      f.write(xml_str)

  def regen_art_mts_files(self, art_run_tests):
    """Regenerate ART MTS definition files."""

    # Remove any previously MTS ART test plan shard (`mts-art-shard-[0-9]+.xml`)
    # and any test list shard (`mts-art-tests-list-user-shard-[0-9]+.xml`).
    old_test_plan_shards = sorted([
        test_plan_shard
        for test_plan_shard in os.listdir(self.mts_config_dir)
        if re.match("^mts-art-(tests-list-user-)?shard-[0-9]+.xml$", test_plan_shard)])
    for shard in old_test_plan_shards:
      shard_path = os.path.join(self.mts_config_dir, shard)
      if os.path.exists(shard_path):
        logging.debug(f"Removing `{shard_path}`.")
        os.remove(shard_path)

    mts_test_shards = []

    # ART test (gtest & run-test) shard(s).
    # TODO: Also handle the case of gtests requiring root access to the device
    # (`art_gtest_eng_only_module_names`).
    art_run_test_module_names = [ART_RUN_TEST_MODULE_NAME_PREFIX + t for t in art_run_tests]
    art_run_test_shards = split_list(art_run_test_module_names, NUM_MTS_ART_RUN_TEST_SHARDS)
    for i in range(len(art_run_test_shards)):
      art_tests_shard_i_tests = art_run_test_shards[i]
      # Append ART gtests to the last ART run-test shard for now.
      # If needed, consider moving them to their own shard to increase
      # the parallelization of code coverage runs.
      if i + 1 == len(art_run_test_shards):
        art_tests_shard_i_tests.extend(art_gtest_mts_user_module_names)
      art_tests_shard_i = self.create_mts_test_shard(
          "ART run-tests", art_tests_shard_i_tests, i, 2020,
          ["TODO(rpl): Find a way to express this list in a more concise fashion."])
      mts_test_shards.append(art_tests_shard_i)

    # CTS Libcore non-OJ tests (`CtsLibcoreTestCases`) shard.
    cts_libcore_tests_shard_num = len(mts_test_shards)
    cts_libcore_tests_shard = self.create_mts_test_shard(
        "CTS Libcore non-OJ tests", ["CtsLibcoreTestCases"], cts_libcore_tests_shard_num, 2020)
    mts_test_shards.append(cts_libcore_tests_shard)

    # Other CTS Libcore tests shard.
    other_cts_libcore_tests_shard_num = len(mts_test_shards)
    other_cts_libcore_tests_shard_tests = [
        "CtsLibcoreApiEvolutionTestCases",
        "CtsLibcoreFileIOTestCases",
        "CtsLibcoreJsr166TestCases",
        "CtsLibcoreLegacy22TestCases",
        "CtsLibcoreOjTestCases",
        "CtsLibcoreWycheproofBCTestCases",
        "MtsLibcoreOkHttpTestCases",
    ]
    other_cts_libcore_tests_shard = self.create_mts_test_shard(
        "CTS Libcore OJ tests", other_cts_libcore_tests_shard_tests,
        other_cts_libcore_tests_shard_num, 2021)
    mts_test_shards.append(other_cts_libcore_tests_shard)

    for s in mts_test_shards:
      s.regen_test_plan_file()
      s.regen_test_list_file()

    self.regen_mts_art_tests_list_user_file(len(mts_test_shards))

  def regen_test_files(self, regen_art_mts):
    """Regenerate ART test files.

    Args:
      regen_art_mts: If true, also regenerate the ART MTS definition.
    """
    run_tests = self.enumerate_run_tests()

    # Create a list of the tests that can currently be built, and for
    # which a Blueprint file is to be generated.
    buildable_tests = list(filter(self.is_buildable, run_tests))

    # Create a list of the tests that can be built and are expected to
    # succeed. These tests are to be added to ART's `TEST_MAPPING`
    # file and also tagged as part of TradeFed's `art-target-run-test`
    # test suite via the `test-suite-tag` option in their
    # configuration file.
    expected_succeeding_tests = list(filter(is_expected_succeeding, buildable_tests))

    # Regenerate Blueprint files.
    # ---------------------------

    self.regen_bp_files(run_tests, buildable_tests)

    buildable_tests_percentage = int(len(buildable_tests) * 100 / len(run_tests))

    print(f"Generated Blueprint files for {len(buildable_tests)} ART run-tests out of"
          f" {len(run_tests)} ({buildable_tests_percentage}%).")

    # Regenerate `TEST_MAPPING` file.
    # -------------------------------

    # Note: We only include ART run-tests expected to succeed for now.

    # Note: We only include a (growing) fraction of the supported ART
    # run-tests (see `presubmit_tests_percentage`) into the
    # `presubmit` test group (the other ART run-tests are added to the
    # `postsubmit` test group), as we initially had issues with
    # Android presubmits when the whole set of supported ART run-tests
    # was included in one go (b/169310621). This progressive rollout
    # allows us to better monitor future potential presubmit failures.
    #
    # Likewise for tests in the `mainline-presubmit` group.
    num_presubmit_run_tests = int(len(expected_succeeding_tests) * presubmit_tests_percentage / 100)
    num_mainline_presubmit_run_tests = int(
        len(expected_succeeding_tests) * mainline_presubmit_tests_percentage / 100)
    self.regen_test_mapping_file(
        expected_succeeding_tests, num_presubmit_run_tests, num_mainline_presubmit_run_tests)

    expected_succeeding_tests_percentage = int(
        len(expected_succeeding_tests) * 100 / len(run_tests))

    mainline_presubmit_gtests_percentage = int(
        len(art_gtests_mainline_presubmit_module_names) * 100 / len(art_gtest_module_names))

    num_postsubmit_tests = len(expected_succeeding_tests) - num_presubmit_run_tests
    postsubmit_tests_percentage = 100 - presubmit_tests_percentage

    print(f"Generated TEST_MAPPING entries for {len(expected_succeeding_tests)} ART run-tests out"
          f" of {len(run_tests)} ({expected_succeeding_tests_percentage}%):")
    for (num_tests, test_kind, tests_percentage, test_group_name) in [
        (num_mainline_presubmit_run_tests, "ART run-tests", mainline_presubmit_tests_percentage,
         "mainline-presubmit"),
        (len(art_gtests_mainline_presubmit_module_names), "ART gtests",
         mainline_presubmit_gtests_percentage, "mainline-presubmit"),
        (num_presubmit_run_tests, "ART run-tests", presubmit_tests_percentage, "presubmit"),
        (len(art_gtest_module_names), "ART gtests", 100, "presubmit"),
        (num_postsubmit_tests, "ART run-tests", postsubmit_tests_percentage, "postsubmit"),
    ]:
      print(
          f"  {num_tests:3d} {test_kind} ({tests_percentage}%) in `{test_group_name}` test group.")

    # Regenerate ART MTS definition (optional).
    # -----------------------------------------

    if regen_art_mts:
      self.regen_art_mts_files(expected_succeeding_tests)
      print(f"Generated ART MTS entries for {len(expected_succeeding_tests)} ART run-tests out"
            f" of {len(run_tests)} ({expected_succeeding_tests_percentage}%).")

def main():
  if "ANDROID_BUILD_TOP" not in os.environ:
    logging.error("ANDROID_BUILD_TOP environment variable is empty; did you forget to run `lunch`?")
    sys.exit(1)

  parser = argparse.ArgumentParser(
      formatter_class=argparse.RawDescriptionHelpFormatter,
      description=textwrap.dedent("Regenerate some ART test related files."),
      epilog=textwrap.dedent("""\
        Regenerate ART run-tests Blueprint files, ART's `TEST_MAPPING` file, and
        optionally the ART MTS (Mainline Test Suite) definition.
        """))
  parser.add_argument("-m", "--regen-art-mts", help="regenerate the ART MTS definition as well",
                      action="store_true")
  parser.add_argument("-v", "--verbose", help="enable verbose output", action="store_true")
  args = parser.parse_args()

  if args.verbose:
    logging.getLogger().setLevel(logging.DEBUG)

  generator = Generator(os.path.join(os.environ["ANDROID_BUILD_TOP"]))
  generator.regen_test_files(args.regen_art_mts)


if __name__ == "__main__":
  main()