protobuf/hpb/bazel/hpb_proto_library.bzl

# Copyright (c) 2024 Google LLC
# All rights reserved.
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file or at
# https://developers.google.com/open-source/licenses/bsd

"""Public rules for using hpb:
  - hpb_proto_library()
"""

load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain", "use_cpp_toolchain")
load("//bazel:upb_proto_library.bzl", "GeneratedSrcsInfo", "UpbWrappedCcInfo", "upb_proto_library_aspect")
load("//bazel/common:proto_common.bzl", "proto_common")

def upb_use_cpp_toolchain():
    return use_cpp_toolchain()

def _filter_none(elems):
    return [e for e in elems if e]

# TODO: b/373443334 - consolidate
def _cc_library_func(ctx, name, hdrs, srcs, copts, dep_ccinfos):
    """Like cc_library(), but callable from rules.

    Args:
      ctx: Rule context.
      name: Unique name used to generate output files.
      hdrs: Public headers that can be #included from other rules.
      srcs: C/C++ source files.
      copts: Additional options for cc compilation.
      dep_ccinfos: CcInfo providers of dependencies we should build/link against.

    Returns:
      CcInfo provider for this compilation.
    """

    compilation_contexts = [info.compilation_context for info in dep_ccinfos]
    linking_contexts = [info.linking_context for info in dep_ccinfos]
    toolchain = find_cpp_toolchain(ctx)
    feature_configuration = cc_common.configure_features(
        ctx = ctx,
        cc_toolchain = toolchain,
        requested_features = ctx.features,
        unsupported_features = ctx.disabled_features,
    )

    (compilation_context, compilation_outputs) = cc_common.compile(
        actions = ctx.actions,
        feature_configuration = feature_configuration,
        cc_toolchain = toolchain,
        name = name,
        srcs = srcs,
        public_hdrs = hdrs,
        user_compile_flags = copts,
        compilation_contexts = compilation_contexts,
    )

    # buildifier: disable=unused-variable
    (linking_context, linking_outputs) = cc_common.create_linking_context_from_compilation_outputs(
        actions = ctx.actions,
        name = name,
        feature_configuration = feature_configuration,
        cc_toolchain = toolchain,
        compilation_outputs = compilation_outputs,
        linking_contexts = linking_contexts,
    )

    return CcInfo(
        compilation_context = compilation_context,
        linking_context = linking_context,
    )

# Dummy rule to expose select() copts to aspects  ##############################

HpbProtoLibraryCoptsInfo = provider(
    "Provides copts for hpb proto targets",
    fields = {
        "copts": "copts for hpb_proto_library()",
    },
)

def hpb_proto_library_copts_impl(ctx):
    return HpbProtoLibraryCoptsInfo(copts = ctx.attr.copts)

hpb_proto_library_copts = rule(
    implementation = hpb_proto_library_copts_impl,
    attrs = {"copts": attr.string_list(default = [])},
)

_UpbCcWrappedCcInfo = provider("Provider for cc_info for hpb", fields = ["cc_info"])
_WrappedCcGeneratedSrcsInfo = provider("Provider for generated sources", fields = ["srcs"])

def _get_lang_toolchain(ctx):
    return ctx.attr._hpb_lang_toolchain[proto_common.ProtoLangToolchainInfo]

def _compile_upb_cc_protos(ctx, proto_info, proto_sources):
    if len(proto_sources) == 0:
        return GeneratedSrcsInfo(srcs = [], hdrs = [])

    srcs = []
    srcs += proto_common.declare_generated_files(
        ctx.actions,
        extension = ".upb.proto.cc",
        proto_info = proto_info,
    )

    hdrs = []
    hdrs += proto_common.declare_generated_files(
        ctx.actions,
        extension = ".upb.proto.h",
        proto_info = proto_info,
    )
    hdrs += proto_common.declare_generated_files(
        ctx.actions,
        extension = ".upb.fwd.h",
        proto_info = proto_info,
    )

    proto_common.compile(
        actions = ctx.actions,
        proto_info = proto_info,
        proto_lang_toolchain_info = _get_lang_toolchain(ctx),
        generated_files = srcs + hdrs,
        experimental_exec_group = "proto_compiler",
    )

    return GeneratedSrcsInfo(srcs = srcs, hdrs = hdrs)

