1# Copyright (C) 2023 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14import argparse
15import functools
16import logging
17import os
18import re
19import shutil
20import uuid
21from pathlib import Path
22from string import Template
23from typing import Callable, Generator, Iterable
24from typing import NewType
25from typing import Optional
26from typing import TextIO
27
28import cuj
29import util
30from cuj import src
31from go_allowlists import GoAllowlistManipulator
32
33_ALLOWLISTS = "build/soong/android/allowlists/allowlists.go"
34
35ModuleType = NewType("ModuleType", str)
36ModuleName = NewType("ModuleName", str)
37
38Filter = Callable[[ModuleType, ModuleName], bool]
39
40
41def module_defs(src_lines: TextIO) -> Generator[tuple[ModuleType, str], None, None]:
42    """
43    Split `scr_lines` (an Android.bp file) into module definitions and discard
44    everything else (e.g. top level comments and assignments)
45    Assumes that the Android.bp file is properly formatted, specifically,
46    for any module definition:
47     1. the first line matches `start_pattern`, e.g. `cc_library {`
48     2. the last line matches a closing curly brace, i.e. '}'
49    """
50    start_pattern = re.compile(r"^(?P<module_type>\w+)\s*\{\s*$")
51    module_type: Optional[ModuleType] = None
52    buffer = ""
53
54    def in_module_def() -> bool:
55        return buffer != ""
56
57    for line in src_lines:
58        # NB: line includes ending newline
59        line = line.replace("$", "$$")  # escape Templating meta-char '$'
60        if not in_module_def():
61            m = start_pattern.match(line)
62            if m:
63                module_type = ModuleType(m.group("module_type"))
64                buffer = line
65        else:
66            buffer += line
67            if line.rstrip() == "}":
68                assert in_module_def()
69                # end of module definition
70                yield module_type, buffer
71                module_type = None
72                buffer = ""
73
74
75def type_in(*module_types: str) -> Filter:
76    def f(t: ModuleType, _: ModuleName) -> bool:
77        return t in module_types
78
79    return f
80
81
82def name_in(*module_names: str) -> Filter:
83    def f(_: ModuleType, n: ModuleName) -> bool:
84        return n in module_names
85
86    return f
87
88
89def _modify_genrule_template(module_name: ModuleName, module_def: str) -> Optional[str]:
90    # assume `out` only occurs as top-level property of a module
91    # assume `out` is always a singleton array
92    p = re.compile(r'[\n\r]\s+out\s*:\s*\[[\n\r]*\s*"[^"]+(?=")', re.MULTILINE)
93    g = p.search(module_def)
94    if g is None:
95        logging.debug('Could not find "out" for "%s"', module_name)
96        return None
97    index = g.end()
98    return f"{module_def[: index]}-${{suffix}}{module_def[index:]}"
99
100
101def _extract_templates_helper(
102    src_lines: TextIO, f: Filter
103) -> dict[ModuleName, Template]:
104    """
105    For `src_lines` from an Android.bp file, find modules that satisfy the
106    Filter `f` and for each such mach return a "template" text that facilitates
107    changing the module's name.
108    """
109    # assume `name` only occurs as top-level property of a module
110    name_pattern = re.compile(r'[\n\r]\s+name:\s*"(?P<name>[^"]+)(?=")', re.MULTILINE)
111    result = dict[ModuleName, Template]()
112    for module_type, module_def in module_defs(src_lines):
113        m = name_pattern.search(module_def)
114        if not m:
115            continue
116        module_name = ModuleName(m.group("name"))
117        if module_name in result:
118            logging.debug(
119                f"{module_name} already exists thus " f"ignoring {module_type}"
120            )
121            continue
122        if not f(module_type, module_name):
123            continue
124        i = m.end()
125        module_def = f"{module_def[: i]}-${{suffix}}{module_def[i:]}"
126        if module_type == ModuleType("genrule"):
127            module_def = _modify_genrule_template(module_name, module_def)
128            if module_def is None:
129                continue
130        result[module_name] = Template(module_def)
131    return result
132
133
134def _extract_templates(
135    bps: dict[Path, Filter]
136) -> dict[Path, dict[ModuleName, Template]]:
137    """
138    If any key is a directory instead of an Android.bp file, expand it is as if it
139    were the glob pattern $key/**/Android.bp, i.e. replace it with all Android.bp
140    files under its tree.
141    """
142    bp2templates = dict[Path, dict[ModuleName, Template]]()
143    with open(src(_ALLOWLISTS), "rt") as af:
144        go_allowlists = GoAllowlistManipulator(af.readlines())
145        alwaysconvert = go_allowlists.locate("Bp2buildModuleAlwaysConvertList")
146
147    def maybe_register(bp: Path):
148        with open(bp, "rt") as src_lines:
149            templates = _extract_templates_helper(src_lines, fltr)
150            if not go_allowlists.is_dir_allowed(bp.parent):
151                templates = {n: v for n, v in templates.items() if n in alwaysconvert}
152            if len(templates) == 0:
153                logging.debug("No matches in %s", k)
154            else:
155                bp2templates[bp] = bp2templates.get(bp, {}) | templates
156
157    for k, fltr in bps.items():
158        if k.name == "Android.bp":
159            maybe_register(k)
160        for root, _, _ in os.walk(k):
161            if Path(root).is_relative_to(util.get_out_dir()):
162                continue
163            file = Path(root).joinpath("Android.bp")
164            if file.exists():
165                maybe_register(file)
166
167    return bp2templates
168
169
170@functools.cache
171def _back_up_path() -> Path:
172    #  a directory to back up files that these CUJs change
173    return util.get_out_dir().joinpath("clone-cuj-backup")
174
175
176def _backup(bps: Iterable[Path]):
177    # if first cuj_step then back up files to restore later
178    if _back_up_path().exists():
179        raise RuntimeError(
180            f"{_back_up_path()} already exists. "
181            f"Did you kill a previous cuj run? "
182            f"Delete {_back_up_path()} and revert changes to "
183            f"allowlists.go and Android.bp files"
184        )
185    for bp in bps:
186        src_path = bp.relative_to(util.get_top_dir())
187        bak_file = _back_up_path().joinpath(src_path)
188        os.makedirs(os.path.dirname(bak_file))
189        shutil.copy(bp, bak_file)
190    src_allowlists = src(_ALLOWLISTS)
191    bak_allowlists = _back_up_path().joinpath(_ALLOWLISTS)
192    os.makedirs(os.path.dirname(bak_allowlists))
193    shutil.copy(src_allowlists, bak_allowlists)
194
195
196def _restore():
197    src(_ALLOWLISTS).touch(exist_ok=True)
198    for root, _, files in os.walk(_back_up_path()):
199        for file in files:
200            bak_file = Path(root).joinpath(file)
201            bak_path = bak_file.relative_to(_back_up_path())
202            src_file = util.get_top_dir().joinpath(bak_path)
203            shutil.copy(bak_file, src_file)
204            # touch to update mtime; ctime is ignored by ninja
205            src_file.touch(exist_ok=True)
206
207
208def _bz_counterpart(bp: Path) -> Path:
209    return (
210        util.get_out_dir()
211        .joinpath("soong", "bp2build", bp.relative_to(util.get_top_dir()))
212        .with_name("BUILD.bazel")
213    )
214
215
216def _make_clones(bp2templates: dict[Path, dict[ModuleName, Template]], n: int):
217    r = f"{str(uuid.uuid4()):.6}"  # cache-busting
218    source_count = 0
219    output = ["\n"]
220
221    with open(src(_ALLOWLISTS), "rt") as f:
222        go_allowlists = GoAllowlistManipulator(f.readlines())
223        mixed_build_enabled_list = go_allowlists.locate("ProdMixedBuildsEnabledList")
224        alwaysconvert = go_allowlists.locate("Bp2buildModuleAlwaysConvertList")
225
226    def _allow():
227        if name not in mixed_build_enabled_list:
228            mixed_build_enabled_list.prepend([name])
229        mixed_build_enabled_list.prepend(clones)
230
231        if name in alwaysconvert:
232            alwaysconvert.prepend(clones)
233
234    for bp, n2t in bp2templates.items():
235        source_count += len(n2t)
236        output.append(
237            f"{n:5,}X{len(n2t):2,} modules = {n * len(n2t):+5,} "
238            f"in {bp.relative_to(util.get_top_dir())}"
239        )
240        with open(bp, "a") as f:
241            for name, t in n2t.items():
242                clones = []
243                for n in range(1, n + 1):
244                    suffix = f"{r}-{n:05d}"
245                    f.write(t.substitute(suffix=suffix))
246                    clones.append(ModuleName(f"{name}-{suffix}"))
247                _allow()
248
249    with open(src(_ALLOWLISTS), "wt") as f:
250        f.writelines(go_allowlists.lines)
251
252    logging.info(
253        f"Cloned {n:,}X{source_count:,} modules = {n * source_count:+,} "
254        f"in {len(bp2templates)} Android.bp files"
255    )
256    logging.debug("\n".join(output))
257
258
259def _display_sizes():
260    file_count = 0
261    orig_tot = 0
262    curr_tot = 0
263    output = ["\n"]
264    for root, _, files in os.walk(_back_up_path()):
265        file_count += len(files)
266        for file in files:
267            backup_file = Path(root).joinpath(file)
268            common_path = backup_file.relative_to(_back_up_path())
269            source_file = util.get_top_dir().joinpath(common_path)
270            curr_size = os.stat(source_file).st_size
271            curr_tot += curr_size
272            orig_size = os.stat(backup_file).st_size
273            orig_tot += orig_size
274            output.append(
275                f"{orig_size:7,} {curr_size - orig_size :+5,} => {curr_size:9,} "
276                f"bytes {source_file.relative_to(util.get_top_dir())}"
277            )
278            if file == "Android.bp":
279                bz = _bz_counterpart(source_file)
280                output.append(
281                    f"{os.stat(bz).st_size:8,} bytes "
282                    f"$OUTDIR/{bz.relative_to(util.get_out_dir())}"
283                )
284    logging.info(
285        f"Affected {file_count} files {orig_tot:,} "
286        f"{curr_tot - orig_tot:+,} => {curr_tot:,} bytes"
287    )
288    logging.debug("\n".join(output))
289
290
291def _name_cuj(count: int, module_count: int, bp_count: int) -> str:
292    match module_count:
293        case 1:
294            name = f"{count}"
295        case _:
296            name = f"{count}x{module_count}"
297    if bp_count > 1:
298        name = f"{name}({bp_count} files)"
299    return name
300
301
302class Clone(cuj.CujGroup):
303    def __init__(self, group_name: str, bps: dict[Path, Filter]):
304        super().__init__(group_name)
305        self.bps = bps
306
307    def get_steps(self) -> Iterable[cuj.CujStep]:
308        bp2templates = _extract_templates(self.bps)
309        bp_count = len(bp2templates)
310        if bp_count == 0:
311            raise RuntimeError(f"No eligible module to clone in {self.bps.keys()}")
312        module_count = sum(len(templates) for templates in bp2templates.values())
313
314        if "CLONE" in os.environ:
315            counts = [int(s) for s in os.environ["CLONE"].split(",")]
316        else:
317            counts = [1, 100, 200, 300, 400]
318            logging.info(
319                f'Will clone {",".join(str(i) for i in counts)} in cujs. '
320                f"You may specify alternative counts with CLONE env var, "
321                f"e.g. CLONE = 1,10,100,1000"
322            )
323        first_bp = next(iter(bp2templates.keys()))
324
325        def modify_bp():
326            with open(first_bp, mode="a") as f:
327                f.write(f"//post clone modification {uuid.uuid4()}\n")
328
329        steps: list[cuj.CujStep] = []
330        for i, count in enumerate(counts):
331            base_name = _name_cuj(count, module_count, bp_count)
332            steps.append(
333                cuj.CujStep(
334                    verb=base_name,
335                    apply_change=cuj.sequence(
336                        functools.partial(_backup, bp2templates.keys())
337                        if i == 0
338                        else _restore,
339                        functools.partial(_make_clones, bp2templates, count),
340                    ),
341                    verify=_display_sizes,
342                )
343            )
344            steps.append(cuj.CujStep(verb=f"bp aft {base_name}", apply_change=modify_bp))
345            if i == len(counts) - 1:
346                steps.append(
347                    cuj.CujStep(
348                        verb="revert",
349                        apply_change=cuj.sequence(
350                            _restore, lambda: shutil.rmtree(_back_up_path())
351                        ),
352                        verify=_display_sizes,
353                    )
354                )
355        return steps
356
357
358def main():
359    """
360    provided only for manual run;
361    use incremental_build.sh to invoke the cuj instead
362    """
363    p = argparse.ArgumentParser()
364    p.add_argument(
365        "--module",
366        "-m",
367        default="adbd",
368        help="name of the module to clone; default=%(default)s",
369    )
370    p.add_argument(
371        "--count",
372        "-n",
373        default=1,
374        type=int,
375        help="number of times to clone; default: %(default)s",
376    )
377    adb_bp = util.get_top_dir().joinpath("packages/modules/adb/Android.bp")
378    p.add_argument(
379        "androidbp",
380        nargs="?",
381        default=adb_bp,
382        type=Path,
383        help="absolute path to Android.bp file; default=%(default)s",
384    )
385    options = p.parse_args()
386    _make_clones(
387        _extract_templates({options.androidbp: name_in(options.module)}), options.count
388    )
389    logging.warning("Changes made to your source tree; TIP: `repo status`")
390
391
392if __name__ == "__main__":
393    logging.root.setLevel(logging.INFO)
394    main()
395