FTL: Add Optional<T>::and_then

Bug: 185536303
Test: ftl_test
Change-Id: Ic907a662a2090ad1363bd329e7d758b2acef55ad
diff --git a/include/ftl/details/optional.h b/include/ftl/details/optional.h
new file mode 100644
index 0000000..bff7c1e
--- /dev/null
+++ b/include/ftl/details/optional.h
@@ -0,0 +1,58 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <functional>
+#include <optional>
+
+#include <ftl/details/type_traits.h>
+
+namespace android::ftl {
+
+template <typename>
+struct Optional;
+
+namespace details {
+
+template <typename>
+struct is_optional : std::false_type {};
+
+template <typename T>
+struct is_optional<std::optional<T>> : std::true_type {};
+
+template <typename T>
+struct is_optional<Optional<T>> : std::true_type {};
+
+template <typename F, typename T>
+struct transform_result {
+  using type = Optional<std::remove_cv_t<std::invoke_result_t<F, T>>>;
+};
+
+template <typename F, typename T>
+using transform_result_t = typename transform_result<F, T>::type;
+
+template <typename F, typename T>
+struct and_then_result {
+  using type = remove_cvref_t<std::invoke_result_t<F, T>>;
+  static_assert(is_optional<type>{}, "and_then function must return an optional");
+};
+
+template <typename F, typename T>
+using and_then_result_t = typename and_then_result<F, T>::type;
+
+}  // namespace details
+}  // namespace android::ftl
diff --git a/include/ftl/details/type_traits.h b/include/ftl/details/type_traits.h
new file mode 100644
index 0000000..7092ec5
--- /dev/null
+++ b/include/ftl/details/type_traits.h
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <type_traits>
+
+namespace android::ftl::details {
+
+// TODO: Replace with std::remove_cvref_t in C++20.
+template <typename U>
+using remove_cvref_t = std::remove_cv_t<std::remove_reference_t<U>>;
+
+}  // namespace android::ftl::details
diff --git a/include/ftl/optional.h b/include/ftl/optional.h
index daf4502..a0a95c4 100644
--- a/include/ftl/optional.h
+++ b/include/ftl/optional.h
@@ -18,9 +18,10 @@
 
 #include <functional>
 #include <optional>
-#include <type_traits>
 #include <utility>
 
+#include <ftl/details/optional.h>
+
 namespace android::ftl {
 
 // Superset of std::optional<T> with monadic operations, as proposed in https://wg21.link/P0798R8.
@@ -37,30 +38,59 @@
   // Returns Optional<U> where F is a function that maps T to U.
   template <typename F>
   constexpr auto transform(F&& f) const& {
-    using U = std::remove_cv_t<std::invoke_result_t<F, decltype(value())>>;
-    if (has_value()) return Optional<U>(std::invoke(std::forward<F>(f), value()));
-    return Optional<U>();
+    using R = details::transform_result_t<F, decltype(value())>;
+    if (has_value()) return R(std::invoke(std::forward<F>(f), value()));
+    return R();
   }
 
   template <typename F>
   constexpr auto transform(F&& f) & {
-    using U = std::remove_cv_t<std::invoke_result_t<F, decltype(value())>>;
-    if (has_value()) return Optional<U>(std::invoke(std::forward<F>(f), value()));
-    return Optional<U>();
+    using R = details::transform_result_t<F, decltype(value())>;
+    if (has_value()) return R(std::invoke(std::forward<F>(f), value()));
+    return R();
   }
 
   template <typename F>
   constexpr auto transform(F&& f) const&& {
-    using U = std::invoke_result_t<F, decltype(std::move(value()))>;
-    if (has_value()) return Optional<U>(std::invoke(std::forward<F>(f), std::move(value())));
-    return Optional<U>();
+    using R = details::transform_result_t<F, decltype(std::move(value()))>;
+    if (has_value()) return R(std::invoke(std::forward<F>(f), std::move(value())));
+    return R();
   }
 
   template <typename F>
   constexpr auto transform(F&& f) && {
-    using U = std::invoke_result_t<F, decltype(std::move(value()))>;
-    if (has_value()) return Optional<U>(std::invoke(std::forward<F>(f), std::move(value())));
-    return Optional<U>();
+    using R = details::transform_result_t<F, decltype(std::move(value()))>;
+    if (has_value()) return R(std::invoke(std::forward<F>(f), std::move(value())));
+    return R();
+  }
+
+  // Returns Optional<U> where F is a function that maps T to Optional<U>.
+  template <typename F>
+  constexpr auto and_then(F&& f) const& {
+    using R = details::and_then_result_t<F, decltype(value())>;
+    if (has_value()) return std::invoke(std::forward<F>(f), value());
+    return R();
+  }
+
+  template <typename F>
+  constexpr auto and_then(F&& f) & {
+    using R = details::and_then_result_t<F, decltype(value())>;
+    if (has_value()) return std::invoke(std::forward<F>(f), value());
+    return R();
+  }
+
+  template <typename F>
+  constexpr auto and_then(F&& f) const&& {
+    using R = details::and_then_result_t<F, decltype(std::move(value()))>;
+    if (has_value()) return std::invoke(std::forward<F>(f), std::move(value()));
+    return R();
+  }
+
+  template <typename F>
+  constexpr auto and_then(F&& f) && {
+    using R = details::and_then_result_t<F, decltype(std::move(value()))>;
+    if (has_value()) return std::invoke(std::forward<F>(f), std::move(value()));
+    return R();
   }
 };
 
