tools/devshell/contrib/lib/rust/rustdoc-link.py

#!/usr/bin/env fuchsia-vendored-python
# Copyright 2024 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Builds all Fuchsia Rust documentation.

Please use `include_rustdoc = true`, as in `$ fx set --args=include_rustdoc=true
core.x64`.

This script does three things:
1. Runs `fx build` to output Rust documentation to target-specific
   directories
2. Runs rustdoc to render files that depend on multiple crates
   (e.g. the search index)
3. Merges all documentation to $FUCHSIA_BUILD_DIR/docs/rust/doc/

You can skip 1. by passing `--no-build`. You can skip 1., 2., and 3.
by passing `--dry-run`, which instead saves a documentation plan to
$FUCHSIA_BUILD_DIR/docs/rust/actions.json. When you provide both `--no-build`
and `--dry-run`, we will exclude build actions from that documentation plan.

Use `--verbose` and `--quiet` to control this script's stderr.

Some options are documented "test only." These are exclusively for use by
//build/rust/tests/rustdoc-link/*.

If you pass specific labels to rustdoc-link, the script will document only
those specific crates.

This script reads out/default/rust_target_mapping.json, which has a list of
rustdoc targets. This script references the FUCHSIA_DIR and FUCHSIA_BUILD_DIR
environment variables, which are typically provided by the devshell wrapper.

# Examples

Build docs specifically for fuchsia-async, and put the generated docs at
out/default/docs/rust/doc/fuchsia-async/:

`$ fx rustdoc-link //src/lib/fuchsia-async`

Build docs for all Rust crates in the build graph, and put the generated docs
at out/default/docs/rust/doc/:

`$ fx rustdoc-link`

Save a documentation plan at out/default/docs/rust/actions.json describing how
to document every Rust crate in the build graph. This plan will assume that you
have already built the docs and just need to link and merge them:

`$ fx rustdoc-link --dry-run --verbose --no-build`

# More info

See //tools/devshell/contrib/lib/rust/rustdoc_link_actions.py for a description
of the JSON generated by --dry-run.

See //build/rust/rust_auxiliary.gni for a description of how rustdoc targets
are built in GN.

See <https://doc.rust-lang.org/rustdoc/> for the rustdoc book.

