/// @file dpf/tree_traits.hpp
/// @brief BGI vs Half-Tree walk policy, selected by the interior PRG.
/// @details Default traits match today's Boyle–Gilboa–Ishai tree. A PRG that
///          defines `half_tree_tag` (e.g. `prg::aes128_ccr`) opts into the
///          Guo et al. Half-Tree mid-level expand / CW / advance, with a
///          two-tweak last level that keeps BGI-style advice packing.
/// @note Default expand: Boyle, Gilboa, and Ishai, CCS 2016 (full version ePrint 2018/707).
/// @note Best known (same model), selected by `half_tree_tag`: Guo, Yang, Wang, Zhang, Xie, Zhang, and Liu, ePrint 2022/1431. Their dealer point key keeps the CCS 2016 length and the n-hash point evaluation; they state about 2n+2 random-permutation calls to generate a key versus about 4n, and 1.5N calls for a full-domain evaluation versus 2N.
/// @copyright Copyright (c) 2019-2026 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_TREE_TRAITS_HPP__
#define LIBDPF_INCLUDE_DPF_TREE_TRAITS_HPP__

#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>

#include "hedley/hedley.h"
#include "simde/simde/x86/avx2.h"
#include "portable-snippets/exact-int/exact-int.h"

#include "dpf/twiddle.hpp"
#include "dpf/utils.hpp"

namespace dpf
{

template <typename P, typename = void>
struct prg_has_indep4 : std::false_type {};

template <typename P>
struct prg_has_indep4<P, std::void_t<decltype(P::eval_indep4(
    std::declval<const typename P::block_type *>(),
    std::declval<const psnip_uint32_t *>(),
    std::declval<typename P::block_type *>()))>> : std::true_type {};

/// @brief Walk policy for interior DPF levels. Specialized when `PRG::half_tree_tag`
/// exists.
/// @tparam PRG pseudorandom generator
template <typename PRG, typename = void>
struct tree_traits
{
    using prg = PRG;
    using node = typename PRG::block_type;

    static constexpr bool is_half_tree = false;
    static constexpr bool stores_mid_advice = true;
    static constexpr bool last_level_differs = false;

    template <typename Sampler>
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1)
    static void root_init(node out[2], Sampler && sample)
    {
        out[0] = dpf::unset_lo_bit(static_cast<node>(sample()));
        out[1] = dpf::set_lo_bit(static_cast<node>(sample()));
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    static bool is_last_level(std::size_t level, std::size_t depth) noexcept
    {
        (void)level;
        (void)depth;
        return false;
    }

    /// @brief BGI expand: clear lo-2bits, then `PRG::eval01`.
    /// @param s the `s`
    /// @return BGI expand: clear lo-2bits, then `PRG::eval01`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto expand(node s, bool /*is_last*/ = false) noexcept
    {
        return PRG::eval01(dpf::unset_lo_2bits(s));
    }

/// @brief Convert / value-CW stretch. BGI: identical to `expand`.
/// @param s the node to stretch
/// @return the same value as `expand`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto expand_value(node s) noexcept
    {
        return expand(s, false);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2, 3)
    static void expand_x4(const node * HEDLEY_RESTRICT seeds,
        node * HEDLEY_RESTRICT left, node * HEDLEY_RESTRICT right,
        bool /*is_last*/ = false) noexcept
    {
        alignas(node) node cleared[4];
        DPF_UNROLL_LOOP
        for (std::size_t i = 0; i < 4; ++i)
            cleared[i] = dpf::unset_lo_2bits(seeds[i]);
        PRG::eval01_x4(cleared, left, right);
    }

    /// @brief Pack CW for eval: embed advice bit `dir` into the lo-bit.
    /// @param cw the `cw`
    /// @param advice the advice bit
    /// @param dir the `dir`
    /// @return the returned `node`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    static node pack_cw(node cw, psnip_uint8_t advice, bool dir,
        bool /*is_last*/ = false) noexcept
    {
        return dpf::set_lo_bit(cw, (advice >> dir) & 1);
    }

