internal/linker/link_node_modules.bzl

"""Helper function and aspect to collect first-party packages.

These are used in node rules to link the node_modules before launching a program.
This supports path re-mapping, to support short module names.
See pathMapping doc: https://github.com/Microsoft/TypeScript/issues/5039

This reads the module_root and module_name attributes from rules in
the transitive closure, rolling these up to provide a mapping to the
linker, which uses the mappings to link a node_modules directory for
runtimes to locate all the first-party packages.
"""

# Can't load from //:providers.bzl directly as that introduces a circular dep
load("//internal/providers:external_npm_package_info.bzl", "ExternalNpmPackageInfo")
load("@rules_nodejs//nodejs:providers.bzl", "LinkablePackageInfo")

def _debug(vars, format, args_tuple):
    if "VERBOSE_LOGS" in vars.keys():
        print("[link_node_modules.bzl]", format % args_tuple)

LinkerPackageMappingInfo = provider(
    doc = """Provider capturing package mappings for the linker to consume.""",
    fields = {
        "mappings": """Depset of structs with mapping info.

Each struct has the following fields:
  package_name: The name of the package.
  package_path: The root path of the node_modules under which this package should be linked.
  link_path: The exec path under which the package is available.

Note: The depset may contain multiple entries per (package_name, package_path) pair.
Consumers should handle these duplicated appropriately. The depset uses topological order to ensure
that a target's mappings come before possibly conflicting mappings of its dependencies.""",
        "node_modules_roots": "Depset of node_module roots.",
    },
)

# Traverse 'srcs' in addition so that we can go across a genrule
_MODULE_MAPPINGS_DEPS_NAMES = ["data", "deps", "srcs"]

def add_arg(args, arg):
    """Add an argument

    Args:
        args: either a list or a ctx.actions.Args object
        arg: string arg to append on the end
    """
    if (type(args) == type([])):
        args.append(arg)
    else:
        args.add(arg)

def _detect_conflicts(module_sets, mapping):
    """Verifies that the new mapping does not conflict with existing mappings in module_sets."""
    if mapping.package_path not in module_sets:
        return
    existing_link_path = module_sets[mapping.package_path].get(mapping.package_name)
    if existing_link_path == None:
        return

    # If we're linking to the output tree be tolerant of linking to different
    # output trees since we can have "static" links that come from cfg="exec" binaries.
    # In the future when we static link directly into runfiles without the linker
    # we can remove this logic.
    link_path_segments = mapping.link_path.split("/")
    existing_link_path_segments = existing_link_path.split("/")
    bin_links = len(link_path_segments) >= 3 and len(existing_link_path_segments) >= 3 and link_path_segments[0] == "bazel-out" and existing_link_path_segments[0] == "bazel-out" and link_path_segments[2] == "bin" and existing_link_path_segments[2] == "bin"
    if bin_links:
        compare_link_path = "/".join(link_path_segments[3:]) if len(link_path_segments) > 3 else ""
        compare_existing_link_path = "/".join(existing_link_path_segments[3:]) if len(existing_link_path_segments) > 3 else ""
    else:
        compare_link_path = mapping.link_path
        compare_existing_link_path = existing_link_path
    if compare_link_path != compare_existing_link_path:
        msg = "conflicting mapping: '%s' (package path: %s) maps to conflicting paths '%s' and '%s'" % (mapping.package_name, mapping.package_path, compare_link_path, compare_existing_link_path)
        fail(msg)

def _flatten_to_module_set(mappings_depset):
    """Convert a depset of mapping to a module sets (modules per package package_path).

    The returned dictionary has the following structure:
    {
        "package_path": {
            "package_name": "link_path",
            ...
        },
        ...
    }
    """

    # FIXME: Flattens a depset during the analysis phase. Ideally, this would be done during the
    # execution phase using an Args object.
    flattens_mappings = mappings_depset.to_list()
    module_sets = {}
    for mapping in flattens_mappings:
        _detect_conflicts(module_sets, mapping)
        if mapping.package_path not in module_sets:
            module_sets[mapping.package_path] = {}
        module_sets[mapping.package_path][mapping.package_name] = mapping.link_path
    return module_sets

