/// @file dpf/flute.hpp
/// @brief Lookup tables as a multi-fan-in inner product.
/// @details A public table `f : {0,1}^δ → {0,1}^σ` is the OR of the input
///          rows whose output bit is 1. Those terms are disjoint, so the OR
///          is an XOR, and the XOR of ANDs is an inner product of one vector
///          per input bit (and the public output column). Complements flip the
///          already-opened masked bit and leave the mask λ alone, so setup
///          builds subset products of the δ input masks once, not once per row.
///
///          Online, each party XORs the public terms into a share of `v` and
///          the parties exchange that share. Two parties send two bits per
///          output bit, independent of δ. The masked output `m_z` is then
///          local, and the semantic bit is `m_z XOR λ_z`.
///
///          The two-party exchange is ΠLUT. The three-party function is the
///          same algebra on three XOR shares of each mask: the paper's
///          construction is the inner product, and the share count is the
///          underlying MPC. This is not a new share domain.
/// @note Andreas Brüggemann, Robin Hundt, Thomas Schneider, Ajith Suresh, and
///       Hossein Yalame, "FLUTE: Fast and Secure Lookup Table Evaluations,"
///       IEEE S&P 2023 (ePrint 2023/499). The masked bits are the ABY2.0 wire
///       (Patra, Schneider, Suresh, and Yalame, USENIX Security 2021).
/// @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_FLUTE_HPP__
#define LIBDPF_INCLUDE_DPF_FLUTE_HPP__

#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <vector>

#include "hedley/hedley.h"

#include "dpf/random.hpp"