def _upb_cc_proto_rule_impl(ctx):
    if len(ctx.attr.deps) != 1:
        fail("only one deps dependency allowed.")
    dep = ctx.attr.deps[0]

    if _WrappedCcGeneratedSrcsInfo in dep:
        srcs = dep[_WrappedCcGeneratedSrcsInfo].srcs
    else:
        fail("proto_library rule must generate _WrappedCcGeneratedSrcsInfo (aspect should have " +
             "handled this).")

    if _UpbCcWrappedCcInfo in dep:
        cc_info = dep[_UpbCcWrappedCcInfo].cc_info
    elif UpbWrappedCcInfo in dep:
        cc_info = dep[UpbWrappedCcInfo].cc_info
    else:
        fail("proto_library rule must generate UpbWrappedCcInfo or " +
             "_UpbCcWrappedCcInfo (aspect should have handled this).")

    lib = cc_info.linking_context.linker_inputs.to_list()[0].libraries[0]
    files = _filter_none([
        lib.static_library,
        lib.pic_static_library,
        lib.dynamic_library,
    ])
    return [
        DefaultInfo(files = depset(files + srcs.hdrs + srcs.srcs)),
        srcs,
        cc_info,
    ]

def _upb_cc_proto_aspect_impl(target, ctx, cc_provider, file_provider):
    proto_info = target[ProtoInfo]
    files = _compile_upb_cc_protos(ctx, proto_info, proto_info.direct_sources)
    deps = ctx.rule.attr.deps + ctx.attr._upbprotos
    dep_ccinfos = [dep[CcInfo] for dep in deps if CcInfo in dep]
    dep_ccinfos += [dep[UpbWrappedCcInfo].cc_info for dep in deps if UpbWrappedCcInfo in dep]
    dep_ccinfos += [dep[_UpbCcWrappedCcInfo].cc_info for dep in deps if _UpbCcWrappedCcInfo in dep]
    if UpbWrappedCcInfo not in target:
        fail("Target should have UpbWrappedCcInfo provider")
    dep_ccinfos.append(target[UpbWrappedCcInfo].cc_info)
    cc_info = _cc_library_func(
        ctx = ctx,
        name = ctx.rule.attr.name + ".upbprotos",
        hdrs = files.hdrs,
        srcs = files.srcs,
        copts = ctx.attr._ccopts[HpbProtoLibraryCoptsInfo].copts,
        dep_ccinfos = dep_ccinfos,
    )
    return [cc_provider(cc_info = cc_info), file_provider(srcs = files)]

def _upb_cc_proto_library_aspect_impl(target, ctx):
    return _upb_cc_proto_aspect_impl(target, ctx, _UpbCcWrappedCcInfo, _WrappedCcGeneratedSrcsInfo)

_upb_cc_proto_library_aspect = aspect(
    attrs = {
        "_ccopts": attr.label(
            default = "//hpb:hpb_proto_library_copts",
        ),
        "_hpb_lang_toolchain": attr.label(
            default = "//src/google/protobuf/compiler/hpb:toolchain",
        ),
        "_cc_toolchain": attr.label(
            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
        ),
        "_upbprotos": attr.label_list(
            default = [
                # TODO: Add dependencies for cc runtime (absl/string etc..)
                "//upb:generated_cpp_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
                "//hpb:generated_hpb_support",
                "@com_google_absl//absl/log:absl_check",
                "@com_google_absl//absl/strings",
                "@com_google_absl//absl/status:statusor",
                "//hpb:repeated_field",
            ],
        ),
    },
    exec_groups = {
        "proto_compiler": exec_group(),
    },
    implementation = _upb_cc_proto_library_aspect_impl,
    provides = [
        _UpbCcWrappedCcInfo,
        _WrappedCcGeneratedSrcsInfo,
    ],
    required_aspect_providers = [
        UpbWrappedCcInfo,
    ],
    attr_aspects = ["deps"],
    fragments = ["cpp"],
    toolchains = upb_use_cpp_toolchain(),
)

upb_cc_proto_library = rule(
    implementation = _upb_cc_proto_rule_impl,
    attrs = {
        "deps": attr.label_list(
            aspects = [
                upb_proto_library_aspect,
                _upb_cc_proto_library_aspect,
            ],
            allow_rules = ["proto_library"],
            providers = [ProtoInfo],
        ),
        "_ccopts": attr.label(
            default = "//hpb:hpb_proto_library_copts",
        ),
    },
)

hpb_proto_library = upb_cc_proto_library