TMaybe Monadic operations

Add monadic operations to TMaybe.
Added methods: AndThen, OrElse, Transform.

2 files changed, 229 insertions, 0 deletions

diff --git a/util/generic/maybe.h b/util/generic/maybe.h
index 8498afda34..b10c07e5fb 100644
--- a/util/generic/maybe.h
+++ b/util/generic/maybe.h
@@ -152,6 +152,29 @@ private:
                                       (!std::is_scalar<T>::value || !std::is_same<UDec, T>::value);
     };
 
+    template <typename>
+    struct TIsMaybe {
+        static constexpr bool value = false;
+    };
+
+    template <typename U, typename P>
+    struct TIsMaybe<TMaybe<U, P>> {
+        static constexpr bool value = true;
+    };
+
+    template <typename F, typename... Args>
+    static auto Call(F&& f, Args&&... args) -> decltype(std::forward<F>(f)(std::forward<Args>(args)...));
+
+    template <typename F, typename... Args>
+    using TCallResult = std::remove_reference_t<std::remove_cv_t<decltype(Call(std::declval<F>(), std::declval<Args>()...))>>;
+
+    template <typename U, typename F>
+    static constexpr F&& CheckReturnsMaybe(F&& f) {
+        using ReturnType = TCallResult<F, U>;
+        static_assert(TIsMaybe<ReturnType>::value, "Function must return TMaybe");
+        return f;
+    }
+
     using TBase = TMaybeBase<T>;
 
 public:
@@ -378,6 +401,118 @@ public:
         return Defined() ? *this : elseValue;
     }
 
