libdpf/include/dpf/sequence_memoizer.hpp

/// @file dpf/sequence_memoizer.hpp
/// @brief Workspaces for a `sequence_recipe` traversal.
/// @details The memoizer stores a reference to the recipe it was built from
///          and later calls must pass that same object (`std::logic_error`
///          otherwise). The factories unwrap `party_key`.
///
///          `inplace_reversing_sequence_memoizer` keeps one level.
///          `double_space_sequence_memoizer` keeps two and is the default
///          inside `eval_sequence(key, recipe, buffer)`.
///          `full_tree_sequence_memoizer` keeps every level.
/// @author Ryan Henry <ryan.henry@ucalgary.ca>
/// @author Christopher Jiang <christopher.jiang@ucalgary.ca>
/// @copyright Copyright (c) 2019-2024 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_SEQUENCE_MEMOIZER_HPP__
#define LIBDPF_INCLUDE_DPF_SEQUENCE_MEMOIZER_HPP__

#include "hedley/hedley.h"

#include <cstddef>
#include <cstring>
#include <type_traits>
#include <functional>
#include <utility>
#include <iterator>
#include <algorithm>
#include <stdexcept>
#include <optional>

#include "dpf/secret_share.hpp"
#include "dpf/sequence_recipe.hpp"

namespace dpf
{

struct sequence_memoizer_tag_ {};

template <typename DpfKey,
          typename ReturnT = typename DpfKey::interior_node *>
struct sequence_recipe_memoizer_base : public sequence_memoizer_tag_
{
  public:
    using dpf_type = DpfKey;
    using return_type = ReturnT;
    using iterator_type = return_type;
    using node_type = typename DpfKey::interior_node;
    const sequence_recipe & recipe;

    // level 0 should access the root
    // level goes up to (and including) depth
    HEDLEY_NO_THROW
    virtual return_type operator[](std::size_t) const noexcept = 0;

    // iterators should access most recently completed level
    HEDLEY_NO_THROW
    virtual return_type begin() const noexcept = 0;
    HEDLEY_NO_THROW
    virtual return_type end() const noexcept = 0;

    virtual std::size_t assign_dpf(const dpf_type & dpf, const sequence_recipe & r)
    {
        if (&recipe != &r)
        {
            throw std::logic_error("memoizer cannot be used with different recipe");
        }

        if (dpf_.has_value() == false || std::memcmp(&dpf_root_, &dpf.root(), sizeof(node_type)) != 0
            || std::memcmp(&dpf_common_part_hash_, &dpf.common_part_hash(), sizeof(digest_type)) != 0)
        {
            if (dpf.depth != recipe.depth())
            {
                throw std::logic_error("incorrect dpf depth");
            }
            this->operator[](0)[0] = dpf.root();
            dpf_ = std::cref(dpf);
            dpf_root_ = dpf.root();
            dpf_common_part_hash_ = dpf.common_part_hash();
            level_index = 1;
        }

        return level_index;
    }

    std::size_t advance_level()
    {
        return ++level_index;
    }

    std::size_t get_nodes_at_level() const
    {
        return get_nodes_at_level(level_index);
    }

    std::size_t get_nodes_at_level(std::size_t level) const
    {
        if (level == size_t(-1))
        {
            return 0;
        }

        if (level == depth)
        {
            return recipe.num_leaf_nodes();
        }

        return recipe.level_endpoints()[level+1] - recipe.level_endpoints()[level];
    }

    // returns true if first traversal should be taken
    // this usually means traversing left, but for the inplace_reversing memoizer
    // if it is working in reverse, this could be a right traversal
    virtual bool traverse_first(std::size_t step) const
    {
        return recipe.recipe_steps()[step] > int8_t(-1);
    }

    // returns true if second traversal should be taken
    // this usually means traversing right, but for the inplace_reversing memoizer
    // if it is working in reverse, this could be a left traversal
    virtual bool traverse_second(std::size_t step) const
    {
        return recipe.recipe_steps()[step] < int8_t(1);
    }

    // returns true if traversal should be done to the right
    // this usually means traversing in the same direction as supplied, but for the
    // inplace_reversing memoizer if it is working in reverse, this could be
    // the opposite of the supplied direction
    virtual bool get_direction(bool right) const
    {
        return right;
    }

  protected:
    std::size_t depth;
    std::size_t level_index;  // indicates current level being built