def write_node_modules_manifest(ctx, extra_data = [], mnemonic = None, link_workspace_root = False):
    """Writes a manifest file read by the linker, containing info about resolving runtime dependencies

    Args:
        ctx: starlark rule execution context
        extra_data: labels to search for npm packages that need to be linked (ctx.attr.deps and ctx.attr.data will always be searched)
        mnemonic: optional action mnemonic, used to differentiate module mapping files from the same rule context
        link_workspace_root: Link the workspace root to the bin_dir to support absolute requires like 'my_wksp/path/to/file'.
            If source files need to be required then they can be copied to the bin_dir with copy_to_bin.
    """

    node_modules_roots = {}

    # Look through data/deps attributes to find the root directories for the third-party node_modules;
    # we'll symlink local "node_modules" to them
    for dep in extra_data + getattr(ctx.attr, "data", []) + getattr(ctx.attr, "deps", []):
        if ExternalNpmPackageInfo in dep:
            path = getattr(dep[ExternalNpmPackageInfo], "path", "")
            workspace = dep[ExternalNpmPackageInfo].workspace
            if path in node_modules_roots:
                other_workspace = node_modules_roots[path]
                if workspace != other_workspace:
                    fail("All npm dependencies at the path '%s' must come from a single workspace. Found '%s' and '%s'." % (path, other_workspace, workspace))
            node_modules_roots[path] = workspace

    direct_mappings = []
    direct_node_modules_roots = []
    if link_workspace_root:
        direct_mappings.append(struct(
            package_name = ctx.workspace_name,
            package_path = "",
            link_path = ctx.bin_dir.path,
        ))
        direct_node_modules_roots.append("")

    transitive_mappings = []
    transitive_node_modules_roots = []
    for dep in extra_data + getattr(ctx.attr, "data", []) + getattr(ctx.attr, "deps", []):
        if not LinkerPackageMappingInfo in dep:
            continue
        transitive_mappings.append(dep[LinkerPackageMappingInfo].mappings)
        transitive_node_modules_roots.append(dep[LinkerPackageMappingInfo].node_modules_roots)

    mappings = depset(direct_mappings, transitive = transitive_mappings, order = "topological")
    module_sets = _flatten_to_module_set(mappings)

    # FIXME: Flattens a depset during the analysis phase. Ideally, this would be done during the
    # execution phase using an Args object.
    linker_node_modules_roots = depset(direct_node_modules_roots, transitive = transitive_node_modules_roots).to_list()
    for node_modules_root in linker_node_modules_roots:
        if node_modules_root not in node_modules_roots:
            node_modules_roots[node_modules_root] = ""

    # Write the result to a file, and use the magic node option --bazel_node_modules_manifest
    # The launcher.sh will peel off this argument and pass it to the linker rather than the program.
    prefix = ctx.label.name
    if mnemonic != None:
        prefix += "_%s" % mnemonic
    modules_manifest = ctx.actions.declare_file("_%s.module_mappings.json" % prefix)
    content = {
        "bin": ctx.bin_dir.path,
        "module_sets": module_sets,
        "roots": node_modules_roots,
        "workspace": ctx.workspace_name,
    }
    ctx.actions.write(modules_manifest, str(content))
    return modules_manifest

def _get_linker_package_mapping_info(target, ctx):
    """Transitively gathers module mappings and node_modules roots from LinkablePackageInfo.

    Args:
      target: target
      ctx: ctx

    Returns:
      A LinkerPackageMappingInfo provider that contains the mappings and roots for the current
      target and all its transitive dependencies.
    """

    transitive_mappings = []
    transitive_node_modules_roots = []
    for name in _MODULE_MAPPINGS_DEPS_NAMES:
        for dep in getattr(ctx.rule.attr, name, []):
            if not LinkerPackageMappingInfo in dep:
                continue
            transitive_mappings.append(dep[LinkerPackageMappingInfo].mappings)
            transitive_node_modules_roots.append(dep[LinkerPackageMappingInfo].node_modules_roots)

    direct_mappings = []
    direct_node_modules_roots = []

    # Look for LinkablePackageInfo mapping in this node
    # LinkablePackageInfo may be provided without a package_name so check for that case as well
    if not LinkablePackageInfo in target:
        _debug(ctx.var, "No LinkablePackageInfo for %s", (target.label))
    elif not target[LinkablePackageInfo].package_name:
        _debug(ctx.var, "No package_name in LinkablePackageInfo for %s", (target.label))
    else:
        linkable_package_info = target[LinkablePackageInfo]
        package_path = linkable_package_info.package_path if hasattr(linkable_package_info, "package_path") else ""
        direct_mappings.append(
            struct(
                package_name = linkable_package_info.package_name,
                package_path = package_path,
                link_path = linkable_package_info.path,
            ),
        )
        _debug(ctx.var, "target %s (package path: %s) adding module mapping %s: %s", (
            target.label,
            package_path,
            linkable_package_info.package_name,
            linkable_package_info.path,
        ))
        direct_node_modules_roots.append(package_path)

    mappings = depset(direct_mappings, transitive = transitive_mappings, order = "topological")
    node_modules_roots = depset(direct_node_modules_roots, transitive = transitive_node_modules_roots)
    return LinkerPackageMappingInfo(
        mappings = mappings,
        node_modules_roots = node_modules_roots,
    )

def _module_mappings_aspect_impl(target, ctx):
    # If the target explicitly provides mapping information, we will not propagate
    # any information. The target already provides explicit mapping information.
    # See details on a concrete use-case: https://github.com/bazelbuild/rules_nodejs/issues/2941.
    if LinkerPackageMappingInfo in target:
        return []

    return [
        _get_linker_package_mapping_info(target, ctx),
    ]

module_mappings_aspect = aspect(
    _module_mappings_aspect_impl,
    attr_aspects = _MODULE_MAPPINGS_DEPS_NAMES,
)