See //BUILD.gn for more about the rust_target_mapping.json.
"""


import json
import os
import shutil
import sys
from argparse import (
    SUPPRESS,
    ArgumentParser,
    BooleanOptionalAction,
    Namespace,
    RawTextHelpFormatter,
)
from collections import defaultdict
from dataclasses import asdict, dataclass
from os import environ
from pathlib import Path
from subprocess import CalledProcessError, run
from sys import argv, stderr
from typing import Any, Optional

from rust import HOST_PLATFORM, GnTarget
from rustdoc_link_actions import (
    Action,
    ActionPath,
    BuildAction,
    CopyAction,
    RustdocAction,
    TargetAction,
    VerifyAction,
    ZipAction,
)


@dataclass(frozen=True)
class Metadata:
    """
    Metadata collected for one of our rustdoc targets from the
    rust_target_mapping walk.
    """

    crate_name: str  # according to --extern
    actual_label: str
    rustdoc_label: str
    rustdoc_out_dir: Path
    disable_rustdoc: bool
    original_label: str
    rustdoc_parts_dir: str
    target_is_fuchsia: bool

    def is_shared_target(self) -> bool:
        return self.actual_label.endswith("-shared)")

    @staticmethod
    def parse(m: dict[str, Any], build_dir: Path) -> Optional["Metadata"]:
        if "extern" not in m:
            return None
        _, crate_name, _ = m["extern"].split("=")
        rustdoc_out_dir = Path(build_dir, m["rustdoc_out_dir"])
        return Metadata(
            crate_name=crate_name,
            actual_label=m["actual_label"],
            rustdoc_label=m["rustdoc_label"],
            rustdoc_out_dir=rustdoc_out_dir,
            disable_rustdoc=m["disable_rustdoc"],
            original_label=m["original_label"],
            rustdoc_parts_dir=m["rustdoc_parts_dir"],
            target_is_fuchsia=m["target_is_fuchsia"],
        )


def read_metadata_file(args: Namespace) -> set[Metadata]:
    """Reads rust_target_mapping.json to get metadata"""
    meta = args.rust_target_mapping.read_bytes()
    meta = json.loads(meta)
    return {Metadata.parse(m, args.build_dir) for m in meta} - {None}


def filter_build_labels(
    meta: list[Metadata],
    build_dir: Path,
    user_labels: list[str],
    print_build_labels: bool,
) -> list[Metadata]:
    """
    If the user provides specific labels to build, we should only build those.
    """
    assert (
        len(user_labels) > 0
    ), "only call this method if the user provided specific labels to build"
    user_labels = set(str(GnTarget(l, build_dir)) for l in user_labels)

    filtered_meta = []
    for m in meta:
        possible_labels = {
            str(GnTarget(m.rustdoc_label, build_dir)),
            str(GnTarget(m.original_label, build_dir)),
            str(GnTarget(m.actual_label, build_dir)),
        }
        matching = not possible_labels.isdisjoint(user_labels)
        if matching:
            filtered_meta.append(m)

    if print_build_labels:
        filter_output = sorted(m.rustdoc_label for m in filtered_meta)
        print("building specific labels", *filter_output, file=stderr)
    return filtered_meta


def dedup_crate_name(
    meta: list[Metadata],
    print_collisions: bool,
) -> list[Metadata]:
    """remove metadata for duplicate crate names"""
    ret = defaultdict(list)
    for m in meta:
        ret[m.crate_name, m.target_is_fuchsia].append(m)

    for (name, _target_is_fuchsia), l in ret.items():
        if len(l) > 1 and print_collisions:
            print(f"crate_name collision: {name}", file=stderr)
            for m in sorted(l, key=lambda m: m.actual_label):
                print(f"  {m.actual_label}", file=stderr)

    # Return alphabetically greatest crates because these are somewhat more likely to be the latest
    # version. We can only choose one version of a crate to document due to rustdoc restrictions
    # about having multiple crates with the same name.
    return [max(l, key=lambda m: m.actual_label) for l in ret.values()]


def generate_target_action(
    dst: Path,
    build: bool,
    argfile: Path,
    extra_rustdoc_args: list[str],
    meta: list[Metadata],
    build_dir: Path,
) -> TargetAction:
    include_parts_dir_args = [
        f"--include-parts-dir={m.rustdoc_parts_dir}" for m in meta
    ]
    # let's document an empty crate to avoid calling rustdoc with no input crate
    target_crate_root = "/dev/null"
    flags = [
        f"--out-dir={dst}",
        "--edition=2021",
        "-Zunstable-options",
        "--merge=finalize",
        *include_parts_dir_args,
        *extra_rustdoc_args,
        target_crate_root,
    ]
    argfile.write_text("\n".join(flags))
    # Use a trailing slash dot to copy the contents of the doc
    # directory instead of the directory itself.
    copy_action = CopyAction(
        srcs=sorted(
            ActionPath(f"{m.rustdoc_out_dir.relative_to(build_dir)}/.")
            for m in meta
        ),
        dst=ActionPath(dst.relative_to(build_dir)),
    )
    rustdoc_action = RustdocAction(
        argfile=ActionPath(argfile.relative_to(build_dir))
    )
    if build and len(meta) > 0:
        build_action = BuildAction(labels=sorted(m.rustdoc_label for m in meta))
    else:
        build_action = None
    return TargetAction(
        copy_action=copy_action,
        rustdoc_action=rustdoc_action,
        build_action=build_action,
    )


def execute_build_action(
    build_action: BuildAction,
    build_executable: Path,
    build_dir: Path,
    use_xargs: bool,
) -> None:
    if use_xargs:
        # We use xargs to avoid calling fx build with too many arguments. xargs checks
        # the environment for maximum argument list size and may potentially make
        # several calls to fx build in the case that we have too many to fit in one.
        labels = "".join(f"{l}\0" for l in build_action.labels)
        result = run(
            ["xargs", "--null", build_executable],
            text=True,
            input=labels,
            cwd=build_dir,
        )
    else:
        result = run(
            [build_executable, *build_action.labels],
            cwd=build_dir,
        )
    result.check_returncode()


def execute_rustdoc_action(
    rustdoc_action: RustdocAction, rustdoc_executable: Path, build_dir: Path
) -> None:
    result = run(
        [rustdoc_executable, f"@{rustdoc_action.argfile}"], cwd=build_dir
    )
    result.check_returncode()


def execute_copy_action(
    copy_action: CopyAction,
    use_xargs: bool,
    merge_parallelism: int,
    build_dir: Path,
) -> None:
    if use_xargs:
        # Let's merge the files into our destination. We use xargs here to take
        # advantage of parallelism. It is much (35%) faster than packing everything
        # into a single cp invocation. It is much much (75%) faster than shutil.copytree.
        result = run(
            [
                "xargs",
                "--no-run-if-empty",
                "--null",
                "--max-args=1",
                "--max-procs",
                str(merge_parallelism),
                "cp",
                "--recursive",
                "--force",
                "--target-directory",
                copy_action.dst,
            ],
            cwd=build_dir,
            text=True,
            input="".join(f"{s}\0" for s in copy_action.srcs),
        )
    else:
        result = run(
            [
                "cp",
                "--recursive",
                "--force",
                "--target-directory",
                copy_action.dst,
                *copy_action.srcs,
            ],
            cwd=build_dir,
        )
    result.check_returncode()


def execute_target_action(
    target_action: TargetAction,
    build_executable: Path,
    rustdoc_executable: Path,
    build_dir: Path,
    use_xargs: bool,
    merge_parallelism: int,
    quiet: bool,
) -> None:
    if target_action.build_action is not None:
        execute_build_action(
            target_action.build_action, build_executable, build_dir, use_xargs
        )
    if not quiet:
        print(
            "fx rustdoc-link: running rustdoc to build files that depend on multiple crates",
            file=stderr,
        )
    execute_rustdoc_action(
        target_action.rustdoc_action,
        rustdoc_executable=rustdoc_executable,
        build_dir=build_dir,
    )
    if not quiet:
        dst = Path(build_dir, target_action.copy_action.dst)
        print("fx rustdoc-link: copying docs to", dst, file=stderr)
    execute_copy_action(
        target_action.copy_action,
        use_xargs,
        merge_parallelism,
        build_dir=build_dir,
    )


def execute_zip_action(zip_action: ZipAction) -> None:
    zip_to_without_suffix = Path(zip_action.zip_to).with_suffix("")
    shutil.make_archive(
        str(zip_to_without_suffix), format="zip", root_dir=zip_action.zip_from
    )
    assert Path(
        zip_action.zip_to
    ).is_file(), "should have created .zip with our intended name"


def execute_verify_action(
    verify_action: VerifyAction,
    build_dir: Path,
) -> None:
    result = run(
        [verify_action.executable, *verify_action.args],
        cwd=build_dir,
    )
    try:
        result.check_returncode()
    except CalledProcessError as e:
        print(
            "fx rustdoc-link: The generated documentation appears broken.",
            file=stderr,
        )
        print(
            "fx rustdoc-link: Check stderr of the above command for more.",
            file=stderr,
        )
        print(f"fx rustdoc-link: see: {e}", file=stderr)
        sys.exit(1)


def execute_action(
    action: Action,
    build_executable: Path,
    rustdoc_executable: Path,
    build_dir: Path,
    fuchsia_dir: Path,
    use_xargs: bool,
    merge_parallelism: int,
    quiet: bool,
) -> None:
    execute_target_action(
        action.host_action,
        build_executable=build_executable,
        rustdoc_executable=rustdoc_executable,
        build_dir=build_dir,
        use_xargs=use_xargs,
        merge_parallelism=merge_parallelism,
        quiet=quiet,
    )
    execute_target_action(
        action.fuchsia_action,
        build_executable=build_executable,
        rustdoc_executable=rustdoc_executable,
        build_dir=build_dir,
        use_xargs=use_xargs,
        merge_parallelism=merge_parallelism,
        quiet=quiet,
    )
    if action.zip_action is not None:
        if not quiet:
            zip_to = action.zip_action.zip_to
            print("fx rustdoc-link: zipping docs to", zip_to, file=stderr)

        execute_zip_action(action.zip_action)
    execute_verify_action(
        verify_action=action.verify_action,
        build_dir=build_dir,
    )


def set_arg_defaults(args: Namespace) -> None:
    """
    need to set defaults here because they depend on the fuchsia_dir and build_dir args
    """
    assert (
        args.fuchsia_dir is not None
    ), "must provide fuchsia dir through --fuchsia-dir or FUCHSIA_DIR envirnoment var. This is normally done through the devshell wrapper."
    assert (
        args.build_dir is not None
    ), "must provide build dir through --build-dir or FUCHSIA_BUILD_DIR environment var. This is normally done through the devshell wrapper."

    if args.rust_target_mapping is None:
        args.rust_target_mapping = Path(
            args.build_dir, "rust_target_mapping.json"
        )
    if args.output_base is None:
        args.output_base = Path(args.build_dir) / "docs" / "rust"
    if args.zip_to is not None:
        args.zip_to = Path(args.build_dir, args.zip_to)
    if args.rustdoc_executable is None:
        args.rustdoc_executable = Path(
            args.fuchsia_dir,
            "prebuilt/third_party/rust",
            HOST_PLATFORM,
            "bin/rustdoc",
        )
    if args.build_executable is None:
        args.build_executable = Path(args.fuchsia_dir, "tools/devshell/build")
    if args.extra_rustdoc_arg is None:
        args.extra_rustdoc_arg = []


def make_output_directories(args: Namespace) -> None:
    """This type of operation is often done by the build system, but since
    we are running outside of that we have to do it manually.
    """
    if not args.dry_run:
        # remove the destination to ensure that we always document into a fresh
        # directory
        shutil.rmtree(args.output_base, ignore_errors=True)

    # In addition to rustdocs, extra files are stored in output_base:
    #
    #   * actions.json
    #   * argfiles directory
    #
    args.output_base.mkdir(parents=True, exist_ok=True)

    # Make argfiles directory for `rustdoc @argfiles` invocations.
    (args.output_base / "argfiles").mkdir(exist_ok=True)

    # Make destination for fuchsia-side docs.
    (args.output_base / "doc").mkdir(exist_ok=True)

    # Make destination for host-side docs.
    (args.output_base / "doc" / "host").mkdir(exist_ok=True)


def generate_action(meta: list[Metadata], args: Namespace) -> Action:
    argfiles_dir = args.output_base / "argfiles"
    doc_out_dir = args.output_base / "doc"
    host_action = generate_target_action(
        meta=[m for m in meta if not m.target_is_fuchsia],
        dst=doc_out_dir / "host",
        argfile=Path(argfiles_dir, "host.args"),
        extra_rustdoc_args=args.extra_rustdoc_arg,
        build=args.build,
        build_dir=args.build_dir,
    )
    fuchsia_action = generate_target_action(
        meta=[m for m in meta if m.target_is_fuchsia],
        dst=doc_out_dir,
        argfile=Path(argfiles_dir, "fuchsia.args"),
        extra_rustdoc_args=args.extra_rustdoc_arg,
        build=args.build,
        build_dir=args.build_dir,
    )
    zip_action = (
        ZipAction(
            zip_from=ActionPath(doc_out_dir.relative_to(args.build_dir)),
            zip_to=ActionPath(args.zip_to.relative_to(args.build_dir)),
        )
        if args.zip_to
        else None
    )

    executable = (
        Path(args.fuchsia_dir)
        / "tools"
        / "devshell"
        / "contrib"
        / "lib"
        / "rust"
        / "rustdoc_link_verify.py"
    )
    verify_action = VerifyAction(
        executable=ActionPath(
            os.path.relpath(sys.executable, start=args.build_dir)
        ),
        args=[
            ActionPath(executable.relative_to(args.build_dir)),
            ActionPath(doc_out_dir.relative_to(args.build_dir)),
        ],
    )
    return Action(
        host_action=host_action,
        fuchsia_action=fuchsia_action,
        zip_action=zip_action,
        verify_action=verify_action,
    )


def main(args: Namespace) -> None:
    set_arg_defaults(args)

    make_output_directories(args)

    if args.verbose:
        print("fx rustdoc-link: using arguments", vars(args), file=stderr)

    meta = read_metadata_file(args)

    # remove shared and disabled targets
    meta = [
        m for m in meta if not m.is_shared_target() and not m.disable_rustdoc
    ]

    # filter according to the specific targets that the user wants to build
    if args.build_labels:
        meta = filter_build_labels(
            meta,
            args.build_dir,
            args.build_labels,
            print_build_labels=args.verbose,
        )
    meta = dedup_crate_name(meta, print_collisions=args.verbose)

    action = generate_action(meta, args)

    if args.dry_run:
        save_actions_to = Path(args.build_dir, "docs", "rust", "actions.json")
        if not args.quiet:
            print(
                f"fx rustdoc-link: dry run, saving actions to {save_actions_to}"
            )
        with save_actions_to.open("w") as save_actions_to:
            json.dump(asdict(action), save_actions_to, indent=4)
    else:
        try:
            execute_action(
                action=action,
                build_executable=args.build_executable,
                rustdoc_executable=args.rustdoc_executable,
                build_dir=args.build_dir,
                fuchsia_dir=args.fuchsia_dir,
                use_xargs=args.use_xargs,
                merge_parallelism=args.merge_parallelism,
                quiet=args.quiet,
            )
        except CalledProcessError as e:
            print(
                f"rustdoc-link: Failed to build and link docs. The call to {e.cmd[0]} was not successful.",
                file=stderr,
            )
            print(
                f"rustdoc-link: We attempted to run",
                *e.cmd,
                f"and got code {e.returncode}. Check stdout and stderr above.",
                file=stderr,
            )


def _main_arg_parser() -> ArgumentParser:
    # __doc__ refers to the doc comment at the top of this file. Use the raw formatter
    # so that newlines in the doc comment are newlines in the --help text.
    parser = ArgumentParser(
        description=__doc__,
        formatter_class=RawTextHelpFormatter,
    )

    # we don't want consumers to rely on these being present, since they are
    # implementation details that may be removed at any time.
    help_debug_only = SUPPRESS

    # path directly to the rustdoc executable. Can be relative to the current working directory.
    parser.add_argument(
        "--rustdoc-executable",
        help=help_debug_only,
    )
    parser.add_argument(
        "--build",
        default=True,
        action=BooleanOptionalAction,
        help="run fx build to build extra rustdoc targets. manually run fx build beforehand if you use --no-build",
    )
    parser.add_argument(
        "--quiet",
        default=False,
        action=BooleanOptionalAction,
        help="do not print status messages. failure messages still printed",
    )
    parser.add_argument(
        "--verbose",
        default=False,
        action=BooleanOptionalAction,
        help="add additional context to status messages",
    )
    parser.add_argument(
        "--dry-run",
        default=False,
        action=BooleanOptionalAction,
        help="do not build, copy, run rustdoc, or zip anything. creates $BUILD_DIR/docs/rust/actions.json instead",
    )
    parser.add_argument(
        "build_labels",
        nargs="*",
        help="build only these labels if supplied",
    )
    # Path to fx build program. Defaults to the in-tree source of `fx build`.
    parser.add_argument(
        "--build-executable",
        type=Path,
        help=help_debug_only,
    )
    # How much parallelism to use in scheduling the tasks that merge docs.
    parser.add_argument(
        "--merge-parallelism",
        default=os.cpu_count(),
        type=int,
        help=help_debug_only,
    )
    parser.add_argument(
        "--out-dir",
        help="ignored for compatibility with rust devshell wrapper",
    )
    parser.add_argument(
        "--use-xargs",
        default=shutil.which("xargs") is not None,
        action=BooleanOptionalAction,
        help=help_debug_only,
    )
    parser.add_argument(
        "--rust-target-mapping",
        type=Path,
        help="Path to rust_target_mapping.json, defaults to arg.build_dir/rust_target_mapping.json",
    )
    parser.add_argument(
        "--output-base",
        type=Path,
        help="Base output directory for writing merged doc and an optional zip, defaults to arg.build_dir.",
    )
    parser.add_argument(
        "--zip-to",
        type=Path,
        help="When set, creates a zip file at the specified location.",
    )
    parser.add_argument(
        "--fuchsia-dir",
        default=environ.get("FUCHSIA_DIR"),
        type=str,
        help="test only. fuchsia checkout directory, e.g. ~/fuchsia. typically provided by the devshell wrapper. Can be relative to the current working directory",
    )
    parser.add_argument(
        "--build-dir",
        default=environ.get("FUCHSIA_BUILD_DIR"),
        type=str,
        help="test only. build directory to rebase path to, e.g. $FUCHSIA_DIR/out/default. typically provided by the devshell wrapper. can be relative to the current working directory",
    )
    parser.add_argument(
        "--extra-rustdoc-arg",
        action="append",
        help=help_debug_only,
    )
    # generates the documentation plan relative to this location
    parser.add_argument(
        "--save-actions-to",
        type=Path,
        help=help_debug_only,
    )
    parser.set_defaults(func=main)
    return parser


if __name__ == "__main__":
    parser = _main_arg_parser()
    args = parser.parse_args(argv[1:])
    args.func(args)