    explicit sequence_recipe_memoizer_base(const sequence_recipe & r)
      : recipe{r},
        depth{recipe.level_endpoints().size()-1},
        level_index{0},
        dpf_{std::nullopt}
    { }

  private:
    std::optional<std::reference_wrapper<const dpf_type>> dpf_;
    node_type dpf_root_;
    digest_type dpf_common_part_hash_;
};

namespace detail
{

template <typename ForwardIterT,
          typename ReverseIterT>
struct pointer_facade
{
  public:
    using forward_iter = ForwardIterT;
    using reverse_iter = ReverseIterT;

    using value_type = typename std::iterator_traits<ForwardIterT>::value_type;
    using reference = value_type &;
    using const_reference = const value_type &;
    using pointer = std::add_pointer_t<value_type>;
    using iterator_category = std::bidirectional_iterator_tag;
    using difference_type = std::pair<std::ptrdiff_t, std::ptrdiff_t>;

    HEDLEY_ALWAYS_INLINE
    pointer_facade(bool flip, forward_iter it, reverse_iter rit)
        : flip_{flip}, it_{it}, rit_{rit}
    { }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    reference operator*() const noexcept
    {
        return flip_ ? *rit_ : *it_;
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    pointer_facade & operator++() noexcept
    {
        ++it_;
        ++rit_;
        return *this;
    }

    HEDLEY_NO_THROW
    pointer_facade operator++(int) noexcept
    {
        auto tmp = *this;
        pointer_facade::operator++();
        return tmp;
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    pointer_facade & operator--() noexcept
    {
        --it_;
        --rit_;
        return *this;
    }

    HEDLEY_NO_THROW
    pointer_facade operator--(int) noexcept
    {
        auto tmp = *this;
        pointer_facade::operator--();
        return tmp;
    }

    pointer_facade & operator+=(std::size_t n) noexcept
    {
        it_ += n;
        rit_ += n;
        return *this;
    }

    HEDLEY_NO_THROW
    pointer_facade operator+(std::size_t n) const noexcept
    {
        return pointer_facade(flip_, it_ + n, rit_ + n);
    }

    HEDLEY_NO_THROW
    pointer_facade & operator-=(std::size_t n) noexcept
    {
        it_ -= n;
        rit_ -= n;
        return *this;
    }

    HEDLEY_NO_THROW
    pointer_facade operator-(std::size_t n) const noexcept
    {
        return pointer_facade(flip_, it_ - n, rit_ - n);
    }

    HEDLEY_NO_THROW
    difference_type operator-(pointer_facade rhs) const noexcept
    {
        return std::make_pair(it_ - rhs.it_, rit_ - rhs.rit_);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    HEDLEY_PURE
    reference operator[](std::size_t i) noexcept
    {
        return flip_ ? rit_[i] : it_[i];
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    HEDLEY_PURE
    const_reference operator[](std::size_t i) const noexcept
    {
        return flip_ ? rit_[i] : it_[i];
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr bool operator==(const pointer_facade & rhs) const noexcept
    {
        return flip_ == rhs.flip_ && it_ == rhs.it_ && rit_ == rhs.rit_;
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr bool operator!=(const pointer_facade & rhs) const noexcept
    {
        return !(*this == rhs);
    }

  private:
    bool flip_;
    forward_iter it_;
    reverse_iter rit_;
};

}  // namespace detail

/// @brief One level. The buffer is traversed in the opposite direction on
/// alternate levels.
/// @tparam DpfKey DPF key type
/// @tparam Allocator allocator type
template <typename DpfKey,
          typename Allocator = aligned_allocator<typename DpfKey::interior_node>>
struct inplace_reversing_sequence_memoizer final
  : public sequence_recipe_memoizer_base<DpfKey,
        detail::pointer_facade<typename DpfKey::interior_node *, std::reverse_iterator<typename DpfKey::interior_node *>>>
{
  public:
    using unique_ptr = typename Allocator::unique_ptr;
    using forward_iter = typename DpfKey::interior_node *;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    using reverse_iter = std::reverse_iterator<forward_iter>;
    using return_type = detail::pointer_facade<forward_iter, reverse_iter>;
  private:
    using parent = sequence_recipe_memoizer_base<DpfKey, return_type>;
HEDLEY_PRAGMA(GCC diagnostic pop)
  public:
    using parent::recipe;
    using parent::depth;
    using parent::level_index;
    using parent::get_nodes_at_level;

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    explicit inplace_reversing_sequence_memoizer(const sequence_recipe & r,
        Allocator alloc = Allocator{})
      : parent::sequence_recipe_memoizer_base(r),
        buf{alloc.allocate_unique_ptr(std::max(r.num_leaf_nodes(), std::size_t{1}))}
    { }
HEDLEY_PRAGMA(GCC diagnostic pop)

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type operator[](std::size_t level) const noexcept override
    {
        // flip false => forward traversal
        bool flip = (depth ^ level) & 1;
        // first check used to determine if previous or current level is being requested
        // second check used to determine if the last layer is being requested
        //     in which case it is setup to always return buf in normal order
        if (level == level_index-1 && level != depth)
        {
            std::size_t nodes_at_level = get_nodes_at_level(level);
            return return_type(!flip, &buf[recipe.num_leaf_nodes()-nodes_at_level],
                std::make_reverse_iterator(&buf[nodes_at_level]));
        }

        return return_type(flip, &buf[0],
            std::make_reverse_iterator(&buf[recipe.num_leaf_nodes()]));
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type begin() const noexcept override
    {
        return this->operator[](level_index - 1);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type end() const noexcept override
    {
        auto it = begin();
        it += get_nodes_at_level(level_index - 1);
        return it;
    }

    bool traverse_first(std::size_t step) const override
    {
        // flip false => forward traversal
        bool flip = (depth ^ level_index) & 1;
        step = !flip ? step : recipe.level_endpoints()[level_index] - step - 1 + recipe.level_endpoints()[level_index-1];
        return !flip ? (recipe.recipe_steps()[step] > int8_t(-1)) : (recipe.recipe_steps()[step] < int8_t(1));
    }

    bool traverse_second(std::size_t step) const override
    {
        // flip false => forward traversal
        bool flip = (depth ^ level_index) & 1;
        step = !flip ? step : recipe.level_endpoints()[level_index] - step - 1 + recipe.level_endpoints()[level_index-1];
        return !flip ? (recipe.recipe_steps()[step] < int8_t(1)) : (recipe.recipe_steps()[step] > int8_t(-1));
    }

    bool get_direction(bool right) const override
    {
        // flip false => forward traversal
        bool flip = (depth ^ level_index) & 1;
        return !flip ? right : !right;
    }

  private:
    unique_ptr buf;
};

/// @brief Two levels, so a level can be built while the previous level is still
/// intact. Default workspace for `eval_sequence` on a recipe.
/// @tparam DpfKey DPF key type
/// @tparam Allocator allocator type
template <typename DpfKey,
          typename Allocator = aligned_allocator<typename DpfKey::interior_node>>
struct double_space_sequence_memoizer final
  : public sequence_recipe_memoizer_base<DpfKey>
{
  HEDLEY_PRAGMA(GCC diagnostic push)
  HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
  private:
    using parent = sequence_recipe_memoizer_base<DpfKey>;
  HEDLEY_PRAGMA(GCC diagnostic pop)
  public:
    using unique_ptr = typename Allocator::unique_ptr;
    using return_type = typename DpfKey::interior_node *;
    using parent::recipe;
    using parent::depth;
    using parent::level_index;
    using parent::get_nodes_at_level;

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    explicit double_space_sequence_memoizer(const sequence_recipe & r, Allocator alloc = Allocator{})
      : parent::sequence_recipe_memoizer_base(r),
        buf{alloc.allocate_unique_ptr(2 * std::max(recipe.num_leaf_nodes(), std::size_t{1}))}
    { }
HEDLEY_PRAGMA(GCC diagnostic pop)

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type operator[](std::size_t level) const noexcept override
    {
        auto b = (depth ^ level) & 1;
        return Allocator::assume_aligned(&buf[recipe.num_leaf_nodes()*b]);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type begin() const noexcept override
    {
        return this->operator[](level_index - 1);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type end() const noexcept override
    {
        return this->operator[](level_index - 1) + get_nodes_at_level(level_index - 1);
    }

  private:
    unique_ptr buf;
};

/// @brief Every level of the recipe's traversal.
/// @tparam DpfKey DPF key type
/// @tparam Allocator allocator type
template <typename DpfKey,
          typename Allocator = aligned_allocator<typename DpfKey::interior_node>>
struct full_tree_sequence_memoizer final
  : public sequence_recipe_memoizer_base<DpfKey>
{
  HEDLEY_PRAGMA(GCC diagnostic push)
  HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
  private:
    using parent = sequence_recipe_memoizer_base<DpfKey>;
  HEDLEY_PRAGMA(GCC diagnostic pop)
  public:
    using unique_ptr = typename Allocator::unique_ptr;
    using return_type = typename DpfKey::interior_node *;
    using parent::recipe;
    using parent::level_index;
    using parent::get_nodes_at_level;

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    explicit full_tree_sequence_memoizer(const sequence_recipe & r, Allocator alloc = Allocator{})
      : parent::sequence_recipe_memoizer_base(r),
        buf{alloc.allocate_unique_ptr(std::max(
            recipe.level_endpoints()[recipe.level_endpoints().size()-1] + recipe.num_leaf_nodes(),
            std::size_t{1}))}
    { }
HEDLEY_PRAGMA(GCC diagnostic pop)

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type operator[](std::size_t level) const noexcept override
    {
        return Allocator::assume_aligned(&buf[recipe.level_endpoints()[level]]);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type begin() const noexcept override
    {
        return this->operator[](level_index - 1);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    return_type end() const noexcept override
    {
        return this->operator[](level_index - 1) + get_nodes_at_level(level_index - 1);
    }

  private:
    unique_ptr buf;
};

namespace detail
{

template <typename MemoizerT>
HEDLEY_ALWAYS_INLINE
auto make_sequence_memoizer(const sequence_recipe & recipe)
{
    return MemoizerT(recipe);
}

}  // namespace detail

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
/// @brief One-level sequence workspace bound to `recipe`.
/// @tparam DpfKey DPF key type
/// @param recipe The object later passed to `eval_sequence`. The memoizer
///        holds a reference to it.
/// @snippet evaluation/memoizers.cpp sequence-memoizer
/// @return One-level sequence workspace bound to `recipe`
template <typename DpfKey>
inline auto make_inplace_reversing_sequence_memoizer(const sequence_recipe & recipe)
{
    using key_t = unwrap_party_key_t<DpfKey>;
    return detail::make_sequence_memoizer<inplace_reversing_sequence_memoizer<key_t>>(recipe);
}

template <typename DpfKey>
inline auto make_inplace_reversing_sequence_memoizer(const DpfKey &, const sequence_recipe & recipe)
{
    return make_inplace_reversing_sequence_memoizer<DpfKey>(recipe);
}

/// @brief Two-level sequence workspace bound to `recipe`.
/// @snippet evaluation/eval_sequence.cpp eval-sequence-recipe
/// @tparam DpfKey DPF key type
/// @param recipe the sequence recipe the memoizer was built from
/// @return Two-level sequence workspace bound to `recipe`
template <typename DpfKey>
inline auto make_double_space_sequence_memoizer(const sequence_recipe & recipe)
{
    using key_t = unwrap_party_key_t<DpfKey>;
    return detail::make_sequence_memoizer<double_space_sequence_memoizer<key_t>>(recipe);
}

template <typename DpfKey>
inline auto make_double_space_sequence_memoizer(const DpfKey &, const sequence_recipe & recipe)
{
    return make_double_space_sequence_memoizer<DpfKey>(recipe);
}

/// @brief Full-tree sequence workspace bound to `recipe`.
/// @tparam DpfKey DPF key type
/// @param recipe the sequence recipe the memoizer was built from
/// @return Full-tree sequence workspace bound to `recipe`
template <typename DpfKey>
inline auto make_full_tree_sequence_memoizer(const sequence_recipe & recipe)
{
    using key_t = unwrap_party_key_t<DpfKey>;
    return detail::make_sequence_memoizer<full_tree_sequence_memoizer<key_t>>(recipe);
}
HEDLEY_PRAGMA(GCC diagnostic pop)

template <typename DpfKey>
inline auto make_full_tree_sequence_memoizer(const DpfKey &, const sequence_recipe & recipe)
{
    return make_full_tree_sequence_memoizer<DpfKey>(recipe);
}

}  // namespace dpf

namespace std
{

template <typename Iterator>
struct iterator_traits<dpf::detail::pointer_facade<Iterator, std::reverse_iterator<Iterator>>>
{
  private:
    using type = dpf::detail::pointer_facade<Iterator, std::reverse_iterator<Iterator>>;
  public:
    using iterator_category = typename type::iterator_category;
    using difference_type = typename type::difference_type;
    using value_type = typename type::value_type;
    using reference = typename type::reference;
    using const_reference = typename type::const_reference;
    using pointer = typename type::pointer;
};

}  // namespace std

#endif  // LIBDPF_INCLUDE_DPF_SEQUENCE_MEMOIZER_HPP__