    /// @brief Off-path child XOR + packed advice `t0|t1` (BGI).
    /// @param cw_out the `cw_out`
    /// @param advice_out the `advice_out`
    /// @param kids0 the `kids0`
    /// @param kids1 the `kids1`
    /// @param bit the bit value or bit index
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void make_cw(node & cw_out, psnip_uint8_t & advice_out,
        const std::array<node, 2> & kids0, const std::array<node, 2> & kids1,
        const node & /*s0*/, const node & /*s1*/, bool bit,
        bool /*is_last*/ = false) noexcept
    {
        const node child[2] = {
            simde_mm_xor_si128(kids0[0], kids1[0]),
            simde_mm_xor_si128(kids0[1], kids1[1])
        };
        const bool t0 = static_cast<bool>(dpf::get_lo_bit(child[0]) ^ !bit);
        const bool t1 = static_cast<bool>(dpf::get_lo_bit(child[1]) ^ bit);
        cw_out = child[!bit];
        advice_out = static_cast<psnip_uint8_t>((t1 << 1) | t0);
    }

    /// @brief `xor_if(child[dir], pack_cw(...), parent_control)`.
    /// @param parent the parent node
    /// @param kids the `kids`
    /// @param cw the `cw`
    /// @param advice the advice bit
    /// @param dir the `dir`
    /// @param parent_control the `parent_control`
    /// @return `xor_if(child[dir], pack_cw(...), parent_control)`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    static node advance(node parent, const std::array<node, 2> & kids,
        node cw, psnip_uint8_t advice, bool dir,
        bool parent_control, bool /*is_last*/ = false) noexcept
    {
        (void)parent;
        const node packed = pack_cw(cw, advice, dir);
        return dpf::xor_if(kids[dir ? 1u : 0u], packed, parent_control);
    }

    /// @brief One child only: `PRG::eval(cleared, dir)` instead of `eval01`.
    /// @details Path walks never need the sibling. Interval/full use
    ///          `traverse01` / `traverse01_x4` when both children are live.
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static node traverse(node parent, node cw_packed, bool dir,
        bool /*is_last*/ = false) noexcept
    {
        const node child = PRG::eval(dpf::unset_lo_2bits(parent),
            static_cast<psnip_uint32_t>(dir ? 1 : 0));
        return dpf::xor_if_lo_bit(child, cw_packed, parent);
    }

    /// @brief Four independent one-child steps. Same result as `traverse`.
    /// @details Interleaves the PRG when `PRG::eval_indep4` exists so a
    ///          strided sequence is not limited by single-block AES latency.
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void traverse4(const node parents[4], const node cws[4],
        const bool dirs[4], node out[4]) noexcept
    {
        if constexpr (prg_has_indep4<PRG>::value)
        {
            node seeds[4];
            psnip_uint32_t pos[4];
            for (int i = 0; i < 4; ++i)
            {
                seeds[i] = dpf::unset_lo_2bits(parents[i]);
                pos[i] = dirs[i] ? 1u : 0u;
            }
            node kids[4];
            PRG::eval_indep4(seeds, pos, kids);
            for (int i = 0; i < 4; ++i)
                out[i] = dpf::xor_if_lo_bit(kids[i], cws[i], parents[i]);
        }
        else
        {
            for (int i = 0; i < 4; ++i)
                out[i] = traverse(parents[i], cws[i], dirs[i]);
        }
    }

    /// @brief Eight independent one-child steps. Same result as `traverse`.
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void traverse8(const node parents[8], const node cws[8],
        const bool dirs[8], node out[8]) noexcept
    {
        if constexpr (prg_has_indep4<PRG>::value)
        {
            node seeds[8];
            psnip_uint32_t pos[8];
            for (int i = 0; i < 8; ++i)
            {
                seeds[i] = dpf::unset_lo_2bits(parents[i]);
                pos[i] = dirs[i] ? 1u : 0u;
            }
            node kids[8];
            PRG::eval_indep8(seeds, pos, kids);
            for (int i = 0; i < 8; ++i)
                out[i] = dpf::xor_if_lo_bit(kids[i], cws[i], parents[i]);
        }
        else
        {
            for (int i = 0; i < 8; ++i)
                out[i] = traverse(parents[i], cws[i], dirs[i]);
        }
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto traverse01(node parent, node cw0, node cw1,
        bool is_last = false) noexcept
    {
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        auto kids = expand(parent, is_last);
        return std::array<node, 2>{
            dpf::xor_if_lo_bit(kids[0], cw0, parent),
            dpf::xor_if_lo_bit(kids[1], cw1, parent)
        };
HEDLEY_PRAGMA(GCC diagnostic pop)
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 4, 5)
    static void traverse01_x4(const node * HEDLEY_RESTRICT parents,
        node cw0, node cw1, node * HEDLEY_RESTRICT left,
        node * HEDLEY_RESTRICT right, bool is_last = false) noexcept
    {
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        expand_x4(parents, left, right, is_last);
HEDLEY_PRAGMA(GCC diagnostic pop)
        DPF_UNROLL_LOOP
        for (std::size_t i = 0; i < 4; ++i)
        {
            left[i] = dpf::xor_if_lo_bit(left[i], cw0, parents[i]);
            right[i] = dpf::xor_if_lo_bit(right[i], cw1, parents[i]);
        }
    }
};

