fix(cxx_common): change selector serialization to work better with pr…

…otos (#4785)

* fix(cxx_common): change selector serialization to work better with protos

* chore: use common implementation for pointers and values

* chore: fix typo

kythe/cxx/extractor/bazel_artifact_selector.cc

@@ -61,10 +61,53 @@ struct FromRange {
 };
 
 template <typename T>
-FromRange(const T&)->FromRange<T>;
+FromRange(const T&) -> FromRange<T>;
 
+template <typename T>
+const T& AsConstRef(const T& value) {
+  return value;
+}
+
+template <typename T>
+const T& AsConstRef(const T* value) {
+  return *value;
+}
+
+template <typename T, typename U>
+absl::Status DeserializeInternal(T& selector, const U& container) {
+  absl::Status error;
+  for (const auto& any : container) {
+    switch (auto status = selector.DeserializeFrom(AsConstRef(any));
+            status.code()) {
+      case absl::StatusCode::kOk:
+      case absl::StatusCode::kUnimplemented:
+        return absl::OkStatus();
+      case absl::StatusCode::kInvalidArgument:
+        return status;
+      case absl::StatusCode::kFailedPrecondition:
+        error = status;
+        continue;
+      default:
+        error = status;
+        LOG(WARNING) << "Unrecognized status code: " << status;
+    }
+  }
+  return error.ok() ? absl::NotFoundError("No state found")
+                    : absl::NotFoundError(
+                          absl::StrCat("No state found: ", error.ToString()));
+}
 }  // namespace
 
