From ec9fc0c64acbd214af4f170061e3e9ff3ff6274b Mon Sep 17 00:00:00 2001
From: Tom Jackson <tjackson@fb.com>
Date: Thu, 29 Jan 2015 10:56:30 -0800
Subject: [PATCH] Making 'just()' reference arguments like 'from()'

Summary: That is, for all inputs except r-values.

Test Plan:
Run tests
@override-test-failures

Reviewed By: ajaymenon@fb.com

Subscribers: trunkagent, folly-diffs@

FB internal diff: D1807789

Signature: t1:1807789:1422410689:442d50ab17f8e3b08e34f33318a4dc7f23b4c6cc
---
 folly/gen/Base-inl.h        | 29 ++++++++++++++++++++++++-----
 folly/gen/Base.h            | 24 +++++++++++++++++++-----
 folly/gen/Parallel-inl.h    |  2 +-
 folly/gen/test/BaseTest.cpp | 25 +++++++++++++++++++++++++
 4 files changed, 69 insertions(+), 11 deletions(-)

diff --git a/folly/gen/Base-inl.h b/folly/gen/Base-inl.h
index 083cf979..e04c246c 100644
--- a/folly/gen/Base-inl.h
+++ b/folly/gen/Base-inl.h
@@ -328,13 +328,32 @@ class Empty : public GenImpl<Value, Empty<Value>> {
   void foreach(Body&&) const {}
 };
 
-template<class Value>
-class Just : public GenImpl<const Value&, Just<Value>> {
+template <class Value>
+class SingleReference : public GenImpl<Value&, SingleReference<Value>> {
+  static_assert(!std::is_reference<Value>::value,
+                "SingleReference requires non-ref types");
+  Value* ptr_;
+ public:
+  explicit SingleReference(Value* ptr) : ptr_(ptr){}
+
+  template <class Handler>
+  bool apply(Handler&& handler) const {
+    return handler(*ptr_);
+  }
+
+  template <class Body>
+  void foreach(Body&& body) const {
+    body(*ptr_);
+  }
+};
+
+template <class Value>
+class SingleCopy : public GenImpl<const Value&, SingleCopy<Value>> {
   static_assert(!std::is_reference<Value>::value,
-                "Just requires non-ref types");
-  const Value value_;
+                "SingleCopy requires non-ref types");
+  Value value_;
  public:
-  explicit Just(Value value) : value_(std::forward<Value>(value)) {}
+  explicit SingleCopy(Value value) : value_(std::forward<Value>(value)) {}
 
   template <class Handler>
   bool apply(Handler&& handler) const {
diff --git a/folly/gen/Base.h b/folly/gen/Base.h
index 12d917ed..bf9b4b53 100644
--- a/folly/gen/Base.h
+++ b/folly/gen/Base.h
@@ -297,7 +297,10 @@ template<class Value>
 class Empty;
 
 template<class Value>
-class Just;
+class SingleReference;
+
+template<class Value>
+class SingleCopy;
 
 /*
  * Operators
@@ -482,14 +485,25 @@ Yield generator(Source&& source) {
 /*
  * empty() - for producing empty sequences.
  */
-template<class Value>
+template <class Value>
 detail::Empty<Value> empty() {
   return {};
 }
 
-template<class Value>
-detail::Just<Value> just(Value value) {
-  return detail::Just<Value>(std::move(value));
+template <class Value, class Just = detail::SingleReference<Value>>
+Just just(Value& value) {
+  return Just(&value);
+}
+
+template <class Value, class Just = detail::SingleReference<const Value>>
+Just just(const Value& value) {
+  return Just(&value);
+}
+
+template <class Value,
+          class Just = detail::SingleCopy<typename std::decay<Value>::type>>
+Just just(Value&& value) {
+  return Just(std::move(value));
 }
 
 /*
diff --git a/folly/gen/Parallel-inl.h b/folly/gen/Parallel-inl.h
index 666acb16..06f1af01 100644
--- a/folly/gen/Parallel-inl.h
+++ b/folly/gen/Parallel-inl.h
@@ -136,7 +136,7 @@ class Sub : public Operator<Sub<Sink>> {
             class Source,
             class Result =
                 decltype(std::declval<Sink>().compose(std::declval<Source>())),
-            class Just = Just<typename std::decay<Result>::type>>
+            class Just = SingleCopy<typename std::decay<Result>::type>>
   Just compose(const GenImpl<Value, Source>& source) const {
     return Just(source | sink_);
   }
diff --git a/folly/gen/test/BaseTest.cpp b/folly/gen/test/BaseTest.cpp
index 2b99f157..d403d99d 100644
--- a/folly/gen/test/BaseTest.cpp
+++ b/folly/gen/test/BaseTest.cpp
@@ -1073,6 +1073,31 @@ TEST(Gen, BatchMove) {
   EXPECT_EQ(expected, actual);
 }
 
+TEST(Gen, Just) {
+  {
+    int x = 3;
+    auto j = just(x);
+    EXPECT_EQ(&x, j | mapped([](const int& v) { return &v; }) | first);
+    x = 4;
+    EXPECT_EQ(4, j | first);
+  }
+  {
+    int x = 3;
+    const int& cx = x;
+    auto j = just(cx);
+    EXPECT_EQ(&x, j | mapped([](const int& v) { return &v; }) | first);
+    x = 5;
+    EXPECT_EQ(5, j | first);
+  }
+  {
+    int x = 3;
+    auto j = just(std::move(x));
+    EXPECT_NE(&x, j | mapped([](const int& v) { return &v; }) | first);
+    x = 5;
+    EXPECT_EQ(3, j | first);
+  }
+}
+
 int main(int argc, char *argv[]) {
   testing::InitGoogleTest(&argc, argv);
   gflags::ParseCommandLineFlags(&argc, &argv, true);
-- 
2.34.1