/// @brief Half-Tree specialization (PRG advertises `half_tree_tag`).
/// @tparam PRG pseudorandom generator
template <typename PRG>
struct tree_traits<PRG, std::void_t<typename PRG::half_tree_tag>>
{
    using prg = PRG;
    using node = typename PRG::block_type;

    static constexpr bool is_half_tree = true;
    static constexpr bool stores_mid_advice = false;
    static constexpr bool last_level_differs = true;

    template <typename Sampler>
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1)
    static void root_init(node out[2], Sampler && sample)
    {
        // Shares of a fixed Δ with lsb(Δ)=1.
        const node s = dpf::unset_lo_bit(static_cast<node>(sample()));
        const node delta = dpf::set_lo_bit(static_cast<node>(sample()));
        out[0] = s;
        out[1] = simde_mm_xor_si128(s, delta);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    static bool is_last_level(std::size_t level, std::size_t depth) noexcept
    {
        return depth != 0 && level + 1 == depth;
    }

    /// @brief Mid: `{H(s), H(s)⊕s}` (keep control bit). Last: two-tweak stretch.
    /// @param s the `s`
    /// @param is_last the `is_last`
    /// @return Mid: `{H(s), H(s)⊕s}` (keep control bit)
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto expand(node s, bool is_last = false) noexcept
    {
        if (is_last)
        {
            // `{H(s|0), H(s|1)}` — LSB forced, matching Guo et al. / myl7.
            const node base = dpf::unset_lo_bit(s);
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
            return std::array<node, 2>{
                PRG::hash(base),
                PRG::hash(dpf::set_lo_bit(base))
            };
HEDLEY_PRAGMA(GCC diagnostic pop)
        }
        return PRG::eval01(s);
    }

/// @brief Convert / value-CW stretch: always two-tweak `{H(s|0), H(s|1)}`.
/// @details Seed walk mid levels keep using half-style `expand`.
/// @param s the node to stretch
/// @return the pair `{H(s|0), H(s|1)}`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto expand_value(node s) noexcept
    {
        return expand(s, true);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2, 3)
    static void expand_x4(const node * HEDLEY_RESTRICT seeds,
        node * HEDLEY_RESTRICT left, node * HEDLEY_RESTRICT right,
        bool is_last = false) noexcept
    {
        if (is_last)
        {
            DPF_UNROLL_LOOP
            for (std::size_t i = 0; i < 4; ++i)
            {
                auto kids = expand(seeds[i], true);
                left[i] = kids[0];
                right[i] = kids[1];
            }
            return;
        }
        PRG::eval01_x4(seeds, left, right);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    static node pack_cw(node cw, psnip_uint8_t advice, bool dir,
        bool is_last = false) noexcept
    {
        if (!is_last)
            return cw;  // mid: full CW, no advice packing
        return dpf::set_lo_bit(cw, (advice >> dir) & 1);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void make_cw(node & cw_out, psnip_uint8_t & advice_out,
        const std::array<node, 2> & kids0, const std::array<node, 2> & kids1,
        const node & s0, const node & s1, bool bit, bool is_last = false) noexcept
    {
        if (is_last)
        {
            // BGI-style advice on the leaf step only.
            const node child[2] = {
                simde_mm_xor_si128(kids0[0], kids1[0]),
                simde_mm_xor_si128(kids0[1], kids1[1])
            };
            const bool t0 = static_cast<bool>(dpf::get_lo_bit(child[0]) ^ !bit);
            const bool t1 = static_cast<bool>(dpf::get_lo_bit(child[1]) ^ bit);
            cw_out = child[!bit];
            advice_out = static_cast<psnip_uint8_t>((t1 << 1) | t0);
            (void)s0;
            (void)s1;
            return;
        }
        // CW = H(s0)⊕H(s1)⊕ᾱ·Δ = off-path children XOR (Half-Tree identity).
        const node child[2] = {
            simde_mm_xor_si128(kids0[0], kids1[0]),
            simde_mm_xor_si128(kids0[1], kids1[1])
        };
        cw_out = child[!bit];
        advice_out = 0;
        (void)s0;
        (void)s1;
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    static node advance(node parent, const std::array<node, 2> & kids,
        node cw, psnip_uint8_t advice, bool dir, bool parent_control,
        bool is_last = false) noexcept
    {
        (void)parent;
        // Mid and last: select child[dir], XOR packed CW if parent control set.
        // Mid Half-Tree: child[1]=H⊕s so this is `h ⊕ (dir?s:0) ⊕ (t?cw:0)`.
        const node packed = pack_cw(cw, advice, dir, is_last);
        return dpf::xor_if(kids[dir ? 1u : 0u], packed, parent_control);
    }

    /// @brief One child only. Mid: one CCR hash (+ optional ⊕s). Last: one
    ///        two-tweak hash instead of both.
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static node traverse(node parent, node cw_packed, bool dir,
        bool is_last = false) noexcept
    {
        node child;
        if (is_last)
        {
            const node base = dpf::unset_lo_bit(parent);
            child = dir ? PRG::hash(dpf::set_lo_bit(base)) : PRG::hash(base);
        }
        else
        {
            const node h = PRG::hash(parent);
            child = dir ? simde_mm_xor_si128(h, parent) : h;
        }
        return dpf::xor_if_lo_bit(child, cw_packed, parent);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto traverse01(node parent, node cw0, node cw1,
        bool is_last = false) noexcept
    {
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        auto kids = expand(parent, is_last);
        return std::array<node, 2>{
            dpf::xor_if_lo_bit(kids[0], cw0, parent),
            dpf::xor_if_lo_bit(kids[1], cw1, parent)
        };
HEDLEY_PRAGMA(GCC diagnostic pop)
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 4, 5)
    static void traverse01_x4(const node * HEDLEY_RESTRICT parents,
        node cw0, node cw1, node * HEDLEY_RESTRICT left,
        node * HEDLEY_RESTRICT right, bool is_last = false) noexcept
    {
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        expand_x4(parents, left, right, is_last);
HEDLEY_PRAGMA(GCC diagnostic pop)
        DPF_UNROLL_LOOP
        for (std::size_t i = 0; i < 4; ++i)
        {
            left[i] = dpf::xor_if_lo_bit(left[i], cw0, parents[i]);
            right[i] = dpf::xor_if_lo_bit(right[i], cw1, parents[i]);
        }
    }
};

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_TREE_TRAITS_HPP__