+absl::Status BazelArtifactSelector::Deserialize(
+    absl::Span<const google::protobuf::Any> state) {
+  return DeserializeInternal(*this, state);
+}
+
+absl::Status BazelArtifactSelector::Deserialize(
+    absl::Span<const google::protobuf::Any* const> state) {
+  return DeserializeInternal(*this, state);
+}
+
 absl::optional<BazelArtifact> AspectArtifactSelector::Select(
     const build_event_stream::BuildEvent& event) {
   absl::optional<BazelArtifact> result = absl::nullopt;
@@ -81,34 +124,33 @@ absl::optional<BazelArtifact> AspectArtifactSelector::Select(
   return result;
 }
 
-absl::optional<google::protobuf::Any> AspectArtifactSelector::Serialize()
-    const {
-  kythe::proto::BazelAspectArtifactSelectorState state;
-  *state.mutable_disposed() = FromRange{state_.disposed};
-  *state.mutable_filesets() = FromRange{state_.filesets};
-  *state.mutable_pending() = FromRange{state_.pending};
+bool AspectArtifactSelector::SerializeInto(google::protobuf::Any& state) const {
+  kythe::proto::BazelAspectArtifactSelectorState raw;
+  *raw.mutable_disposed() = FromRange{state_.disposed};
+  *raw.mutable_filesets() = FromRange{state_.filesets};
+  *raw.mutable_pending() = FromRange{state_.pending};
 
-  google::protobuf::Any result;
-  result.PackFrom(state);
-  return result;
+  state.PackFrom(raw);
+  return true;
 }
 
-absl::Status AspectArtifactSelector::Deserialize(
-    absl::Span<const google::protobuf::Any> state) {
-  for (const auto& any : state) {
-    kythe::proto::BazelAspectArtifactSelectorState proto_state;
-    if (any.UnpackTo(&proto_state)) {
-      state_ = {
-          .disposed = FromRange{proto_state.disposed()},
-          .filesets = FromRange{proto_state.filesets()},
-          .pending = FromRange{proto_state.pending()},
-      };
-
-      return absl::OkStatus();
-    }
+absl::Status AspectArtifactSelector::DeserializeFrom(
+    const google::protobuf::Any& state) {
+  kythe::proto::BazelAspectArtifactSelectorState raw;
+  if (state.UnpackTo(&raw)) {
+    state_ = {
+        .disposed = FromRange{raw.disposed()},
+        .filesets = FromRange{raw.filesets()},
+        .pending = FromRange{raw.pending()},
+    };
+    return absl::OkStatus();
+  }
+  if (state.Is<kythe::proto::BazelAspectArtifactSelectorState>()) {
+    return absl::InvalidArgumentError(
+        "Malformed kythe.proto.BazelAspectArtifactSelectorState");
   }
-  return absl::NotFoundError(
-      "No entry found for kythe.proto.BazelAspectArtifactSelectorState");
+  return absl::FailedPreconditionError(
+      "State not of type kythe.proto.BazelAspectArtifactSelectorState");
 }
 
 absl::optional<BazelArtifact> AspectArtifactSelector::SelectFileSet(

kythe/cxx/extractor/bazel_artifact_selector.h

@@ -47,25 +47,34 @@ class BazelArtifactSelector {
   virtual absl::optional<BazelArtifact> Select(
       const build_event_stream::BuildEvent& event) = 0;
 
-  /// \brief Returns an Any-encoded protobuf with per-stream selector state.
+  /// \brief Encodes per-stream selector state into the Any protobuf.
   /// Stateful selectors should serialize any per-stream state into a
   /// suitable protocol buffer, encoded as an Any. If no state has been
   /// accumulated, they should return an empty protocol buffer of the
-  /// appropriate type.
-  /// Stateless selectors should return nullopt.
-  virtual absl::optional<google::protobuf::Any> Serialize() const {
-    return absl::nullopt;
+  /// appropriate type and return true.
+  /// Stateless selectors should return false.
+  virtual bool SerializeInto(google::protobuf::Any& state) const {
+    return false;
   }
 
-  /// \brief Finds and updates any per-stream state from the provided list.
-  /// Stateless selectors should unconditionally return OK.
-  /// Stateful selectors should return NotFound if the custom protobuf is not
-  /// found in the provided list.
-  virtual absl::Status Deserialize(
-      absl::Span<const google::protobuf::Any> state) {
-    return absl::OkStatus();
+  /// \brief Updates any per-stream state from the provided proto.
+  /// Stateless selectors should unconditionally return a kUnimplemented status.
+  /// Stateful selectors should return OK if the provided state contains a
+  /// suitable proto, InvalidArgument if the proto is of the right type but
+  /// cannot be decoded or FailedPrecondition if the proto is of the wrong type.
+  virtual absl::Status DeserializeFrom(const google::protobuf::Any& state) {
+    return absl::UnimplementedError("stateless selector");
   }
 
+  /// \brief Finds and updates any per-stream state from the provided list.
+  /// Returns OK if the selector is stateless or if the requisite state was
+  /// found in the list.
+  /// Returns NotFound for a stateful selector whose state was not present
+  /// or InvalidArgument if the state was present but couldn't be decoded.
+  absl::Status Deserialize(absl::Span<const google::protobuf::Any> state);
+  absl::Status Deserialize(
+      absl::Span<const google::protobuf::Any* const> state);
+
  protected:
   // Not publicly copyable or movable to avoid slicing, but subclasses may be.
   BazelArtifactSelector() = default;
@@ -74,8 +83,8 @@ class BazelArtifactSelector {
 };
 
 /// \brief A type-erased value-type implementation of the BazelArtifactSelector
-/// interface. Does not derive from BazelArtifactSelector.
-class AnyArtifactSelector {
+/// interface.
+class AnyArtifactSelector final : public BazelArtifactSelector {
  public:
   /// \brief Constructs an AnyArtifactSelector which delegates to the provided
   /// argument, which must derive from BazelArtifactSelector.
@@ -104,13 +113,13 @@ class AnyArtifactSelector {
   }
 
   /// \brief Forwards serialization to the contained BazelArtifactSelector.
-  absl::optional<google::protobuf::Any> Serialize() const {
-    return get_().Serialize();
+  bool SerializeInto(google::protobuf::Any& state) const final {
+    return get_().SerializeInto(state);
   }
 
   /// \brief Forwards deserialization to the contained BazelArtifactSelector.
-  absl::Status Deserialize(absl::Span<const google::protobuf::Any> state) {
-    return get_().Deserialize(state);
+  absl::Status DeserializeFrom(const google::protobuf::Any& state) final {
+    return get_().DeserializeFrom(state);
   }
 
  private:
@@ -152,11 +161,11 @@ class AspectArtifactSelector final : public BazelArtifactSelector {
   /// \brief Serializes the accumulated state into the return value, which will
   /// always be non-empty and of type
   /// `kythe.proto.BazelAspectArtifactSelectorState`.
-  absl::optional<google::protobuf::Any> Serialize() const final;
+  bool SerializeInto(google::protobuf::Any& state) const final;
 
   /// \brief Deserializes accumulated stream state from an Any of type
-  /// `kythe.proto.BazelAspectArtifactSelectorState` or returns NotFound.
-  absl::Status Deserialize(absl::Span<const google::protobuf::Any> state) final;
+  /// `kythe.proto.BazelAspectArtifactSelectorState`.
+  absl::Status DeserializeFrom(const google::protobuf::Any& state) final;
 
  private:
   struct State {