FTL: Add Optional<T>::and_then

Bug: 185536303
Test: ftl_test
Change-Id: Ic907a662a2090ad1363bd329e7d758b2acef55ad

4 files changed, 131 insertions, 19 deletions

diff --git a/include/ftl/details/optional.h b/include/ftl/details/optional.h
new file mode 100644
index 0000000000..bff7c1e000
--- /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 0000000000..7092ec50fb
--- /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 daf4502bd2..a0a95c4b9a 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 @@ struct Optional final : std::optional<T> {
   // 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 339726e4ea..11294c3ac8 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 @@ class SmallVector final : details::ArrayTraits<T>, details::ArrayComparators<Sma
   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 @@ class SmallVector final : details::ArrayTraits<T>, details::ArrayComparators<Sma
 
   // 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)...) {}