/// @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.
/// @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
{

/// @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
    {
        const node packed = pack_cw(cw, advice, dir);
        return dpf::xor_if(kids[dir ? 1u : 0u], packed, parent_control);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static node traverse(node parent, node cw_packed, bool dir,
        bool /*is_last*/ = false) noexcept
    {
        auto kids = expand(parent, false);
        return dpf::xor_if_lo_bit(kids[dir ? 1u : 0u], 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]);
        }
    }
};

/// @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
    {
        // 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);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static node traverse(node parent, node cw_packed, bool dir,
        bool is_last = false) noexcept
    {
        auto kids = expand(parent, is_last);
        return dpf::xor_if_lo_bit(kids[dir ? 1u : 0u], 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__