+    template <typename F>
+    constexpr auto AndThen(F&& func) & {
+        using ReturnType = TCallResult<F, T&>;
+
+        if (Defined()) {
+            return std::forward<F>(CheckReturnsMaybe<T&>(func))(*Data());
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr auto AndThen(F&& func) const& {
+        using ReturnType = TCallResult<F, const T&>;
+
+        if (Defined()) {
+            return std::forward<F>(CheckReturnsMaybe<const T&>(func))(*Data());
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr auto AndThen(F&& func) && {
+        using ReturnType = TCallResult<F, T&&>;
+
+        if (Defined()) {
+            return std::forward<F>(CheckReturnsMaybe<T&&>(func))(std::move(*Data()));
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr auto AndThen(F&& func) const&& {
+        using ReturnType = TCallResult<F, const T&&>;
+
+        if (Defined()) {
+            return std::forward<F>(CheckReturnsMaybe<const T&&>(func))(std::move(*Data()));
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr auto Transform(F&& func) & {
+        using ReturnType = TMaybe<TCallResult<F, T&>>;
+
+        if (Defined()) {
+            return ReturnType(std::forward<F>(func)(*Data()));
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr auto Transform(F&& func) const& {
+        using ReturnType = TMaybe<TCallResult<F, const T&>>;
+
+        if (Defined()) {
+            return ReturnType(std::forward<F>(func)(*Data()));
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr auto Transform(F&& func) && {
+        using ReturnType = TMaybe<TCallResult<F, T&&>>;
+
+        if (Defined()) {
+            return ReturnType(std::forward<F>(func)(std::move(*Data())));
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr auto Transform(F&& func) const&& {
+        using ReturnType = TMaybe<TCallResult<F, const T&&>>;
+
+        if (Defined()) {
+            return ReturnType(std::forward<F>(func)(std::move(*Data())));
+        }
+
+        return ReturnType{};
+    }
+
+    template <typename F>
+    constexpr TMaybe Or(F&& func) const& {
+        using ResultType = TCallResult<F>;
+        static_assert(std::is_same<ResultType, TMaybe>::value, "Function must return TMaybe with the same type");
+
+        if (Defined()) {
+            return *this;
+        }
+
+        return std::forward<F>(func)();
+    }
+
+    template <typename F>
+    constexpr TMaybe Or(F&& func) && {
+        using ResultType = TCallResult<F>;
+        static_assert(std::is_same<ResultType, TMaybe>::value, "Function must return TMaybe with the same type");
+
+        if (Defined()) {
+            return std::move(*this);
+        }
+
+        return std::forward<F>(func)();
+    }
+
     template <typename U>
     TMaybe<U, Policy> Cast() const {
         return Defined() ? TMaybe<U, Policy>(*Data()) : TMaybe<U, Policy>();
diff --git a/util/generic/maybe_ut.cpp b/util/generic/maybe_ut.cpp
index 2c1a425c5e..a7c5407e81 100644
--- a/util/generic/maybe_ut.cpp
+++ b/util/generic/maybe_ut.cpp
@@ -239,6 +239,100 @@ Y_UNIT_TEST_SUITE(TMaybeTest) {
         }
     }
 
+    Y_UNIT_TEST(TestAndThen) {
+        {
+            auto f = [](int i) -> TMaybe<int> {
+                if (i % 2 == 0) {
+                    return i / 2;
+                }
+                return {};
+            };
+
+            UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{4}.AndThen(f), TMaybe<int>{2});
+            UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{5}.AndThen(f), TMaybe<int>{});
+            UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{}.AndThen(f), TMaybe<int>{});
+        }
+
+        {
+            // Move semantics: value stealing
+            TMaybe<TString> x{"Hello, "};
+            UNIT_ASSERT_VALUES_EQUAL(std::move(x).AndThen([](TString s) { return MakeMaybe(s + "world!"); }), TMaybe<TString>{"Hello, world!"});
+            UNIT_ASSERT_VALUES_EQUAL(*x, "");
+        }
+    }
+
+    Y_UNIT_TEST(TestTransform) {
+        auto f = [](int i) {
+            return i + 1;
+        };
+
+        UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{1}.Transform(f), TMaybe<int>{2});
+        UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{}.Transform(f), TMaybe<int>{});
+
+        auto f2 = [](int i) {
+            return ToString(i);
+        };
+
+        UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{1}.Transform(f2), TMaybe<TString>{"1"});
+        UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{}.Transform(f2), TMaybe<TString>{});
+
+        {
+            // Move semantics: value stealing
+            TMaybe<TString> x{"Hello, "};
+            UNIT_ASSERT_VALUES_EQUAL(std::move(x).Transform([](TString s) { return s + "world!"; }), TMaybe<TString>("Hello, world!"));
+            UNIT_ASSERT_VALUES_EQUAL(*x, "");
+        }
+
+        {
+            struct Data {
+                Data(int x)
+                    : value(x)
+                {
+                }
+
+                int value;
+                int GetValue() {
+                    return value;
+                }
+            };
+
+            UNIT_ASSERT_VALUES_EQUAL(MakeMaybe<Data>(5).Transform(std::mem_fn(&Data::value)), TMaybe<int>{5});
+            UNIT_ASSERT_VALUES_EQUAL(MakeMaybe<Data>(5).Transform(std::mem_fn(&Data::GetValue)), TMaybe<int>{5});
+        }
+    }
+
+    Y_UNIT_TEST(TestOr) {
+        {
+            auto f = []() {
+                return TMaybe<int>{5};
+            };
+
+            UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{1}.Or(f), TMaybe<int>{1});
+            UNIT_ASSERT_VALUES_EQUAL(TMaybe<int>{}.Or(f), TMaybe<int>{5});
+        }
+
+        {
+            bool flag = false;
+
+            auto f = [&flag]() {
+                flag = true;
+                return TMaybe<int>{5};
+            };
+
+            TMaybe<int>{5}.Or(f);
+            UNIT_ASSERT(!flag);
+
+            TMaybe<int>{}.Or(f);
+            UNIT_ASSERT(flag);
+        }
+
+        {
+            TMaybe<TString> x{"abacaba"};
+            std::move(x).Or([]() { return TMaybe<TString>{}; });
+            UNIT_ASSERT_VALUES_EQUAL(*x, "");
+        }
+    }
+
     /*
   ==
   !=