namespace dpf
{
namespace flute
{

inline constexpr unsigned k_max_delta = 8;

/// @brief One evaluation. `opened[w] = masked[w] XOR` the mask shares.
struct result
{
    std::vector<std::uint8_t> opened;
    std::vector<std::uint8_t> masked;
    /// @brief XOR shares of `λ_z`, one vector per party. Size 2 or 3.
    std::vector<std::vector<std::uint8_t>> mask;
    /// @brief Bits exchanged online: one bit per party per output bit.
    std::size_t online_bits = 0;
};

namespace detail
{

inline std::uint8_t bit_of(std::uint32_t row, unsigned i)
{
    return static_cast<std::uint8_t>((row >> i) & 1u);
}

/// @brief Masked literal of input `i` on row `j`. Complement flips `m` only.
inline std::uint8_t masked_lit(std::uint8_t m, std::uint8_t encoding)
{
    return encoding ? m : static_cast<std::uint8_t>(m ^ 1u);
}

inline std::uint8_t and_subset(std::uint32_t subset, std::uint32_t row,
    unsigned delta, const std::uint8_t * m)
{
    std::uint8_t acc = 1;
    for (unsigned i = 0; i < delta; ++i)
    {
        if (((subset >> i) & 1u) == 0)
            continue;
        acc = static_cast<std::uint8_t>(
            acc & masked_lit(m[i], bit_of(row, i)));
    }
    return acc;
}

inline std::uint8_t dot_column(std::uint32_t subset, unsigned delta,
    const std::uint8_t * m, const std::uint8_t * column)
{
    const std::uint32_t rows = 1u << delta;
    std::uint8_t acc = 0;
    for (std::uint32_t j = 0; j < rows; ++j)
    {
        if (column[j] == 0)
            continue;
        acc = static_cast<std::uint8_t>(
            acc ^ and_subset(subset, j, delta, m));
    }
    return acc;
}

inline std::uint8_t rand_bit()
{
    return static_cast<std::uint8_t>(dpf::uniform_sample<std::uint8_t>() & 1u);
}

struct setup
{
    std::vector<std::uint8_t> m;
    std::vector<std::uint8_t> lam;
    /// @brief `share[party][subset] =` that party's XOR share of AND_{i in subset} λ_i.
    std::vector<std::vector<std::uint8_t>> share;
    std::vector<std::vector<std::uint8_t>> lamz;
};

inline setup make_setup(unsigned delta, unsigned n_out, unsigned parties,
    const std::uint8_t * x)
{
    if (parties != 2 && parties != 3)
        throw std::invalid_argument("flute: parties");
    setup s;
    s.m.resize(delta);
    s.lam.resize(delta);
    s.share.assign(parties, std::vector<std::uint8_t>(1u << delta, 0));
    s.lam.resize(delta);
    for (unsigned i = 0; i < delta; ++i)
    {
        if (x[i] > 1)
            throw std::invalid_argument("flute: input bit");
        std::uint8_t lam = 0;
        for (unsigned p = 0; p < parties; ++p)
        {
            const std::uint8_t sh = rand_bit();
            s.share[p][1u << i] = sh;
            lam = static_cast<std::uint8_t>(lam ^ sh);
        }
        s.lam[i] = lam;
        s.m[i] = static_cast<std::uint8_t>(x[i] ^ lam);
    }
    s.share[0][0] = 1;
    const std::uint32_t nsub = 1u << delta;
    for (std::uint32_t subset = 1; subset < nsub; ++subset)
    {
        if ((subset & (subset - 1u)) == 0)
            continue;
        for (unsigned p = 0; p < parties; ++p)
        {
            // Dealer shares the AND. Party 0 holds the product of the full
            // masks adjusted by the other parties' shares of this subset.
            s.share[p][subset] = 0;
        }
        std::uint8_t prod = 1;
        for (unsigned i = 0; i < delta; ++i)
            if (((subset >> i) & 1u) != 0)
                prod = static_cast<std::uint8_t>(prod & s.lam[i]);
        for (unsigned p = 1; p < parties; ++p)
            s.share[p][subset] = rand_bit();
        std::uint8_t rest = prod;
        for (unsigned p = 1; p < parties; ++p)
            rest = static_cast<std::uint8_t>(rest ^ s.share[p][subset]);
        s.share[0][subset] = rest;
    }
    s.lamz.assign(parties, std::vector<std::uint8_t>(n_out, 0));
    for (unsigned w = 0; w < n_out; ++w)
        for (unsigned p = 0; p < parties; ++p)
            s.lamz[p][w] = rand_bit();
    return s;
}

inline std::uint8_t party_v(const setup & s, unsigned party, unsigned delta,
    std::uint32_t full, const std::uint8_t * column, std::uint8_t lamz_share)
{
    std::uint8_t v = lamz_share;
    for (std::uint32_t subset = 0; subset < full; ++subset)
    {
        const std::uint32_t rest = full ^ subset;
        const std::uint8_t t = dot_column(subset, delta, s.m.data(), column);
        v = static_cast<std::uint8_t>(
            v ^ (t & s.share[party][rest]));
    }
    return v;
}

inline result run(unsigned delta, unsigned n_out, unsigned parties,
    const std::uint8_t * columns, const std::uint8_t * x)
{
    if (delta < 1 || delta > k_max_delta)
        throw std::invalid_argument("flute: delta");
    if (n_out == 0 || columns == nullptr || x == nullptr)
        throw std::invalid_argument("flute: table");
    const std::uint32_t rows = 1u << delta;
    const std::uint32_t full = rows - 1u;
    auto s = make_setup(delta, n_out, parties, x);
    result out;
    out.mask.assign(parties, std::vector<std::uint8_t>(n_out, 0));
    out.masked.resize(n_out);
    out.opened.resize(n_out);
    out.online_bits = static_cast<std::size_t>(parties) * n_out;
    for (unsigned w = 0; w < n_out; ++w)
    {
        const std::uint8_t * column = columns + static_cast<std::size_t>(w) * rows;
        std::uint8_t v = 0;
        for (unsigned p = 0; p < parties; ++p)
        {
            const std::uint8_t share = party_v(s, p, delta, full, column, s.lamz[p][w]);
            v = static_cast<std::uint8_t>(v ^ share);
            out.mask[p][w] = s.lamz[p][w];
        }
        const std::uint8_t t_full = dot_column(full, delta, s.m.data(), column);
        out.masked[w] = static_cast<std::uint8_t>(v ^ t_full);
        std::uint8_t lam = 0;
        for (unsigned p = 0; p < parties; ++p)
            lam = static_cast<std::uint8_t>(lam ^ out.mask[p][w]);
        out.opened[w] = static_cast<std::uint8_t>(out.masked[w] ^ lam);
    }
    return out;
}

}  // namespace detail

/// @brief Clear table. `x` is δ bits, least-significant bit first.
///        `columns[w * 2^δ + row]` is output bit `w` on that row.
HEDLEY_WARN_UNUSED_RESULT
inline std::vector<std::uint8_t> eval_plain(unsigned delta, unsigned n_out,
    const std::uint8_t * columns, const std::uint8_t * x)
{
    if (delta < 1 || delta > k_max_delta)
        throw std::invalid_argument("flute: delta");
    std::uint32_t row = 0;
    for (unsigned i = 0; i < delta; ++i)
    {
        if (x[i] > 1)
            throw std::invalid_argument("flute: input bit");
        row |= static_cast<std::uint32_t>(x[i]) << i;
    }
    const std::uint32_t rows = 1u << delta;
    std::vector<std::uint8_t> out(n_out);
    for (unsigned w = 0; w < n_out; ++w)
        out[w] = columns[static_cast<std::size_t>(w) * rows + row] & 1u;
    return out;
}

/// @brief Two-party ΠLUT. Online cost is two bits per output bit.
HEDLEY_WARN_UNUSED_RESULT
inline result eval_pair(unsigned delta, unsigned n_out,
    const std::uint8_t * columns, const std::uint8_t * x)
{
    return detail::run(delta, n_out, 2, columns, x);
}

/// @brief The same ΠLUT inner product on three XOR shares of each mask.
HEDLEY_WARN_UNUSED_RESULT
inline result eval_trio(unsigned delta, unsigned n_out,
    const std::uint8_t * columns, const std::uint8_t * x)
{
    return detail::run(delta, n_out, 3, columns, x);
}

}  // namespace flute
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_FLUTE_HPP__