1"""Public rules for using upb protos:
2  - upb_proto_library()
3  - upb_proto_reflection_library()
4"""
5
6load("@bazel_skylib//lib:paths.bzl", "paths")
7load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain")
8load("@rules_proto//proto:defs.bzl", "ProtoInfo")  # copybara:strip_for_google3
9
10# Generic support code #########################################################
11
12_is_bazel = True  # copybara:replace_for_google3 _is_bazel = False
13
14def _get_real_short_path(file):
15    # For some reason, files from other archives have short paths that look like:
16    #   ../com_google_protobuf/google/protobuf/descriptor.proto
17    short_path = file.short_path
18    if short_path.startswith("../"):
19        second_slash = short_path.index("/", 3)
20        short_path = short_path[second_slash + 1:]
21
22    # Sometimes it has another few prefixes like:
23    #   _virtual_imports/any_proto/google/protobuf/any.proto
24    #   benchmarks/_virtual_imports/100_msgs_proto/benchmarks/100_msgs.proto
25    # We want just google/protobuf/any.proto.
26    virtual_imports = "_virtual_imports/"
27    if virtual_imports in short_path:
28        short_path = short_path.split(virtual_imports)[1].split("/", 1)[1]
29    return short_path
30
31def _get_real_root(file):
32    real_short_path = _get_real_short_path(file)
33    return file.path[:-len(real_short_path) - 1]
34
35def _generate_output_file(ctx, src, extension):
36    real_short_path = _get_real_short_path(src)
37    real_short_path = paths.relativize(real_short_path, ctx.label.package)
38    output_filename = paths.replace_extension(real_short_path, extension)
39    ret = ctx.actions.declare_file(output_filename)
40    return ret
41
42def _filter_none(elems):
43    out = []
44    for elem in elems:
45        if elem:
46            out.append(elem)
47    return out
48
49def _cc_library_func(ctx, name, hdrs, srcs, copts, dep_ccinfos):
50    """Like cc_library(), but callable from rules.
51
52    Args:
53      ctx: Rule context.
54      name: Unique name used to generate output files.
55      hdrs: Public headers that can be #included from other rules.
56      srcs: C/C++ source files.
57      dep_ccinfos: CcInfo providers of dependencies we should build/link against.
58
59    Returns:
60      CcInfo provider for this compilation.
61    """
62
63    compilation_contexts = [info.compilation_context for info in dep_ccinfos]
64    linking_contexts = [info.linking_context for info in dep_ccinfos]
65    toolchain = find_cpp_toolchain(ctx)
66    feature_configuration = cc_common.configure_features(
67        ctx = ctx,
68        cc_toolchain = toolchain,
69        requested_features = ctx.features,
70        unsupported_features = ctx.disabled_features,
71    )
72
73    blaze_only_args = {}
74
75    if not _is_bazel:
76        blaze_only_args["grep_includes"] = ctx.file._grep_includes
77
78    (compilation_context, compilation_outputs) = cc_common.compile(
79        actions = ctx.actions,
80        feature_configuration = feature_configuration,
81        cc_toolchain = toolchain,
82        name = name,
83        srcs = srcs,
84        public_hdrs = hdrs,
85        user_compile_flags = copts,
86        compilation_contexts = compilation_contexts,
87        **blaze_only_args
88    )
89    (linking_context, linking_outputs) = cc_common.create_linking_context_from_compilation_outputs(
90        actions = ctx.actions,
91        name = name,
92        feature_configuration = feature_configuration,
93        cc_toolchain = toolchain,
94        compilation_outputs = compilation_outputs,
95        linking_contexts = linking_contexts,
96        **blaze_only_args
97    )
98
99    return CcInfo(
100        compilation_context = compilation_context,
101        linking_context = linking_context,
102    )
103
104# Build setting for whether fasttable code generation is enabled ###############
105
106_FastTableEnabled = provider(
107    fields = {
108        "enabled": "whether fasttable is enabled",
109    },
110)
111
112def fasttable_enabled_impl(ctx):
113    raw_setting = ctx.build_setting_value
114
115    if raw_setting:
116        # TODO(haberman): check that the target CPU supports fasttable.
117        pass
118
119    return _FastTableEnabled(enabled = raw_setting)
120
121upb_fasttable_enabled = rule(
122    implementation = fasttable_enabled_impl,
123    build_setting = config.bool(flag = True),
124)
125
126# Dummy rule to expose select() copts to aspects  ##############################
127
128_UpbProtoLibraryCopts = provider(
129    fields = {
130        "copts": "copts for upb_proto_library()",
131    },
132)
133
134def upb_proto_library_copts_impl(ctx):
135    return _UpbProtoLibraryCopts(copts = ctx.attr.copts)
136
137upb_proto_library_copts = rule(
138    implementation = upb_proto_library_copts_impl,
139    attrs = {"copts": attr.string_list(default = [])},
140)
141
142# upb_proto_library / upb_proto_reflection_library shared code #################
143
144GeneratedSrcsInfo = provider(
145    fields = {
146        "srcs": "list of srcs",
147        "hdrs": "list of hdrs",
148    },
149)
150
151_UpbWrappedCcInfo = provider(fields = ["cc_info"])
152_UpbDefsWrappedCcInfo = provider(fields = ["cc_info"])
153_WrappedGeneratedSrcsInfo = provider(fields = ["srcs"])
154_WrappedDefsGeneratedSrcsInfo = provider(fields = ["srcs"])
155
156def _compile_upb_protos(ctx, generator, proto_info, proto_sources):
157    if len(proto_sources) == 0:
158        return GeneratedSrcsInfo(srcs = [], hdrs = [])
159
160    ext = "." + generator
161    tool = getattr(ctx.executable, "_gen_" + generator)
162    srcs = [_generate_output_file(ctx, name, ext + ".c") for name in proto_sources]
163    hdrs = [_generate_output_file(ctx, name, ext + ".h") for name in proto_sources]
164    transitive_sets = proto_info.transitive_descriptor_sets.to_list()
165    fasttable_enabled = (hasattr(ctx.attr, "_fasttable_enabled") and
166                         ctx.attr._fasttable_enabled[_FastTableEnabled].enabled)
167    codegen_params = "fasttable:" if fasttable_enabled else ""
168    ctx.actions.run(
169        inputs = depset(
170            direct = [proto_info.direct_descriptor_set],
171            transitive = [proto_info.transitive_descriptor_sets],
172        ),
173        tools = [tool],
174        outputs = srcs + hdrs,
175        executable = ctx.executable._protoc,
176        arguments = [
177                        "--" + generator + "_out=" + codegen_params + _get_real_root(srcs[0]),
178                        "--plugin=protoc-gen-" + generator + "=" + tool.path,
179                        "--descriptor_set_in=" + ctx.configuration.host_path_separator.join([f.path for f in transitive_sets]),
180                    ] +
181                    [_get_real_short_path(file) for file in proto_sources],
182        progress_message = "Generating upb protos for :" + ctx.label.name,
183    )
184    return GeneratedSrcsInfo(srcs = srcs, hdrs = hdrs)
185
186def _upb_proto_rule_impl(ctx):
187    if len(ctx.attr.deps) != 1:
188        fail("only one deps dependency allowed.")
189    dep = ctx.attr.deps[0]
190
191    if _WrappedDefsGeneratedSrcsInfo in dep:
192        srcs = dep[_WrappedDefsGeneratedSrcsInfo].srcs
193    elif _WrappedGeneratedSrcsInfo in dep:
194        srcs = dep[_WrappedGeneratedSrcsInfo].srcs
195    else:
196        fail("proto_library rule must generate _WrappedGeneratedSrcsInfo or " +
197             "_WrappedDefsGeneratedSrcsInfo (aspect should have handled this).")
198
199    if _UpbDefsWrappedCcInfo in dep:
200        cc_info = dep[_UpbDefsWrappedCcInfo].cc_info
201    elif _UpbWrappedCcInfo in dep:
202        cc_info = dep[_UpbWrappedCcInfo].cc_info
203    else:
204        fail("proto_library rule must generate _UpbWrappedCcInfo or " +
205             "_UpbDefsWrappedCcInfo (aspect should have handled this).")
206
207    lib = cc_info.linking_context.linker_inputs.to_list()[0].libraries[0]
208    files = _filter_none([
209        lib.static_library,
210        lib.pic_static_library,
211        lib.dynamic_library,
212    ])
213    return [
214        DefaultInfo(files = depset(files + srcs.hdrs + srcs.srcs)),
215        srcs,
216        cc_info,
217    ]
218
219def _upb_proto_aspect_impl(target, ctx, generator, cc_provider, file_provider):
220    proto_info = target[ProtoInfo]
221    files = _compile_upb_protos(ctx, generator, proto_info, proto_info.direct_sources)
222    deps = ctx.rule.attr.deps + getattr(ctx.attr, "_" + generator)
223    dep_ccinfos = [dep[CcInfo] for dep in deps if CcInfo in dep]
224    dep_ccinfos += [dep[_UpbWrappedCcInfo].cc_info for dep in deps if _UpbWrappedCcInfo in dep]
225    dep_ccinfos += [dep[_UpbDefsWrappedCcInfo].cc_info for dep in deps if _UpbDefsWrappedCcInfo in dep]
226    if generator == "upbdefs":
227        if _UpbWrappedCcInfo not in target:
228            fail("Target should have _UpbDefsWrappedCcInfo provider")
229        dep_ccinfos += [target[_UpbWrappedCcInfo].cc_info]
230    cc_info = _cc_library_func(
231        ctx = ctx,
232        name = ctx.rule.attr.name + "." + generator,
233        hdrs = files.hdrs,
234        srcs = files.srcs,
235        copts = ctx.attr._copts[_UpbProtoLibraryCopts].copts,
236        dep_ccinfos = dep_ccinfos,
237    )
238    return [cc_provider(cc_info = cc_info), file_provider(srcs = files)]
239
240def _upb_proto_library_aspect_impl(target, ctx):
241    return _upb_proto_aspect_impl(target, ctx, "upb", _UpbWrappedCcInfo, _WrappedGeneratedSrcsInfo)
242
243def _upb_proto_reflection_library_aspect_impl(target, ctx):
244    return _upb_proto_aspect_impl(target, ctx, "upbdefs", _UpbDefsWrappedCcInfo, _WrappedDefsGeneratedSrcsInfo)
245
246def _maybe_add(d):
247    if not _is_bazel:
248        d["_grep_includes"] = attr.label(
249            allow_single_file = True,
250            cfg = "host",
251            default = "//tools/cpp:grep-includes",
252        )
253    return d
254
255# upb_proto_library() ##########################################################
256
257_upb_proto_library_aspect = aspect(
258    attrs = _maybe_add({
259        "_copts": attr.label(
260            default = "//:upb_proto_library_copts__for_generated_code_only_do_not_use",
261        ),
262        "_gen_upb": attr.label(
263            executable = True,
264            cfg = "host",
265            default = "//upbc:protoc-gen-upb",
266        ),
267        "_protoc": attr.label(
268            executable = True,
269            cfg = "host",
270            default = "@com_google_protobuf//:protoc",
271        ),
272        "_cc_toolchain": attr.label(
273            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
274        ),
275        "_upb": attr.label_list(default = [
276            "//:generated_code_support__only_for_generated_code_do_not_use__i_give_permission_to_break_me",
277            "//:upb",
278        ]),
279        "_fasttable_enabled": attr.label(default = "//:fasttable_enabled"),
280    }),
281    implementation = _upb_proto_library_aspect_impl,
282    provides = [
283        _UpbWrappedCcInfo,
284        _WrappedGeneratedSrcsInfo,
285    ],
286    attr_aspects = ["deps"],
287    fragments = ["cpp"],
288    toolchains = ["@bazel_tools//tools/cpp:toolchain_type"],
289    incompatible_use_toolchain_transition = True,
290)
291
292upb_proto_library = rule(
293    output_to_genfiles = True,
294    implementation = _upb_proto_rule_impl,
295    attrs = {
296        "deps": attr.label_list(
297            aspects = [_upb_proto_library_aspect],
298            allow_rules = ["proto_library"],
299            providers = [ProtoInfo],
300        ),
301    },
302)
303
304# upb_proto_reflection_library() ###############################################
305
306_upb_proto_reflection_library_aspect = aspect(
307    attrs = _maybe_add({
308        "_copts": attr.label(
309            default = "//:upb_proto_library_copts__for_generated_code_only_do_not_use",
310        ),
311        "_gen_upbdefs": attr.label(
312            executable = True,
313            cfg = "host",
314            default = "//upbc:protoc-gen-upbdefs",
315        ),
316        "_protoc": attr.label(
317            executable = True,
318            cfg = "host",
319            default = "@com_google_protobuf//:protoc",
320        ),
321        "_cc_toolchain": attr.label(
322            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
323        ),
324        "_upbdefs": attr.label_list(
325            default = [
326                "//:upb",
327                "//:reflection",
328            ],
329        ),
330    }),
331    implementation = _upb_proto_reflection_library_aspect_impl,
332    provides = [
333        _UpbDefsWrappedCcInfo,
334        _WrappedDefsGeneratedSrcsInfo,
335    ],
336    required_aspect_providers = [
337        _UpbWrappedCcInfo,
338        _WrappedGeneratedSrcsInfo,
339    ],
340    attr_aspects = ["deps"],
341    fragments = ["cpp"],
342    toolchains = ["@bazel_tools//tools/cpp:toolchain_type"],
343    incompatible_use_toolchain_transition = True,
344)
345
346upb_proto_reflection_library = rule(
347    output_to_genfiles = True,
348    implementation = _upb_proto_rule_impl,
349    attrs = {
350        "deps": attr.label_list(
351            aspects = [
352                _upb_proto_library_aspect,
353                _upb_proto_reflection_library_aspect,
354            ],
355            allow_rules = ["proto_library"],
356            providers = [ProtoInfo],
357        ),
358    },
359)
360