Project import generated by Copybara.

GitOrigin-RevId: 73d686c40057684f8bfaca285368bf1813f9fc26

mediapipe/framework/api2/BUILD

@@ -134,6 +134,7 @@ cc_test(
     deps = [
         ":packet",
         "//mediapipe/framework/port:gtest_main",
+        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
     ],
 )

mediapipe/framework/api2/builder.h

@@ -313,8 +313,8 @@ template <class Calc>
 class Node : public NodeBase {
  public:
   Node() : NodeBase(Calc::kCalculatorName) {}
-  // Overrides the built-in calculator type std::string with the provided
-  // argument. Can be used to create nodes from pure interfaces.
+  // Overrides the built-in calculator type string with the provided argument.
+  // Can be used to create nodes from pure interfaces.
   // TODO: only use this for pure interfaces
   Node(const std::string& type_override) : NodeBase(type_override) {}
 
@@ -377,6 +377,29 @@ class PacketGenerator {
     return *options_.MutableExtension(T::ext);
   }
 
+  template <typename B, typename T, bool kIsOptional, bool kIsMultiple>
+  auto operator[](const PortCommon<B, T, kIsOptional, kIsMultiple>& port) {
+    using PayloadT =
+        typename PortCommon<B, T, kIsOptional, kIsMultiple>::PayloadT;
+    if constexpr (std::is_same_v<B, SideOutputBase>) {
+      auto* base = &out_sides_[port.Tag()];
+      if constexpr (kIsMultiple) {
+        return MultiSideSource<PayloadT>(base);
+      } else {
+        return SideSource<PayloadT>(base);
+      }
+    } else if constexpr (std::is_same_v<B, SideInputBase>) {
+      auto* base = &in_sides_[port.Tag()];
+      if constexpr (kIsMultiple) {
+        return MultiSideDestination<PayloadT>(base);
+      } else {
+        return SideDestination<PayloadT>(base);
+      }
+    } else {
+      static_assert(dependent_false<B>::value, "Type not supported.");
+    }
+  }
+
  private:
   std::string type_;
   TagIndexMap<DestinationBase> in_sides_;
@@ -402,7 +425,7 @@ class Graph {
   }
 
   // Creates a node of a specific type. Should be used for pure interfaces,
-  // which do not have a built-in type std::string.
+  // which do not have a built-in type string.
   template <class Calc>
   Node<Calc>& AddNode(const std::string& type) {
     auto node = std::make_unique<Node<Calc>>(type);

mediapipe/framework/api2/const_str.h

@@ -6,8 +6,8 @@
 namespace mediapipe {
 namespace api2 {
 
-// This class stores a constant std::string that can be inspected at compile
-// time in constexpr code.
+// This class stores a constant string that can be inspected at compile time
+// in constexpr code.
 class const_str {
  public:
   constexpr const_str(std::size_t size, const char* data)

mediapipe/framework/api2/packet.h

@@ -215,6 +215,7 @@ class Packet : public Packet<internal::Generic> {
     return typed_payload->data();
   }
   const T& operator*() const { return Get(); }
+  const T* operator->() const { return &Get(); }
 
   template <typename U>
   T GetOr(U&& v) const {

mediapipe/framework/api2/packet_test.cc

@@ -1,5 +1,6 @@
 #include "mediapipe/framework/api2/packet.h"
 
+#include "absl/memory/memory.h"
 #include "absl/strings/str_cat.h"
 #include "mediapipe/framework/port/gmock.h"
 #include "mediapipe/framework/port/gtest.h"
@@ -18,6 +19,17 @@ class LiveCheck {
   bool& alive_;
 };
 
+class Base {
+ public:
+  virtual ~Base() = default;
+  virtual absl::string_view name() const { return "Base"; }
+};
+
+class Derived : public Base {
+ public:
+  absl::string_view name() const override { return "Derived"; }
+};
+
 TEST(PacketTest, PacketBaseDefault) {
   PacketBase p;
   EXPECT_TRUE(p.IsEmpty());
@@ -242,6 +254,16 @@ TEST(PacketTest, OneOfConsume) {
   EXPECT_TRUE(p.IsEmpty());
 }
 
+TEST(PacketTest, Polymorphism) {
+  Packet<Base> base = PacketAdopting<Base>(absl::make_unique<Derived>());
+  EXPECT_EQ(base->name(), "Derived");
+  // Since packet contents are implicitly immutable, if you need mutability the
+  // current recommendation is still to wrap the contents in a unique_ptr.
+  Packet<std::unique_ptr<Base>> mutable_base =
+      MakePacket<std::unique_ptr<Base>>(absl::make_unique<Derived>());
+  EXPECT_EQ((**mutable_base).name(), "Derived");
+}
+
 }  // namespace
 }  // namespace api2
 }  // namespace mediapipe

mediapipe/framework/api2/port.h

@@ -172,9 +172,14 @@ inline void SetType<NoneType>(CalculatorContract* cc, PacketType& pt) {
   pt.SetNone();
 }
 
+template <typename... T>
+inline void SetTypeOneOf(OneOf<T...>, CalculatorContract* cc, PacketType& pt) {
+  pt.SetOneOf<T...>();
+}
+
 template <typename T, typename std::enable_if<IsOneOf<T>{}, int>::type = 0>
 inline void SetType(CalculatorContract* cc, PacketType& pt) {
-  pt.SetAny();
+  SetTypeOneOf(T{}, cc, pt);
 }
 
 template <typename ValueT>
@@ -294,14 +299,26 @@ struct SideBase<InputBase> {
   using type = SideInputBase;
 };
 
+// TODO: maybe return a PacketBase instead of a Packet<internal::Generic>?
+template <typename T, class = void>
+struct ActualPayloadType {
+  using type = T;
+};
+
+template <typename T>
+struct ActualPayloadType<
+    T, std::enable_if_t<std::is_base_of<DynamicType, T>{}, void>> {
+  using type = internal::Generic;
+};
+
 }  // namespace internal
 
-// TODO: maybe return a PacketBase instead of a Packet<internal::Generic>?
-template <typename T, typename std::enable_if<
-                          !std::is_base_of<DynamicType, T>{}, int>::type = 0>
-auto ActualValueT(T) -> T;
+// Maps special port value types, such as AnyType, to internal::Generic.
+template <typename T>
+using ActualPayloadT = typename internal::ActualPayloadType<T>::type;
 
-auto ActualValueT(DynamicType) -> internal::Generic;
+static_assert(std::is_same_v<ActualPayloadT<int>, int>, "");
+static_assert(std::is_same_v<ActualPayloadT<AnyType>, internal::Generic>, "");
 
 template <typename Base, typename ValueT, bool IsOptional = false,
           bool IsMultiple = false>
@@ -325,7 +342,7 @@ class PortCommon : public Base {
   explicit constexpr PortCommon(const char (&tag)[N])
       : Base(N, tag, &get_type_hash<ValueT>, IsOptionalV, IsMultipleV) {}
 
-  using PayloadT = decltype(ActualValueT(std::declval<ValueT>()));
+  using PayloadT = ActualPayloadT<ValueT>;
 
   auto operator()(CalculatorContext* cc) const {
     return internal::AccessPort<PayloadT>(
@@ -385,7 +402,7 @@ class SideFallbackT : public Base {
   static constexpr bool kOptional = IsOptionalV;
   static constexpr bool kMultiple = IsMultipleV;
   using Optional = SideFallbackT<Base, ValueT, true, IsMultipleV>;
-  using PayloadT = decltype(ActualValueT(std::declval<ValueT>()));
+  using PayloadT = ActualPayloadT<ValueT>;
 
   const char* Tag() const { return stream_port.Tag(); }
 
@@ -499,6 +516,10 @@ class OutputShardAccess : public OutputShardAccessBase {
     Send(std::move(payload), context_.InputTimestamp());
   }
 
+  void SetHeader(const PacketBase& header) {
+    if (output_) output_->SetHeader(ToOldPacket(header));
+  }
+
  private:
   OutputShardAccess(const CalculatorContext& cc, OutputStreamShard* output)
       : OutputShardAccessBase(cc, output) {}

mediapipe/framework/api2/port_test.cc

@@ -21,6 +21,25 @@ TEST(PortTest, Tag) {
   EXPECT_EQ(std::string(port.Tag()), "FOO");
 }
 
+struct DeletedCopyType {
+  DeletedCopyType(const DeletedCopyType&) = delete;
+  DeletedCopyType& operator=(const DeletedCopyType&) = delete;
+};
+
+TEST(PortTest, DeletedCopyConstructorInput) {
+  static constexpr Input<DeletedCopyType> kInputPort{"INPUT"};
+  EXPECT_EQ(std::string(kInputPort.Tag()), "INPUT");
+
+  static constexpr Output<DeletedCopyType> kOutputPort{"OUTPUT"};
+  EXPECT_EQ(std::string(kOutputPort.Tag()), "OUTPUT");
+
+  static constexpr SideInput<DeletedCopyType> kSideInputPort{"SIDE_INPUT"};
+  EXPECT_EQ(std::string(kSideInputPort.Tag()), "SIDE_INPUT");
+
+  static constexpr SideOutput<DeletedCopyType> kSideOutputPort{"SIDE_OUTPUT"};
+  EXPECT_EQ(std::string(kSideOutputPort.Tag()), "SIDE_OUTPUT");
+}
+
 }  // namespace
 }  // namespace api2
 }  // namespace mediapipe

mediapipe/framework/api2/tag_test.cc

@@ -26,8 +26,8 @@ TEST(TagTest, String) {
   EXPECT_EQ(kBAR.str(), "BAR");
 }
 
-// Separate invocations of MPP_TAG with the same std::string produce objects of
-// the same type.
+// Separate invocations of MPP_TAG with the same string produce objects of the
+// same type.
 TEST(TagTest, SameType) { EXPECT_TRUE(same_type(kFOO, kFOO2)); }
 
 // Different tags have different types.