Project import generated by Copybara.

GitOrigin-RevId: 1610e588e497817fae2d9a458093ab6a370e2972

mediapipe/framework/api2/builder.h

@@ -148,18 +148,18 @@ class SourceImpl {
 
   explicit SourceImpl(std::vector<std::unique_ptr<Base>>* vec)
       : SourceImpl(&GetWithAutoGrow(vec, 0)) {}
-  explicit SourceImpl(SourceBase* base) : base_(*base) {}
+  explicit SourceImpl(SourceBase* base) : base_(base) {}
 
   template <typename U,
             typename std::enable_if<AllowConnection<U>{}, int>::type = 0>
   Src& AddTarget(const Dst<U>& dest) {
     CHECK(dest.base_.source == nullptr);
-    dest.base_.source = &base_;
-    base_.dests_.emplace_back(&dest.base_);
+    dest.base_.source = base_;
+    base_->dests_.emplace_back(&dest.base_);
     return *this;
   }
   Src& SetName(std::string name) {
-    base_.name_ = std::move(name);
+    base_->name_ = std::move(name);
     return *this;
   }
   template <typename U>
@@ -168,7 +168,8 @@ class SourceImpl {
   }
 
  private:
-  SourceBase& base_;
+  // Never null.
+  SourceBase* base_;
 };
 
 template <bool IsSide, typename T>

mediapipe/framework/api2/builder_test.cc

@@ -1,5 +1,7 @@
 #include "mediapipe/framework/api2/builder.h"
 
+#include <functional>
+
 #include "absl/strings/substitute.h"
 #include "mediapipe/framework/api2/node.h"
 #include "mediapipe/framework/api2/packet.h"
@@ -46,6 +48,88 @@ TEST(BuilderTest, BuildGraph) {
   EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
 }
 
+TEST(BuilderTest, CopyableSource) {
+  builder::Graph graph;
+  builder::Source<false, int> a = graph[Input<int>("A")];
+  a.SetName("a");
+  builder::Source<false, int> b = graph[Input<int>("B")];
+  b.SetName("b");
+  builder::SideSource<false, float> side_a = graph[SideInput<float>("SIDE_A")];
+  side_a.SetName("side_a");
+  builder::SideSource<false, float> side_b = graph[SideInput<float>("SIDE_B")];
+  side_b.SetName("side_b");
+  builder::Destination<false, int> out = graph[Output<int>("OUT")];
+  builder::SideDestination<false, float> side_out =
+      graph[SideOutput<float>("SIDE_OUT")];
+
+  builder::Source<false, int> input = a;
+  input = b;
+  builder::SideSource<false, float> side_input = side_b;
+  side_input = side_a;
+
+  input >> out;
+  side_input >> side_out;
+
+  CalculatorGraphConfig expected =
+      mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+        input_stream: "A:a"
+        input_stream: "B:b"
+        output_stream: "OUT:b"
+        input_side_packet: "SIDE_A:side_a"
+        input_side_packet: "SIDE_B:side_b"
+        output_side_packet: "SIDE_OUT:side_a"
+      )pb");
+  EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
+}
+
+TEST(BuilderTest, BuildGraphWithFunctions) {
+  builder::Graph graph;
+
+  builder::Source<false, int> base = graph[Input<int>("IN")];
+  base.SetName("base");
+  builder::SideSource<false, float> side = graph[SideInput<float>("SIDE")];
+  side.SetName("side");
+
+  auto foo_fn = [](builder::Source<false, int> base,
+                   builder::SideSource<false, float> side,
+                   builder::Graph& graph) {
+    auto& foo = graph.AddNode("Foo");
+    base >> foo[Input<int>("BASE")];
+    side >> foo[SideInput<float>("SIDE")];
+    return foo[Output<double>("OUT")];
+  };
+  builder::Source<false, double> foo_out = foo_fn(base, side, graph);
+
+  auto bar_fn = [](builder::Source<false, double> in, builder::Graph& graph) {
+    auto& bar = graph.AddNode("Bar");
+    in >> bar[Input<double>("IN")];
+    return bar[Output<double>("OUT")];
+  };
+  builder::Source<false, double> bar_out = bar_fn(foo_out, graph);
+  bar_out.SetName("out");
+
+  bar_out >> graph[Output<double>("OUT")];
+
+  CalculatorGraphConfig expected =
+      mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+        input_stream: "IN:base"
+        input_side_packet: "SIDE:side"
+        output_stream: "OUT:out"
+        node {
+          calculator: "Foo"
+          input_stream: "BASE:base"
+          input_side_packet: "SIDE:side"
+          output_stream: "OUT:__stream_0"
+        }
+        node {
+          calculator: "Bar"
+          input_stream: "IN:__stream_0"
+          output_stream: "OUT:out"
+        }
+      )pb");
+  EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
+}
+
 template <class FooT>
 void BuildGraphTypedTest() {
   builder::Graph graph;