1#!/usr/bin/python3
2#
3# Copyright (C) 2021 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17#
18# This script runs dex2oat on the host to compile a provided JAR or APK.
19#
20
21import argparse
22import itertools
23import shlex
24import subprocess
25import os
26import os.path
27
28def run_print(lst):
29  return " ".join(map(shlex.quote, lst))
30
31
32def parse_args():
33  parser = argparse.ArgumentParser(
34      description="compile dex or jar files",
35      epilog="Unrecognized options are passed on to dex2oat unmodified.")
36  parser.add_argument(
37      "--dex2oat",
38      action="store",
39      default=os.path.expandvars("$ANDROID_HOST_OUT/bin/dex2oatd64"),
40      help="selects the dex2oat to use.")
41  parser.add_argument(
42      "--debug",
43      action="store_true",
44      default=False,
45      help="launches dex2oatd with lldb-server g :5039. Connect using vscode or remote lldb"
46  )
47  parser.add_argument(
48      "--profman",
49      action="store",
50      default=os.path.expandvars("$ANDROID_HOST_OUT/bin/profmand"),
51      help="selects the profman to use.")
52  parser.add_argument(
53      "--debug-profman",
54      action="store_true",
55      default=False,
56      help="launches profman with lldb-server g :5039. Connect using vscode or remote lldb"
57  )
58  profs = parser.add_mutually_exclusive_group()
59  profs.add_argument(
60      "--profile-file",
61      action="store",
62      help="Use this profile file. Probably want to pass --compiler-filter=speed-profile with this."
63  )
64  profs.add_argument(
65      "--profile-line",
66      action="append",
67      default=[],
68      help="functions to add to a profile. Probably want to pass --compiler-filter=speed-profile with this. All functions are marked as 'hot'. Use --profile-file for more control."
69  )
70  parser.add_argument(
71      "--add-bcp",
72      action="append",
73      default=[],
74      nargs=2,
75      metavar=("BCP_FILE", "BCP_LOCATION"),
76      help="File and location to add to the boot-class-path. Note no deduplication is attempted."
77  )
78  parser.add_argument(
79      "--arch",
80      action="store",
81      choices=["arm", "arm64", "x86", "x86_64", "host64", "host32"],
82      default="host64",
83      help="architecture to compile for. Defaults to host64")
84  parser.add_argument(
85      "--odex-file",
86      action="store",
87      help="odex file to write. File discarded if not set",
88      default=None)
89  parser.add_argument(
90      "--save-profile",
91      action="store",
92      type=argparse.FileType("w"),
93      default=None,
94      help="File path to store the profile to")
95  parser.add_argument(
96      "dex_files", help="dex/jar files", nargs="+", metavar="DEX")
97  return parser.parse_known_args()
98
99
100def get_bcp_runtime_args(additions, image, arch):
101  add_files = map(lambda a: a[0], additions)
102  add_locs = map(lambda a: a[1], additions)
103  if arch != "host32" and arch != "host64":
104    args = [
105        "art/tools/host_bcp.sh",
106        os.path.expandvars(
107            "${{OUT}}/system/framework/oat/{}/services.odex".format(arch)),
108    ]
109    print("Running: {}".format(run_print(args)))
110    print("=START=======================================")
111    res = subprocess.run(args, capture_output=True, text=True)
112    print("=END=========================================")
113    if res.returncode != 0:
114      print("Falling back to ART boot image: {}".format(res))
115      args = [
116          "art/tools/host_bcp.sh",
117          os.path.expandvars(
118              "${{OUT}}/apex/art_boot_images/javalib/{}/boot.oat".format(arch)),
119      ]
120      print("Running: {}".format(run_print(args)))
121      print("=START=======================================")
122      res = subprocess.run(args, capture_output=True, text=True)
123      print("=END=========================================")
124      res.check_returncode()
125    segments = res.stdout.split()
126    def extend_bcp(segment: str):
127      # TODO We should make the bcp have absolute paths.
128      if segment.startswith("-Xbootclasspath:"):
129        return ":".join(itertools.chain((segment,), add_files))
130      elif segment.startswith("-Xbootclasspath-locations:"):
131        return ":".join(itertools.chain((segment,), add_locs))
132      else:
133        return segment
134    return list(map(extend_bcp, segments))
135  else:
136    # Host we just use the bcp locations for both.
137    res = open(
138        os.path.expandvars(
139            "$ANDROID_HOST_OUT/apex/art_boot_images/javalib/{}/boot.oat".format(
140                "x86" if arch == "host32" else "x86_64")), "rb").read()
141    bcp_tag = b"bootclasspath\0"
142    bcp_start = res.find(bcp_tag) + len(bcp_tag)
143    bcp = res[bcp_start:bcp_start + res[bcp_start:].find(b"\0")]
144    img_bcp = bcp.decode()
145    # TODO We should make the str_bcp have absolute paths.
146    str_bcp = ":".join(itertools.chain((img_bcp,), add_files))
147    str_bcp_loc = ":".join(itertools.chain((img_bcp,), add_locs))
148    return [
149        "--runtime-arg", "-Xbootclasspath:{}".format(str_bcp),
150        "--runtime-arg", "-Xbootclasspath-locations:{}".format(str_bcp_loc)
151    ]
152
153
154def fdfile(fd):
155  return "/proc/{}/fd/{}".format(os.getpid(), fd)
156
157
158def get_profile_args(args, location_base):
159  """Handle all the profile file options."""
160  if args.profile_file is None and len(args.profile_line) == 0:
161    return []
162  if args.profile_file:
163    with open(args.profile_file, "rb") as prof:
164      prof_magic = prof.read(4)
165      if prof_magic == b'pro\0':
166        # Looks like the profile-file is a binary profile. Just use it directly
167        return ['--profile-file={}'.format(args.profile_file)]
168  if args.debug_profman:
169    profman_args = ["lldb-server", "g", ":5039", "--", args.profman]
170  else:
171    profman_args = [args.profman]
172  if args.save_profile:
173    prof_out_fd = args.save_profile.fileno()
174    os.set_inheritable(prof_out_fd, True)
175  else:
176    prof_out_fd = os.memfd_create("reference_prof", flags=0)
177  if args.debug_profman:
178    profman_args.append("--reference-profile-file={}".format(
179        fdfile(prof_out_fd)))
180  else:
181    profman_args.append("--reference-profile-file-fd={}".format(prof_out_fd))
182  if args.profile_file:
183    profman_args.append("--create-profile-from={}".format(args.profile_file))
184  else:
185    prof_in_fd = os.memfd_create("input_prof", flags=0)
186    # Why on earth does fdopen take control of the fd and not mention it in the docs.
187    with os.fdopen(os.dup(prof_in_fd), "w") as prof_in:
188      for l in args.profile_line:
189        print(l, file=prof_in)
190    profman_args.append("--create-profile-from={}".format(fdfile(prof_in_fd)))
191  for f in args.dex_files:
192    profman_args.append("--apk={}".format(f))
193    profman_args.append("--dex-location={}".format(
194        os.path.join(location_base, os.path.basename(f))))
195  print("Running: {}".format(run_print(profman_args)))
196  print("=START=======================================")
197  subprocess.run(profman_args, close_fds=False).check_returncode()
198  print("=END=========================================")
199  if args.debug:
200    return ["--profile-file={}".format(fdfile(prof_out_fd))]
201  else:
202    return ["--profile-file={}".format(fdfile(prof_out_fd))]
203
204
205def main():
206  args, extra = parse_args()
207  if args.arch == "host32" or args.arch == "host64":
208    location_base = os.path.expandvars("${ANDROID_HOST_OUT}/framework/")
209    real_arch = "x86" if args.arch == "host32" else "x86_64"
210    boot_image = os.path.expandvars(
211        "$ANDROID_HOST_OUT/apex/art_boot_images/javalib/boot.art")
212    android_root = os.path.expandvars("$ANDROID_HOST_OUT")
213    for f in args.dex_files:
214      extra.append("--dex-location={}".format(
215          os.path.join(location_base, os.path.basename(f))))
216      extra.append("--dex-file={}".format(f))
217  else:
218    location_base = "/system/framework"
219    real_arch = args.arch
220    boot_image = os.path.expandvars(":".join([
221        "${OUT}/apex/art_boot_images/javalib/boot.art",
222        "${OUT}/system/framework/boot-framework.art"
223    ]))
224    android_root = os.path.expandvars("$OUT/system")
225    for f in args.dex_files:
226      extra.append("--dex-location={}".format(
227          os.path.join(location_base, os.path.basename(f))))
228      extra.append("--dex-file={}".format(f))
229  extra += get_bcp_runtime_args(args.add_bcp, boot_image, args.arch)
230  extra += get_profile_args(args, location_base)
231  extra.append("--instruction-set={}".format(real_arch))
232  extra.append("--boot-image={}".format(boot_image))
233  extra.append("--android-root={}".format(android_root))
234  extra += ["--runtime-arg", "-Xms64m", "--runtime-arg", "-Xmx512m"]
235  if args.odex_file is not None:
236    extra.append("--oat-file={}".format(args.odex_file))
237  else:
238    if args.debug:
239      raise Exception("Debug requires a real output file. :(")
240    extra.append("--oat-fd={}".format(os.memfd_create("odex_fd", flags=0)))
241    extra.append("--oat-location={}".format("/tmp/odex_fd.odex"))
242    extra.append("--output-vdex-fd={}".format(
243        os.memfd_create("vdex_fd", flags=0)))
244  pre_args = []
245  if args.debug:
246    pre_args = ["lldb-server", "g", ":5039", "--"]
247  pre_args.append(args.dex2oat)
248  print("Running: {}".format(run_print(pre_args + extra)))
249  print("=START=======================================")
250  subprocess.run(pre_args + extra, close_fds=False).check_returncode()
251  print("=END=========================================")
252
253
254if __name__ == "__main__":
255  main()
256