Migrate proto_descriptor_set

bazel/BUILD.bazel

@@ -64,6 +64,12 @@ bzl_library(
     deps = ["//bazel/private:upb_proto_library_internal_bzl"],
 )
 
+bzl_library(
+    name = "proto_descriptor_set_bzl",
+    srcs = ["proto_descriptor_set.bzl"],
+    visibility = ["//visibility:public"],
+)
+
 # The data in this target is exposed in //bazel/private:for_bazel_tests
 filegroup(
     name = "for_bazel_tests",

bazel/private/file_concat/BUILD.bazel

@@ -0,0 +1,7 @@
+cc_binary(
+    name = "file_concat",
+    srcs = [
+        "main.cc",
+    ],
+    visibility = ["//bazel:__subpackages__"],
+)

bazel/private/file_concat/main.cc

@@ -0,0 +1,65 @@
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <fstream>
+#include <iostream>
+#include <string>
+
+namespace {
+
+constexpr size_t kBufferSize = 4096;  // 4kB
+
+// Return codes.
+constexpr int kOk = 0;
+constexpr int kUsageError = 1;
+constexpr int kIOError = 2;
+
+}  // namespace
+
+int main(int argc, const char* argv[]) {
+  if (argc < 2) {
+    std::cout << "Usage: " << argv[0] << " <output> <inputs...>" << std::endl;
+    return kUsageError;
+  }
+
+  std::string output_path(argv[1]);
+  std::ofstream output(output_path, std::ofstream::binary);
+  if (!output) {
+    std::cerr << "Could not open output file " << output_path << std::endl;
+    return kIOError;
+  }
+
+  for (int i = 2; i < argc; i++) {
+    std::string input_path(argv[i]);
+    std::ifstream input(input_path, std::ifstream::binary);
+    if (!input) {
+      std::cerr << "Could not open input file " << output_path << std::endl;
+      return kIOError;
+    }
+
+    char buffer[kBufferSize];
+    while (input) {
+      if (!input.read(buffer, kBufferSize) && !input.eof()) {
+        std::cerr << "Error reading from " << input_path << std::endl;
+        return kIOError;
+      }
+      if (!output.write(buffer, input.gcount())) {
+        std::cerr << "Error writing to " << output_path << std::endl;
+        return kIOError;
+      }
+    }
+  }
+
+  return kOk;
+}

bazel/proto_descriptor_set.bzl

@@ -0,0 +1,69 @@
+# Copyright 2020 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A rule for generating a `FileDescriptorSet` with all transitive dependencies.
+
+This module contains the definition of `proto_descriptor_set`, a rule that
+collects all `FileDescriptorSet`s from its transitive dependencies and generates
+a single `FileDescriptorSet` containing all the `FileDescriptorProto` from them.
+"""
+
+load("//bazel/common:proto_info.bzl", "ProtoInfo")
+
+def _proto_descriptor_set_impl(ctx):
+    args = ctx.actions.args()
+
+    output = ctx.actions.declare_file("{}.pb".format(ctx.attr.name))
+    args.add(output)
+
+    descriptor_sets = depset(
+        transitive = [dep[ProtoInfo].transitive_descriptor_sets for dep in ctx.attr.deps],
+    )
+    args.add_all(descriptor_sets)
+
+    ctx.actions.run(
+        executable = ctx.executable._file_concat,
+        mnemonic = "ConcatFileDescriptorSet",
+        inputs = descriptor_sets,
+        outputs = [output],
+        arguments = [args],
+    )
+
+    return [
+        DefaultInfo(
+            files = depset([output]),
+            runfiles = ctx.runfiles(files = [output]),
+        ),
+    ]
+
+proto_descriptor_set = rule(
+    implementation = _proto_descriptor_set_impl,
+    attrs = {
+        "deps": attr.label_list(
+            mandatory = False,
+            providers = [ProtoInfo],
+            doc = """
+Sequence of `ProtoInfo`s to collect `FileDescriptorSet`s from.
+""".strip(),
+        ),
+        "_file_concat": attr.label(
+            default = "//bazel/private/file_concat:file_concat",
+            executable = True,
+            cfg = "exec",
+        ),
+    },
+    doc = """
+Collects all `FileDescriptorSet`s from `deps` and combines them into a single
+`FileDescriptorSet` containing all the `FileDescriptorProto`.
+""".strip(),
+)

bazel/tests/BUILD

@@ -4,6 +4,8 @@ load(":proto_common_check_collocated_tests.bzl", "proto_common_check_collocated_
 load(":proto_common_compile_tests.bzl", "proto_common_compile_test_suite")
 load(":proto_common_declare_generated_files_tests.bzl", "proto_common_declare_generated_files_test_suite")
 load(":proto_common_should_generate_tests.bzl", "proto_common_should_generate_test_suite")
+load("//bazel:proto_descriptor_set.bzl", "proto_descriptor_set")
+load("//bazel:proto_library.bzl", "proto_library")
 
 package(default_applicable_licenses = ["//:license"])
 
@@ -18,3 +20,50 @@ proto_common_check_collocated_test_suite(name = "proto_common_check_collocated_t
 bazel_proto_library_test_suite(name = "bazel_proto_library_test_suite")
 
 java_proto_library_test_suite(name = "java_proto_library_test_suite")
+
+
+proto_library(
+    name = "empty_proto_library",
+)
+
+proto_descriptor_set(
+    name = "no_protos",
+    deps = [
+        ":empty_proto_library",
+    ],
+)
+
+proto_descriptor_set(
+    name = "well_known_protos",
+    deps = [
+        "//:any_proto",
+        "//:api_proto",
+        "//:compiler_plugin_proto",
+        "//:descriptor_proto",
+        "//:duration_proto",
+        "//:empty_proto",
+        "//:field_mask_proto",
+        "//:source_context_proto",
+        "//:struct_proto",
+        "//:timestamp_proto",
+        "//:type_proto",
+        "//:wrappers_proto",
+    ],
+)
+
+cc_test(
+    name = "proto_descriptor_set_test",
+    srcs = [
+        "proto_descriptor_set_test.cc",
+    ],
+    data = [
+        ":no_protos",
+        ":well_known_protos",
+    ],
+    deps = [
+        "@bazel_tools//tools/cpp/runfiles",
+        "@googletest//:gtest",
+        "@googletest//:gtest_main",
+        "//:protobuf",
+    ],
+)

bazel/tests/proto_descriptor_set_test.cc

@@ -0,0 +1,108 @@
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <algorithm>
+#include <fstream>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "tools/cpp/runfiles/runfiles.h"
+#include "google/protobuf/descriptor.pb.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+
+using bazel::tools::cpp::runfiles::Runfiles;
+using google::protobuf::FileDescriptorProto;
+using google::protobuf::FileDescriptorSet;
+
+namespace rulesproto {
+constexpr char kWorkspaceRlocation[] = "protobuf/";
+constexpr char kWorkspaceRlocationBzlmod[] = "_main/";
+
+namespace {
+
+std::string GetRlocation(const std::string& file) {
+  static std::unique_ptr<Runfiles> runfiles(Runfiles::CreateForTest());
+  std::string path =
+      runfiles->Rlocation(rulesproto::kWorkspaceRlocation + file);
+  std::ifstream input(path, std::ifstream::binary);
+  if (!input) {
+    path = runfiles->Rlocation(rulesproto::kWorkspaceRlocationBzlmod + file);
+  }
+  return path;
+}
+
+template <typename T, typename K>
+bool Contains(const T& container, const K& key) {
+  return container.find(key) != container.end();
+}
+
+std::vector<std::string> ReadFileDescriptorSet(const std::string& path) {
+  std::ifstream input(path, std::ifstream::binary);
+  EXPECT_TRUE(input) << "Could not open " << path;
+
+  FileDescriptorSet file_descriptor_set;
+  EXPECT_TRUE(file_descriptor_set.ParseFromIstream(&input));
+
+  std::set<std::string> unordered_proto_files;
+  for (FileDescriptorProto file_descriptor : file_descriptor_set.file()) {
+    EXPECT_FALSE(Contains(unordered_proto_files, file_descriptor.name()))
+        << "Already saw " << file_descriptor.name();
+    unordered_proto_files.insert(file_descriptor.name());
+  }
+
+  std::vector<std::string> proto_files(unordered_proto_files.begin(),
+                                       unordered_proto_files.end());
+  std::sort(proto_files.begin(), proto_files.end());
+  return proto_files;
+}
+
+void AssertFileDescriptorSetContains(
+    const std::string& path,
+    const std::vector<std::string>& expected_proto_files) {
+  std::vector<std::string> actual_proto_files =
+      ReadFileDescriptorSet(GetRlocation(path));
+  EXPECT_THAT(actual_proto_files,
+              ::testing::IsSupersetOf(expected_proto_files));
+}
+
+}  // namespace
+
+TEST(ProtoDescriptorSetTest, NoProtos) {
+  AssertFileDescriptorSetContains(
+      "bazel/tests/no_protos.pb", {});
+}
+
+TEST(ProtoDescriptorSetTest, WellKnownProtos) {
+  AssertFileDescriptorSetContains(
+      "bazel/tests/well_known_protos.pb",
+      {
+          "google/protobuf/any.proto",
+          "google/protobuf/api.proto",
+          "google/protobuf/descriptor.proto",
+          "google/protobuf/duration.proto",
+          "google/protobuf/empty.proto",
+          "google/protobuf/field_mask.proto",
+          "google/protobuf/source_context.proto",
+          "google/protobuf/struct.proto",
+          "google/protobuf/timestamp.proto",
+          "google/protobuf/type.proto",
+          "google/protobuf/wrappers.proto",
+      });
+}
+
+}  // namespace rulesproto