diff --git a/include/ftl/small_vector.h b/include/ftl/small_vector.h
index 339726e..11294c3 100644
--- a/include/ftl/small_vector.h
+++ b/include/ftl/small_vector.h
@@ -21,11 +21,12 @@
 
 #include <algorithm>
 #include <iterator>
-#include <type_traits>
 #include <utility>
 #include <variant>
 #include <vector>
 
+#include <ftl/details/type_traits.h>
+
 namespace android::ftl {
 
 template <typename>
@@ -80,10 +81,6 @@
   using Static = StaticVector<T, N>;
   using Dynamic = SmallVector<T, 0>;
 
-  // TODO: Replace with std::remove_cvref_t in C++20.
-  template <typename U>
-  using remove_cvref_t = std::remove_cv_t<std::remove_reference_t<U>>;
-
  public:
   FTL_ARRAY_TRAIT(T, value_type);
   FTL_ARRAY_TRAIT(T, size_type);
@@ -104,7 +101,7 @@
 
   // Constructs at most N elements. See StaticVector for underlying constructors.
   template <typename Arg, typename... Args,
-            typename = std::enable_if_t<!is_small_vector<remove_cvref_t<Arg>>{}>>
+            typename = std::enable_if_t<!is_small_vector<details::remove_cvref_t<Arg>>{}>>
   SmallVector(Arg&& arg, Args&&... args)
       : vector_(std::in_place_type<Static>, std::forward<Arg>(arg), std::forward<Args>(args)...) {}
 
diff --git a/libs/ftl/optional_test.cpp b/libs/ftl/optional_test.cpp
index ede159a..ad8d3cf 100644
--- a/libs/ftl/optional_test.cpp
+++ b/libs/ftl/optional_test.cpp
@@ -20,6 +20,7 @@
 #include <ftl/unit.h>
 #include <gtest/gtest.h>
 
+#include <cstdlib>
 #include <functional>
 #include <numeric>
 #include <utility>
@@ -82,4 +83,62 @@
                      .transform([](const std::string& s) { return s.length(); }));
 }
 
+namespace {
+
+Optional<int> parse_int(const std::string& str) {
+  if (const int i = std::atoi(str.c_str())) return i;
+  return std::nullopt;
+}
+
+}  // namespace
+
+TEST(Optional, AndThen) {
+  // Empty.
+  EXPECT_EQ(std::nullopt, Optional<int>().and_then([](int) -> Optional<int> { return 0; }));
+  EXPECT_EQ(std::nullopt, Optional<int>().and_then([](int) { return Optional<int>(); }));
+
+  // By value.
+  EXPECT_EQ(0, Optional(0).and_then([](int x) { return Optional(x); }));
+  EXPECT_EQ(123, Optional("123").and_then(parse_int));
+  EXPECT_EQ(std::nullopt, Optional("abc").and_then(parse_int));
+
+  // By reference.
+  {
+    Optional opt = 'x';
+    EXPECT_EQ('z', opt.and_then([](char& c) {
+      c = 'y';
+      return Optional('z');
+    }));
+
+    EXPECT_EQ('y', opt);
+  }
+
+  // By rvalue reference.
+  {
+    std::string out;
+    EXPECT_EQ("xyz"s, Optional("abc"s).and_then([&out](std::string&& str) {
+      out = std::move(str);
+      return Optional("xyz"s);
+    }));
+
+    EXPECT_EQ(out, "abc"s);
+  }
+
+  // Chaining.
+  using StringVector = StaticVector<std::string, 3>;
+  EXPECT_EQ(14u, Optional(StaticVector{"-"s, "1"s})
+                     .and_then([](StringVector&& v) -> Optional<StringVector> {
+                       if (v.push_back("4"s)) return v;
+                       return {};
+                     })
+                     .and_then([](const StringVector& v) -> Optional<std::string> {
+                       if (v.full()) return std::accumulate(v.begin(), v.end(), std::string());
+                       return {};
+                     })
+                     .and_then(parse_int)
+                     .and_then([](int i) {
+                       return i > 0 ? std::nullopt : std::make_optional(static_cast<unsigned>(-i));
+                     }));
+}
+
 }  // namespace android::test