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, *args):
    if "VERBOSE_LOGS" in vars.keys():
        print("[link_node_modules.bzl]", *args)

LinkerPackageMappingInfo = provider(
    doc = """Provider capturing package mappings for the linker to consume.""",
    fields = {
        "mappings": "Dictionary of mappings. Maps package names to an exec path.",
    },
)

# 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 _link_mapping(label, mappings, k, v):
    # Check that two package name mapping do not map to two different source paths
    k_segs = k.split(":")
    package_name = k_segs[0]
    package_path = k_segs[1] if len(k_segs) > 1 else ""
    link_path = v

    for iter_key, iter_values in mappings.items():
        # Map key is of format "package_name:package_path"
        # Map values are of format [deprecated, link_path]
        iter_segs = iter_key.split(":")
        iter_package_name = iter_segs[0]
        iter_package_path = iter_segs[1] if len(iter_segs) > 1 else ""
        iter_source_path = iter_values
        if package_name == iter_package_name and package_path == iter_package_path:
            # 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 = link_path.split("/")
            iter_source_path_segments = iter_source_path.split("/")
            bin_links = len(link_path_segments) >= 3 and len(iter_source_path_segments) >= 3 and link_path_segments[0] == "bazel-out" and iter_source_path_segments[0] == "bazel-out" and link_path_segments[2] == "bin" and iter_source_path_segments[2] == "bin"
            if bin_links:
                compare_link_path = "/".join(link_path_segments[3:]) if len(link_path_segments) > 3 else ""
                compare_iter_source_path = "/".join(iter_source_path_segments[3:]) if len(iter_source_path_segments) > 3 else ""
            else:
                compare_link_path = link_path
                compare_iter_source_path = iter_source_path
            if compare_link_path != compare_iter_source_path:
                msg = "conflicting mapping at '%s': '%s' and '%s' map to conflicting %s and %s" % (label, k, iter_key, compare_link_path, compare_iter_source_path)
                fail(msg)

    return True

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.
    """

    mappings = {ctx.workspace_name: ctx.bin_dir.path} if link_workspace_root else {}
    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

    # Look through data/deps attributes to find first party deps to link
    for dep in extra_data + getattr(ctx.attr, "data", []) + getattr(ctx.attr, "deps", []):
        if not LinkerPackageMappingInfo in dep:
            continue

        for k, v in dep[LinkerPackageMappingInfo].mappings.items():
            map_key_split = k.split(":")
            package_name = map_key_split[0]
            package_path = map_key_split[1] if len(map_key_split) > 1 else ""
            if package_path not in node_modules_roots:
                node_modules_roots[package_path] = ""
            if _link_mapping(dep.label, mappings, k, v):
                _debug(ctx.var, "Linking %s: %s" % (k, v))
                mappings[k] = v

    # Convert mappings to a module sets (modules per package package_path)
    # {
    #     "package_path": {
    #             "package_name": "link_path",
    #             ...
    #     },
    #     ...
    # }
    module_sets = {}
    for k, v in mappings.items():
        map_key_split = k.split(":")
        package_name = map_key_split[0]
        package_path = map_key_split[1] if len(map_key_split) > 1 else ""
        link_path = v
        if package_path not in module_sets:
            module_sets[package_path] = {}
        module_sets[package_path][package_name] = link_path

    # 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_module_mappings(target, ctx):
    """Gathers module mappings from LinkablePackageInfo which maps "package_name:package_path" to link_path.

    Args:
      target: target
      ctx: ctx

    Returns:
        Returns module mappings of shape:
        { "package_name:package_path": link_path, ... }
    """
    mappings = {}

    # Propagate transitive mappings
    for name in _MODULE_MAPPINGS_DEPS_NAMES:
        for dep in getattr(ctx.rule.attr, name, []):
            if not LinkerPackageMappingInfo in dep:
                continue

            for k, v in dep[LinkerPackageMappingInfo].mappings.items():
                if _link_mapping(target.label, mappings, k, v):
                    _debug(ctx.var, "target %s propagating module mapping %s: %s" % (dep.label, k, v))
                    mappings[k] = v

    # Look for LinkablePackageInfo mapping in this node
    if not LinkablePackageInfo in target:
        # No mappings contributed here, short-circuit with the transitive ones we collected
        _debug(ctx.var, "No LinkablePackageInfo for", target.label)
        return mappings

    linkable_package_info = target[LinkablePackageInfo]

    # LinkablePackageInfo may be provided without a package_name so check for that case as well
    if not linkable_package_info.package_name:
        # No mappings contributed here, short-circuit with the transitive ones we collected
        _debug(ctx.var, "No package_name in LinkablePackageInfo for", target.label)
        return mappings

    package_path = linkable_package_info.package_path if hasattr(linkable_package_info, "package_path") else ""
    map_key = "%s:%s" % (linkable_package_info.package_name, package_path)
    map_value = linkable_package_info.path

    if _link_mapping(target.label, mappings, map_key, map_value):
        _debug(ctx.var, "target %s adding module mapping %s: %s" % (target.label, map_key, map_value))
        mappings[map_key] = map_value

    # Returns mappings of shape:
    # {
    #   "package_name:package_path": link_path,
    #   ...
    # }
    return mappings

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 [
        LinkerPackageMappingInfo(mappings = _get_module_mappings(target, ctx)),
    ]

module_mappings_aspect = aspect(
    _module_mappings_aspect_impl,
    attr_aspects = _MODULE_MAPPINGS_DEPS_NAMES,
)