Component selection using latent Bazel deps

BUILD.bazel

load(":component_selection.bzl", "binary", "component", "library")

# Implementations. Since this is just a demonstration, empty filegroups will do.
filegroup(name = "impl_mock_foo")
filegroup(name = "impl_real_foo")

component(
    name = "mock_foo",
    description = "A: mock, preferred",
    interface = "foo",
    implementation = [":impl_mock_foo"],
)

component(
    name = "real_foo",
    description = "Z: real, only when mock is not available",
    interface = "foo",
    implementation = [":impl_real_foo"],
)

# Some library to show that indirect dependencies work, too
library(name = "lfa", deps = [":real_foo"])
library(name = "lfb", deps = [":mock_foo"])

# A library with all the code in the binary. In real life, this will probably be created from the same macro
# the library is.
library(name = "link", deps = [":lfa", ":lfb"])

binary(name = "binary", link = ":link")

component_selection.bzl

# All the available components in the transitive closure.
# The "components" field is a list of (interface, description, latent dep) triplets
ComponentInfo = provider(fields = ["components"])

# The components chosen to be linked.
# "components" is a list of ConfiguredTarget objects.
LinkInfo = provider(fields = ["components"])

# Merges multiple ComponentInfo objects. Returns a pair of ComponentInfo and
# LatentDepInfo.
# The result can contain duplicates in case the dependency graph is not a tree
# but that's OK because duplicates do not affect the component choice algorithm.
def merge(cis):
    components = []
    latent_deps = []
    for ci in cis:
        components += ci.components
        for c in ci.components:
            latent_deps += c[2]
    return ComponentInfo(components = components), LatentDepInfo(latent = latent_deps)

# Choose a component to link for each interface, Returns a list of
# (interface, chosen latent dep) pairs.
def choose(ci):
    # sort components by interface. This will be map from interfaces to
    # (description, latent dep) tuples.
    by_interface = {}
    for i, d, l in ci.components:
        if i not in by_interface:
            by_interface[i] = []
        by_interface[i].append((d, l))

    # ...then for each interface, choose the implementation whose description is
    # alphabetically the first
    def f(t):
        return t[0]

    return [(k, sorted(v, key = f)[0][1]) for k, v in by_interface.items()]

# The ComponentInfo of a single component is describes that one component.
def _component_impl(ctx):
    latent_deps = ctx.attr.implementation[0][LatentDepInfo].latent
    return [
        DefaultInfo(),
        ComponentInfo(components = [
            (ctx.attr.interface, ctx.attr.description, latent_deps),
        ],
        LatentDepInfo(latent=[ctx.attr.implementation[0][LatentDepInfo].latent]),
    ]

component = rule(
    implementation = _component_impl,
    attrs = {
        "interface": attr.string(),
        "description": attr.string(),
        "implementation": attr.latent_label_list(),
    },
)

# Propagate the (interface, description, latent dep) tuples up the dependency
# graph. Otherwise, do nothing.
def _library_impl(ctx):
    component_info, latent_dep_info = merge([d[ComponentInfo] for d in ctx.attr.deps])
    return [DefaultInfo(), component_info, latent_dep_info]

library = rule(
    implementation = _library_impl,
    attrs = {"deps": attr.label_list()},
)

# Materialize the chosen components and wrap them in a LinkInfo so that the
# binary can access them. ctx.attr._materialized will contain the actual
# analyzed configured targets.
def _link_aspect_impl(target, ctx):
    return [LinkInfo(components = ctx.attr._materialized)]

# Take the components available to the library this aspect is attached to
# then run the choice algorithm on them.
def _materializer(target):
    chosen = choose(target[ComponentInfo])
    return [c[1] for c in chosen]

link_aspect = aspect(
    implementation = _link_aspect_impl,
    attrs = {"_materialized": attr.label_list(materializer = _materializer)},
)

def _binary_impl(ctx):
    # Print the result of linking
    print("Linked: %s" % ctx.attr.link[LinkInfo].components)
    return [DefaultInfo()]
    
binary = rule(
    implementation = _binary_impl,
    # Attach the link aspect to the single dependency
    attrs = {"link": attr.label(aspects = [link_aspect])},
)