Project import generated by Copybara.

GitOrigin-RevId: 33adfdf31f3a5cbf9edc07ee1ea583e95080bdc5

mediapipe/framework/api2/builder.h

@@ -17,6 +17,14 @@ namespace mediapipe {
 namespace api2 {
 namespace builder {
 
+// Workaround for static_assert(false). Example:
+//   dependent_false<T>::value returns false.
+// For more information, see:
+// https://en.cppreference.com/w/cpp/language/if#Constexpr_If
+// TODO: migrate to a common utility when available.
+template <class T>
+struct dependent_false : std::false_type {};
+
 template <typename T>
 T& GetWithAutoGrow(std::vector<std::unique_ptr<T>>* vecp, int index) {
   auto& vec = *vecp;
@@ -209,6 +217,21 @@ class NodeBase {
     return SideDestination<true>(&in_sides_[tag]);
   }
 
+  template <typename B, typename T, bool kIsOptional, bool kIsMultiple>
+  auto operator[](const PortCommon<B, T, kIsOptional, kIsMultiple>& port) {
+    if constexpr (std::is_same_v<B, OutputBase>) {
+      return Source<kIsMultiple, T>(&out_streams_[port.Tag()]);
+    } else if constexpr (std::is_same_v<B, InputBase>) {
+      return Destination<kIsMultiple, T>(&in_streams_[port.Tag()]);
+    } else if constexpr (std::is_same_v<B, SideOutputBase>) {
+      return SideSource<kIsMultiple, T>(&out_sides_[port.Tag()]);
+    } else if constexpr (std::is_same_v<B, SideInputBase>) {
+      return SideDestination<kIsMultiple, T>(&in_sides_[port.Tag()]);
+    } else {
+      static_assert(dependent_false<B>::value, "Type not supported.");
+    }
+  }
+
   // Convenience methods for accessing purely index-based ports.
   Source<false> Out(int index) { return Out("")[index]; }
 
@@ -429,6 +452,24 @@ class Graph {
     return Dst(&graph_boundary_.in_sides_[graph_output.Tag()]);
   }
 
+  template <typename B, typename T, bool kIsOptional, bool kIsMultiple>
+  auto operator[](const PortCommon<B, T, kIsOptional, kIsMultiple>& port) {
+    if constexpr (std::is_same_v<B, OutputBase>) {
+      return Destination<kIsMultiple, T>(
+          &graph_boundary_.in_streams_[port.Tag()]);
+    } else if constexpr (std::is_same_v<B, InputBase>) {
+      return Source<kIsMultiple, T>(&graph_boundary_.out_streams_[port.Tag()]);
+    } else if constexpr (std::is_same_v<B, SideOutputBase>) {
+      return SideDestination<kIsMultiple, T>(
+          &graph_boundary_.in_sides_[port.Tag()]);
+    } else if constexpr (std::is_same_v<B, SideInputBase>) {
+      return SideSource<kIsMultiple, T>(
+          &graph_boundary_.out_sides_[port.Tag()]);
+    } else {
+      static_assert(dependent_false<B>::value, "Type not supported.");
+    }
+  }
+
   // Returns the graph config. This can be used to instantiate and run the
   // graph.
   CalculatorGraphConfig GetConfig() {

mediapipe/framework/api2/builder_test.cc

@@ -138,6 +138,35 @@ TEST(BuilderTest, TypedMultiple) {
   EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
 }
 
+TEST(BuilderTest, TypedByPorts) {
+  builder::Graph graph;
+  auto& foo = graph.AddNode<Foo>();
+  auto& adder = graph.AddNode<FloatAdder>();
+
+  graph[FooBar1::kIn].SetName("base") >> foo[Foo::kBase];
+  foo[Foo::kOut] >> adder[FloatAdder::kIn][0];
+  foo[Foo::kOut] >> adder[FloatAdder::kIn][1];
+  adder[FloatAdder::kOut].SetName("out") >> graph[FooBar1::kOut];
+
+  CalculatorGraphConfig expected =
+      mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+        input_stream: "IN:base"
+        output_stream: "OUT:out"
+        node {
+          calculator: "Foo"
+          input_stream: "BASE:base"
+          output_stream: "OUT:__stream_0"
+        }
+        node {
+          calculator: "FloatAdder"
+          input_stream: "IN:0:__stream_0"
+          input_stream: "IN:1:__stream_0"
+          output_stream: "OUT:out"
+        }
+      )pb");
+  EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
+}
+
 TEST(BuilderTest, PacketGenerator) {
   builder::Graph graph;
   auto& generator = graph.AddPacketGenerator("FloatGenerator");