Memoization.hpp

// Reaktoro is a unified framework for modeling chemically reactive systems.
//
// Copyright © 2014-2024 Allan Leal
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this library. If not, see <http://www.gnu.org/licenses/>.
 
#pragma once
 
// Reaktoro includes
#include <Reaktoro/Common/Meta.hpp>
#include <Reaktoro/Common/TraitsUtils.hpp>
#include <Reaktoro/Common/Types.hpp>
 
namespace Reaktoro {
 
template<typename T>
struct MemoizationTraits
{
    using Type = Decay<T>;
 
    using CacheType = Type;
 
    static auto equal(const CacheType& a, const Type& b)
    {
        return a == b;
    }
 
    static auto assign(CacheType& a, const Type& b)
    {
        a = b;
    }
};
 
namespace detail {
 
template<typename CacheType, typename T>
constexpr auto sameValue(const CacheType& a, const T& b)
{
    return MemoizationTraits<Decay<T>>::equal(a, b);
};
 
template<typename CacheType, typename T>
constexpr auto assignValue(CacheType& a, const T& b)
{
    return MemoizationTraits<Decay<T>>::assign(a, b);
};
 
template<typename Tuple1, typename Tuple2>
auto sameValues(const Tuple1& tuple1, const Tuple2& tuple2)
{
    using std::tuple_size_v;
    using std::get;
    constexpr auto N1 = tuple_size_v<Tuple1>;
    constexpr auto N2 = tuple_size_v<Tuple2>;
    static_assert(N1 == N2);
    bool res = true;
    For<N1>([&](auto i) constexpr {
        res = res && sameValue(get<i>(tuple1), get<i>(tuple2));
    });
    return res;
}
 
template<typename Tuple1, typename Tuple2>
auto assignValues(Tuple1& tuple1, const Tuple2& tuple2)
{
    using std::tuple_size_v;
    using std::get;
    constexpr auto N1 = tuple_size_v<Tuple1>;
    constexpr auto N2 = tuple_size_v<Tuple2>;
    static_assert(N1 == N2);
    For<N1>([&](auto i) constexpr {
        assignValue(get<i>(tuple1), get<i>(tuple2));
    });
}
 
template<typename T>
using CacheType = typename MemoizationTraits<Decay<T>>::CacheType;
 
} // namespace detail
 
class Memoization
{
public:
    static auto isEnabled() -> bool;
 
    static auto isDisabled() -> bool;
 
    static auto enable() -> void;
 
    static auto disable() -> void;
 
    Memoization() = delete;
};
 
template<typename Ret, typename... Args>
auto memoize(Fn<Ret(Args...)> f) -> Fn<Ret(Args...)>
{
    auto cache = std::make_shared<Map<Tuple<Args...>, Ret>>();
    return [=](Args... args) mutable -> Ret
    {
        if(Memoization::isDisabled())
            return f(args...);
        Tuple<Args...> t(args...);
        if(cache->find(t) == cache->end())
            (*cache)[t] = f(args...);
        return (*cache)[t];
    };
}
 
template<typename Fun, Requires<!isFunction<Fun>> = true>
auto memoize(Fun f)
{
    return memoize(asFunction(f));
}
 
template<typename Ret, typename... Args>
auto memoizeLast(Fn<Ret(Args...)> f) -> Fn<Ret(Args...)>
{
    Tuple<detail::CacheType<Args>...> cache;
    Ret result = Ret();
    auto firsttime = true;
    return [=](Args... args) mutable -> Ret
    {
        if(Memoization::isDisabled())
            return f(args...);
        if(detail::sameValues(cache, std::tie(args...)) && !firsttime)
            return Ret(result);
        detail::assignValues(cache, std::tie(args...));
        firsttime = false;
        return result = f(args...);
    };
}
 
template<typename Fun, Requires<!isFunction<Fun>> = true>
auto memoizeLast(Fun f)
{
    return memoizeLast(asFunction(f));
}
 
template<typename Ret, typename RetRef, typename... Args>
auto memoizeLastUsingRef(Fn<void(RetRef, Args...)> f) -> Fn<void(RetRef, Args...)>
{
    Tuple<detail::CacheType<Args>...> cache;
    Ret result = Ret();
    auto firsttime = true;
    return [=](RetRef res, Args... args) mutable -> void
    {
        if(Memoization::isDisabled())
            f(res, args...);
        else if(detail::sameValues(cache, std::tie(args...)) && !firsttime)
            res = result;
        else
        {
            f(res, args...);
            result = res;
            detail::assignValues(cache, std::tie(args...));
        }
        firsttime = false;
    };
}
 
template<typename Ret, typename Fun, Requires<!isFunction<Fun>> = true>
auto memoizeLastUsingRef(Fun f)
{
    return memoizeLastUsingRef<Ret>(asFunction(f));
}
 
template<typename Ret, typename... Args>
auto memoizeLastUsingRef(Fn<void(Ret&, Args...)> f) -> Fn<void(Ret&, Args...)>
{
    return memoizeLastUsingRef<Ret, Ret&>(f);
}
 
template<typename Fun, Requires<!isFunction<Fun>> = true>
auto memoizeLastUsingRef(Fun f)
{
    return memoizeLastUsingRef(asFunction(f));
}
 
} // namespace Reaktoro