#include <gtest/gtest.h>

#include <cstdint>
#include <stdexcept>
#include <tuple>
#include <vector>

#include "dpf/beaver.hpp"
#include "dpf/buffered_prg.hpp"
#include "dpf/modint.hpp"

namespace
{

struct Counter
{
    int draws = 0;

    std::uint64_t operator()()
    {
        ++draws;
        return 0x9e3779b97f4a7c15ull * static_cast<std::uint64_t>(draws);
    }
};

template <typename Ring>
struct Seq
{
    unsigned n = 1;

    Ring operator()()
    {
        return Ring{static_cast<typename Ring::integral_type>(n++ * 17u + 3u)};
    }
};

template <typename Ring>
struct XorSeq
{
    unsigned n = 1;

    Ring operator()()
    {
        using u = typename Ring::value_type;
        return Ring{static_cast<u>(n++ * 0x9e3779b9u)};
    }
};

using u64 = std::uint64_t;
using session64 = dpf::beavers::session<u64>;
using wire64 = session64::wire;

struct PolyRun
{
    u64 value = 0;
    int rounds = 0;
    std::size_t prep = 0;
    std::size_t monos = 0;
    std::size_t wires = 0;
    int draws = 0;
};

template <typename Formula>
PolyRun run_poly(std::size_t ninputs, Formula formula, const std::vector<u64> & secrets)
{
    session64 s;
    std::vector<wire64> in;
    in.reserve(ninputs);
    for (std::size_t i = 0; i < ninputs; ++i)
        in.push_back(s.input());
    auto y = formula(s, in);
    PolyRun out;
    out.rounds = s.round_of(y);
    out.prep = s.preprocessing_count();
    out.monos = s.monomial_count();
    out.wires = s.wire_count();
    Counter rng;
    s.sample(rng);
    out.draws = rng.draws;
    for (std::size_t i = 0; i < ninputs; ++i)
        s.bind(in[i], secrets[i], rng);
    s.evaluate();
    out.value = s.open(y);
    return out;
}

u64 mpow(u64 base, unsigned exp)
{
    u64 acc = 1;
    for (unsigned i = 0; i < exp; ++i)
        acc *= base;
    return acc;
}

}  // namespace

TEST(Beaver, ProductTwo)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * y);
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(rng.draws, 5);  // two input blinds * 2 draws + one product share
    EXPECT_EQ(s.monomial_count(), 1u);
    EXPECT_EQ(s.round_of(z), 1);
    EXPECT_EQ(s.monomial({{x, 1u}, {y, 1u}}).open(),
        s.lambda(x).open() * s.lambda(y).open());

    s.bind(x, u64{7}, rng);
    s.bind(y, u64{9}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z), 63u);
    EXPECT_EQ(s.delta(x), u64{7} + s.lambda(x).open());
    EXPECT_EQ(dpf::reconstruct(s.value(z).party0(), s.value(z).party1()), 63u);
    EXPECT_NE(s.value(z).p0, s.open(z));
}

TEST(Beaver, ProductThreeAndSquare)
{
    dpf::beavers::session<u64> s;
    auto a = s.input();
    auto b = s.input();
    auto c = s.input();
    auto p = s.product(a, b, c);
    auto sq = s(b * b);
    Counter rng;
    s.sample(rng);
    // wires a,b,c,p,sq = 5 * 2, monomials of a*b*c are the three pairs + abc
    // (singletons live on the wires), plus λb²
    EXPECT_EQ(s.monomial_count(), 5u);
    EXPECT_EQ(s.round_of(p), 1);
    EXPECT_EQ(s.round_of(sq), 1);
    s.bind(a, u64{2}, rng);
    s.bind(b, u64{3}, rng);
    s.bind(c, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(p), 30u);
    EXPECT_EQ(s.open(sq), 9u);
}

TEST(Beaver, MulSquareIsOneRound)
{
    dpf::beavers::session<u64> s;
    auto a = s.input();
    auto x = s.input();
    auto z = s(a * x * x);
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(s.round_of(z), 1);
    EXPECT_EQ(s.wire_count(), 3u);
    // λx², λa λx, λa λx². One blind for both x factors.
    EXPECT_EQ(s.monomial_count(), 3u);
    EXPECT_EQ(rng.draws, 7);  // two input blinds * 2 + three product shares
    EXPECT_EQ(s.monomial({{x, 2u}}).open(),
        s.lambda(x).open() * s.lambda(x).open());
    EXPECT_EQ(s.monomial({{a, 1u}, {x, 2u}}).open(),
        s.lambda(a).open() * s.lambda(x).open() * s.lambda(x).open());
    s.bind(a, u64{4}, rng);
    s.bind(x, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z), 100u);
}

TEST(Beaver, ChainedSquareUsesTwoRounds)
{
    dpf::beavers::session<u64> s;
    auto a = s.input();
    auto x = s.input();
    auto x2 = s(x * x);
    auto z = s(a * x2);
    EXPECT_EQ(s.round_of(x2), 1);
    EXPECT_EQ(s.round_of(z), 2);
    EXPECT_EQ(s.monomial_count(), 2u);
    Counter rng;
    s.sample(rng);
    s.bind(a, u64{4}, rng);
    s.bind(x, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(x2), 25u);
    EXPECT_EQ(s.open(z), 100u);
}

TEST(Beaver, DotAggregatesCrossTerm)
{
    dpf::beavers::session<u64> s;
    auto x0 = s.input();
    auto x1 = s.input();
    auto x2 = s.input();
    auto y0 = s.input();
    auto y1 = s.input();
    auto y2 = s.input();
    auto z = s.dot({x0, x1, x2}, {y0, y1, y2});
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(s.monomial_count(), 0u);
    EXPECT_EQ(s.round_of(z), 1);
    // 6 input blinds, two draws each, plus one aggregated cross draw
    EXPECT_EQ(rng.draws, 13);
    auto cross = s.lambda(x0).open() * s.lambda(y0).open()
        + s.lambda(x1).open() * s.lambda(y1).open()
        + s.lambda(x2).open() * s.lambda(y2).open();
    EXPECT_EQ(s.dot_cross(z).open(), cross);
    s.bind(x0, u64{1}, rng);
    s.bind(x1, u64{2}, rng);
    s.bind(x2, u64{3}, rng);
    s.bind(y0, u64{4}, rng);
    s.bind(y1, u64{5}, rng);
    s.bind(y2, u64{6}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z), 32u);
}

TEST(Beaver, DotReusesAPair)
{
    dpf::beavers::session<u64> s;
    auto a = s.input();
    auto b = s.input();
    auto z = s.dot({a, a}, {b, b});
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(s.monomial_count(), 0u);
    s.bind(a, u64{3}, rng);
    s.bind(b, u64{4}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z), 24u);
}

TEST(Beaver, ScaleSharesScalarBlind)
{
    dpf::beavers::session<u64> s;
    auto scalar = s.input();
    auto v0 = s.input();
    auto v1 = s.input();
    auto v2 = s.input();
    auto v3 = s.input();
    auto z = s.scale(scalar, {v0, v1, v2, v3});
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(z.size(), 4u);
    EXPECT_EQ(s.monomial_count(), 4u);
    EXPECT_EQ(s.round_of(z[0]), 1);
    // 1 scalar + 4 lanes + 4 outputs = 9 wires * 2, plus 4 cross terms
    EXPECT_EQ(rng.draws, 14);
    s.bind(scalar, u64{9}, rng);
    s.bind(v0, u64{0}, rng);
    s.bind(v1, u64{1}, rng);
    s.bind(v2, u64{0}, rng);
    s.bind(v3, u64{0}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z[0]), 0u);
    EXPECT_EQ(s.open(z[1]), 9u);
    EXPECT_EQ(s.open(z[2]), 0u);
    EXPECT_EQ(s.open(z[3]), 0u);
}

TEST(Beaver, ScaleDedupsARepeatedLane)
{
    dpf::beavers::session<u64> s;
    auto scalar = s.input();
    auto v = s.input();
    auto z = s.scale(scalar, {v, v});
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(s.monomial_count(), 1u);
    EXPECT_EQ(rng.draws, 5);
    s.bind(scalar, u64{6}, rng);
    s.bind(v, u64{7}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z[0]), 42u);
    EXPECT_EQ(s.open(z[1]), 42u);
}

TEST(Beaver, BitMulAndMux)
{
    dpf::beavers::session<u64> s;
    auto b = s.bit();
    auto x = s.input();
    auto y = s.input();
    auto prod = s.bit_mul(b, x);
    auto picked = s.mux(b, x, y);
    Counter rng;
    s.sample(rng);
    s.bind(b, u64{1}, rng);
    s.bind(x, u64{42}, rng);
    s.bind(y, u64{7}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(prod), 42u);
    EXPECT_EQ(s.open(picked), 42u);
    EXPECT_EQ(s.round_of(prod), 1);
    EXPECT_EQ(s.round_of(picked), 1);

    dpf::beavers::session<u64> t;
    auto b0 = t.bit();
    auto x0 = t.input();
    auto y0 = t.input();
    auto prod0 = t.bit_mul(b0, x0);
    auto picked0 = t.mux(b0, x0, y0);
    t.sample(rng);
    t.bind(b0, u64{0}, rng);
    t.bind(x0, u64{42}, rng);
    t.bind(y0, u64{7}, rng);
    t.evaluate();
    EXPECT_EQ(t.open(prod0), 0u);
    EXPECT_EQ(t.open(picked0), 7u);
}

TEST(Beaver, BatchOneRoundSharesOperands)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s.input();
    auto xy = s(x * y);
    auto yz = s(y * z);
    auto xx = s(x * x);
    EXPECT_EQ(s.monomial_count(), 3u);
    EXPECT_EQ(s.round_of(xy), 1);
    EXPECT_EQ(s.round_of(yz), 1);
    EXPECT_EQ(s.round_of(xx), 1);
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(rng.draws, 9);
    s.bind(x, u64{2}, rng);
    s.bind(y, u64{3}, rng);
    s.bind(z, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(xy), 6u);
    EXPECT_EQ(s.open(yz), 15u);
    EXPECT_EQ(s.open(xx), 4u);
    // y's blind was opened once and serves both products
    EXPECT_EQ(s.delta(y), u64{3} + s.lambda(y).open());
}

TEST(Beaver, BlindSurvivesALaterRound)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto xy = s(x * y);
    s.pin(xy);
    Counter rng;
    s.sample(rng);
    auto lx = s.lambda(x);
    int draws = rng.draws;
    s.sample(rng);
    EXPECT_EQ(rng.draws, draws);
    EXPECT_EQ(s.lambda(x), lx);
    s.bind(x, u64{4}, rng);
    s.bind(y, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(xy), 20u);
    auto lxy = s.lambda(xy);
    draws = rng.draws;

    // Next batch: z is new, xy already has a blind. Depth is 2.
    auto z = s.input();
    auto xyz = s(xy * z);
    EXPECT_EQ(s.round_of(xyz), 2);
    s.sample(rng);
    EXPECT_EQ(s.lambda(x), lx);
    EXPECT_EQ(s.lambda(xy), lxy);
    EXPECT_EQ(s.monomial_count(), 2u);
    // z's blind (2 draws) and one new product λ_xy λ_z
    EXPECT_EQ(rng.draws, draws + 3);
    s.bind(z, u64{6}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(xy), 20u);
    EXPECT_EQ(s.open(xyz), 120u);
    EXPECT_EQ(s.lambda(x), lx);
}

TEST(Beaver, XorWrapperAndMux)
{
    using W = dpf::xor_wrapper<std::uint32_t>;
    dpf::beavers::session<W> s;
    auto x = s.input();
    auto y = s.input();
    auto bit = s.bit();
    auto prod = s(x * y);
    auto sq = s(x * x);
    auto picked = s.mux(bit, x, y);
    XorSeq<W> rng;
    s.sample(rng);
    s.bind(x, W{0b1100u}, rng);
    s.bind(y, W{0b1010u}, rng);
    s.bind(bit, dpf::beavers::ring_traits<W>::one(), rng);
    s.evaluate();
    EXPECT_EQ(s.open(prod), W{0b1000u});
    EXPECT_EQ(s.open(sq), W{0b1100u});
    EXPECT_EQ(s.open(picked), W{0b1100u});
}

TEST(Beaver, ModintProduct)
{
    using M = dpf::modint<17>;
    dpf::beavers::session<M> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s.mul_square(x, y);
    Seq<M> rng;
    s.sample(rng);
    s.bind(x, M{6}, rng);
    s.bind(y, M{5}, rng);
    s.evaluate();
    EXPECT_EQ(static_cast<M::integral_type>(s.open(z)),
        static_cast<M::integral_type>(M{6} * M{5} * M{5}));
}

TEST(Beaver, OneShotTriplesMatchProducts)
{
    Counter rng;
    auto p2 = dpf::beavers::sample_beaver2<u64>(rng);
    EXPECT_EQ(p2.ab.open(), p2.a.open() * p2.b.open());

    auto p3 = dpf::beavers::sample_beaver3<u64>(rng);
    EXPECT_EQ(p3.ab.open(), p3.a.open() * p3.b.open());
    EXPECT_EQ(p3.ac.open(), p3.a.open() * p3.c.open());
    EXPECT_EQ(p3.bc.open(), p3.b.open() * p3.c.open());
    EXPECT_EQ(p3.abc.open(), p3.a.open() * p3.b.open() * p3.c.open());

    auto sq = dpf::beavers::sample_square<u64>(rng);
    EXPECT_EQ(sq.x2.open(), sq.x.open() * sq.x.open());

    auto ax = dpf::beavers::sample_mul_square<u64>(rng);
    EXPECT_EQ(ax.x2.open(), ax.x.open() * ax.x.open());
    EXPECT_EQ(ax.ax.open(), ax.a.open() * ax.x.open());
    EXPECT_EQ(ax.ax2.open(), ax.a.open() * ax.x.open() * ax.x.open());

    auto dot = dpf::beavers::sample_dot<u64>(3, rng);
    u64 cross = 0;
    for (std::size_t i = 0; i < 3; ++i)
        cross += dot.x[i].open() * dot.y[i].open();
    EXPECT_EQ(dot.cross.open(), cross);

    auto sc = dpf::beavers::sample_scale<u64>(3, rng);
    for (std::size_t i = 0; i < 3; ++i)
        EXPECT_EQ(sc.cross[i].open(), sc.scalar.open() * sc.lanes[i].open());

    auto bm = dpf::beavers::sample_bit_mul<u64>(rng);
    EXPECT_EQ(bm.product.open(), bm.bit.open() * bm.scalar.open());

    auto mx = dpf::beavers::sample_mux<u64>(rng);
    EXPECT_EQ(mx.bit_when1.open(), mx.bit.open() * mx.when1.open());
    EXPECT_EQ(mx.bit_when0.open(), mx.bit.open() * mx.when0.open());

    auto fresh = dpf::beavers::sample_fresh<4, u64>(rng);
    EXPECT_EQ(fresh.subset[(1u << 4) - 2].open(),
        fresh.in[0].open() * fresh.in[1].open() * fresh.in[2].open()
            * fresh.in[3].open());
}

TEST(Beaver, DefaultSamplerOpens)
{
    auto t = dpf::beavers::sample_beaver2<u64>();
    EXPECT_EQ(t.ab.open(), t.a.open() * t.b.open());
}

TEST(BufferedPrg, SeedReplaysLanesInAnyOrder)
{
    using block = dpf::prg::aes128::block_type;
    block seed = simde_mm_set_epi64x(0x1111, 0x2222);
    dpf::beavers::oracle<u64> first(seed, 4);
    dpf::beavers::oracle<u64> second(seed, 4);
    u64 late = first.blind(1, 10);
    u64 early = first.blind(0, 0);
    EXPECT_EQ(second.blind(1, 10), late);
    EXPECT_EQ(second.blind(0, 0), early);
    EXPECT_EQ(first.mask(0, 10), second.mask(0, 10));
    EXPECT_NE(first.blind(0, 0), first.blind(1, 0));

    dpf::randomness::aes_buffered_prg<u64, u64> streamed(seed, 8);
    u64 s0 = streamed.get<0>();
    u64 s1 = streamed.get<1>();
    dpf::randomness::aes_buffered_prg<u64, u64> replay(seed, 8);
    EXPECT_EQ(replay.at<0>(0), s0);
    EXPECT_EQ(replay.at<1>(0), s1);
    EXPECT_EQ(replay.get<0>(), s0);

    dpf::randomness::buffered_prg<dpf::prg::dummy, u64> dummy_stream(seed, 4);
    u64 d0 = dummy_stream.get<0>();
    dpf::randomness::buffered_prg<dpf::prg::dummy, u64> dummy_replay(seed, 4);
    EXPECT_EQ(dummy_replay.at<0>(0), d0);
    EXPECT_EQ(dummy_replay.at<0>(2), d0);
}

TEST(Beaver, OracleCopiesStayOnOneLanePerBlind)
{
    using block = dpf::prg::aes128::block_type;
    block seed = simde_mm_set_epi64x(0x9, 0x9);
    dpf::beavers::oracle<u64> src(seed, 4);
    dpf::beavers::beaver2<u64> bulk[32];
    dpf::beavers::fill_beaver2(src, 0, bulk, 32);

    dpf::beavers::oracle<u64> again(seed, 4);
    auto at31 = dpf::beavers::beaver2_at(again, 31);
    auto at4 = dpf::beavers::beaver2_at(src, 4);
    EXPECT_EQ(bulk[31].a, at31.a);
    EXPECT_EQ(bulk[31].b, at31.b);
    EXPECT_EQ(bulk[31].ab, at31.ab);
    EXPECT_EQ(bulk[31].out, at31.out);
    EXPECT_EQ(bulk[4].a, at4.a);
    EXPECT_EQ(bulk[4].ab, at4.ab);
    EXPECT_EQ(bulk[0].ab.open(), bulk[0].a.open() * bulk[0].b.open());
    EXPECT_EQ(bulk[31].ab.open(), bulk[31].a.open() * bulk[31].b.open());

    dpf::beavers::session<u64> formula;
    auto x = formula.input();
    auto y = formula.input();
    auto z = formula(x * y);
    auto m = formula.material_at(src, 7);
    auto one = dpf::beavers::beaver2_at(src, 7);
    EXPECT_EQ(m.lambda[0], one.a);
    EXPECT_EQ(m.lambda[1], one.b);
    EXPECT_EQ(m.lambda[2], one.out);
    EXPECT_EQ(m.bundles[0].open(), one.ab.open());
    EXPECT_EQ(m.lambda[0], bulk[7].a);

    formula.sample_from(src, 7);
    formula.bind(x, u64{6});
    formula.bind(y, u64{7});
    formula.evaluate();
    EXPECT_EQ(formula.open(z), 42u);
    EXPECT_EQ(formula.lambda(x), one.a);
}

TEST(Beaver, GrottoAppendixEPreprocessing)
{
    dpf::beavers::session<u64> linear;
    auto x = linear.input();
    auto sgn = linear.input();
    auto a0 = linear.input();
    auto a1 = linear.input();
    auto lin = linear(sgn * (a1 * x + a0));
    // Two-round column of Table 3: the sign is a later multiply, and the
    // constant coefficient is added from its value share.
    EXPECT_EQ(linear.round_of(lin), 2);
    EXPECT_EQ(linear.preprocessing_count(), 6u);
    Counter lin_rng;
    linear.sample(lin_rng);
    linear.bind(x, u64{2}, lin_rng);
    linear.bind(sgn, u64{3}, lin_rng);
    linear.bind(a0, u64{4}, lin_rng);
    linear.bind(a1, u64{5}, lin_rng);
    linear.evaluate();
    EXPECT_EQ(linear.open(lin), 3u * (5u * 2u + 4u));

    dpf::beavers::session<u64> quad;
    auto x2 = quad.input();
    auto s2 = quad.input();
    auto b0 = quad.input();
    auto b1 = quad.input();
    auto b2 = quad.input();
    auto q = quad(s2 * (b2 * pow(x2, 2) + b1 * x2 + b0));
    EXPECT_EQ(quad.round_of(q), 2);
    EXPECT_EQ(quad.preprocessing_count(), 9u);
    Counter quad_rng;
    quad.sample(quad_rng);
    quad.bind(x2, u64{2}, quad_rng);
    quad.bind(s2, u64{3}, quad_rng);
    quad.bind(b0, u64{4}, quad_rng);
    quad.bind(b1, u64{5}, quad_rng);
    quad.bind(b2, u64{6}, quad_rng);
    quad.evaluate();
    EXPECT_EQ(quad.open(q), 3u * (6u * 4u + 5u * 2u + 4u));

    dpf::beavers::session<u64> cube;
    auto x3 = cube.input();
    auto s3 = cube.input();
    auto c0 = cube.input();
    auto c1 = cube.input();
    auto c2 = cube.input();
    auto c3 = cube.input();
    auto y = cube(s3 * (c3 * pow(x3, 3) + c2 * pow(x3, 2) + c1 * x3 + c0));
    EXPECT_EQ(cube.round_of(y), 2);
    // Table 3 lists 13. The constant coefficient is a value share, so its
    // mask is not part of the preprocessing.
    EXPECT_EQ(cube.preprocessing_count(), 12u);

    Counter rng;
    cube.sample(rng);
    cube.bind(x3, u64{2}, rng);
    cube.bind(s3, u64{3}, rng);
    cube.bind(c0, u64{4}, rng);
    cube.bind(c1, u64{5}, rng);
    cube.bind(c2, u64{6}, rng);
    cube.bind(c3, u64{7}, rng);
    cube.evaluate();
    EXPECT_EQ(cube.open(y), 3u * (7u * 8u + 6u * 4u + 5u * 2u + 4u));
}

TEST(Beaver, MultivariatePolynomialsSharePowers)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s.input();
    auto sgn = s.input();
    auto p = s(u64{2} + u64{3} * x + u64{4} * y + u64{5} * x * y
        + u64{6} * pow(x, 2) + pow(x, 2) * y + monomial(u64{9}, x, y, z));
    EXPECT_EQ(s.round_of(p), 1);
    // Fused buckets: fewer shares than one subset product per monomial.
    EXPECT_EQ(s.monomial_count(), 5u);
    auto q = s(sgn * (x * y + pow(x, 2)));
    EXPECT_EQ(s.round_of(q), 2);
    // The sign is a later multiply, so the new shares are that product
    // rather than a second one-round crossing of every power.
    EXPECT_EQ(s.monomial_count(), 8u);
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{2}, rng);
    s.bind(y, u64{3}, rng);
    s.bind(z, u64{4}, rng);
    s.bind(sgn, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(p),
        2u + 3u * 2u + 4u * 3u + 5u * 2u * 3u + 6u * 4u + 4u * 3u + 9u * 2u * 3u * 4u);
    EXPECT_EQ(s.open(q), 5u * (2u * 3u + 4u));
}

TEST(Beaver, ElementaryLogFactorsMatchingPowers)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto z = s.input();
    auto a0 = s.input();
    auto a1 = s.input();
    auto a2 = s.input();
    auto a3 = s.input();
    auto y = s(a3 * pow(x, 3) * pow(z, 3) + a2 * pow(x, 2) * pow(z, 2)
        + a1 * x * z + a0);
    EXPECT_EQ(s.round_of(y), 2);
    // m = x*z, then a cubic in m. Same preprocessing as the two-round cubic.
    EXPECT_EQ(s.preprocessing_count(), 12u);
    EXPECT_EQ(s.monomial_count(), 6u);
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{2}, rng);
    s.bind(z, u64{3}, rng);
    s.bind(a0, u64{1}, rng);
    s.bind(a1, u64{1}, rng);
    s.bind(a2, u64{1}, rng);
    s.bind(a3, u64{1}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y), 259u);
}

TEST(Beaver, ElementaryExpScalesAfterTheCubic)
{
    dpf::beavers::session<u64> s;
    auto r = s.input();
    auto c = s.input();
    auto a0 = s.input();
    auto a1 = s.input();
    auto a2 = s.input();
    auto a3 = s.input();
    auto y = s(c * (a3 * pow(r, 3) + a2 * pow(r, 2) + a1 * r + a0));
    EXPECT_EQ(s.round_of(y), 2);
    EXPECT_EQ(s.preprocessing_count(), 12u);
    Counter rng;
    s.sample(rng);
    s.bind(r, u64{2}, rng);
    s.bind(c, u64{3}, rng);
    s.bind(a0, u64{4}, rng);
    s.bind(a1, u64{5}, rng);
    s.bind(a2, u64{6}, rng);
    s.bind(a3, u64{7}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y), 282u);
}

TEST(Beaver, ElementaryScaleAfterFactoredPower)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto z = s.input();
    auto sgn = s.input();
    auto a0 = s.input();
    auto a1 = s.input();
    auto a2 = s.input();
    auto a3 = s.input();
    auto y = s(sgn * (a3 * pow(x, 3) * pow(z, 3) + a2 * pow(x, 2) * pow(z, 2)
        + a1 * x * z + a0));
    EXPECT_EQ(s.round_of(y), 3);
    EXPECT_EQ(s.preprocessing_count(), 15u);
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{2}, rng);
    s.bind(z, u64{3}, rng);
    s.bind(sgn, u64{3}, rng);
    s.bind(a0, u64{1}, rng);
    s.bind(a1, u64{1}, rng);
    s.bind(a2, u64{1}, rng);
    s.bind(a3, u64{1}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y), 3u * 259u);
}

TEST(Beaver, PolynomialsSharePowers)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto quad = s(u64{1} + u64{2} * x + u64{3} * pow(x, 2));
    auto cube = s.horner(x, {u64{4}, u64{0}, u64{5}, u64{6}});
    EXPECT_EQ(s.round_of(quad), 1);
    EXPECT_EQ(s.round_of(cube), 1);
    // λx² is shared. The cubic also needs λx³, and one fused bucket.
    EXPECT_EQ(s.monomial_count(), 3u);
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{3}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(quad), 1u + 2u * 3u + 3u * 9u);
    EXPECT_EQ(s.open(cube), 4u + 5u * 9u + 6u * 27u);
}

TEST(Beaver, SignMultipliesThePolynomialInOneRound)
{
    dpf::beavers::session<u64> s;
    auto sgn = s.input();
    auto x = s.input();
    auto y = s(sgn * (u64{1} + x + pow(x, 2)));
    auto z = s.horner(sgn, x, {u64{1}, u64{1}, u64{1}});
    EXPECT_EQ(s.round_of(y), 1);
    EXPECT_EQ(s.round_of(z), 1);
    // λx², λsgn λx, λsgn λx². The two polynomials share them.
    EXPECT_EQ(s.monomial_count(), 3u);
    Counter rng;
    s.sample(rng);
    s.bind(sgn, u64{2}, rng);
    s.bind(x, u64{3}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y), 2u * (1u + 3u + 9u));
    EXPECT_EQ(s.open(z), s.open(y));
}

TEST(Beaver, LikeTermsCollapse)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto p = s(x * y + y * x);
    auto q = s(u64{5} * pow(x, 2) + u64{3} * x * x);
    EXPECT_EQ(s.monomial_count(), 2u);
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{4}, rng);
    s.bind(y, u64{6}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(p), 2u * 4u * 6u);
    EXPECT_EQ(s.open(q), 8u * 16u);
}

TEST(Beaver, ElementaryLogAgreesWithAHandSplit)
{
    const u64 x = 4, z = 5, a0 = 6, a1 = 7, a2 = 8, a3 = 9;
    const u64 m = x * z;
    const u64 expect = a0 + a1 * m + a2 * m * m + a3 * m * m * m;

    session64 automatic;
    auto ax = automatic.input();
    auto az = automatic.input();
    auto aa0 = automatic.input();
    auto aa1 = automatic.input();
    auto aa2 = automatic.input();
    auto aa3 = automatic.input();
    auto ay = automatic(aa3 * pow(ax, 3) * pow(az, 3) + aa2 * pow(ax, 2) * pow(az, 2)
        + aa1 * ax * az + aa0);

    session64 manual;
    auto mx = manual.input();
    auto mz = manual.input();
    auto ma0 = manual.input();
    auto ma1 = manual.input();
    auto ma2 = manual.input();
    auto ma3 = manual.input();
    auto mm = manual(mx * mz);
    auto my = manual(ma3 * pow(mm, 3) + ma2 * pow(mm, 2) + ma1 * mm + ma0);

    EXPECT_EQ(automatic.round_of(ay), manual.round_of(my));
    EXPECT_EQ(automatic.preprocessing_count(), manual.preprocessing_count());
    EXPECT_EQ(automatic.monomial_count(), manual.monomial_count());
    EXPECT_EQ(automatic.wire_count(), manual.wire_count());

    Counter ar, mr;
    automatic.sample(ar);
    manual.sample(mr);
    EXPECT_EQ(ar.draws, mr.draws);
    automatic.bind(ax, x, ar);
    automatic.bind(az, z, ar);
    automatic.bind(aa0, a0, ar);
    automatic.bind(aa1, a1, ar);
    automatic.bind(aa2, a2, ar);
    automatic.bind(aa3, a3, ar);
    manual.bind(mx, x, mr);
    manual.bind(mz, z, mr);
    manual.bind(ma0, a0, mr);
    manual.bind(ma1, a1, mr);
    manual.bind(ma2, a2, mr);
    manual.bind(ma3, a3, mr);
    automatic.evaluate();
    manual.evaluate();
    EXPECT_EQ(automatic.open(ay), expect);
    EXPECT_EQ(manual.open(my), expect);
    EXPECT_EQ(automatic.open(ay), manual.open(my));
}

TEST(Beaver, ElementaryLogManyPoints)
{
    const u64 xs[] = {0, 1, 2, 7};
    const u64 zs[] = {0, 1, 3, 6};
    const u64 coeffs[] = {0, 1, 4};
    for (u64 x : xs)
    for (u64 z : zs)
    for (u64 a0 : coeffs)
    for (u64 a3 : coeffs)
    {
        const u64 a1 = a0 + 2;
        const u64 a2 = a3 + 1;
        const u64 m = x * z;
        const u64 expect = a0 + a1 * m + a2 * m * m + a3 * mpow(m, 3);
        auto got = run_poly(6, [](session64 & s, const std::vector<wire64> & in) {
            return s(in[5] * pow(in[0], 3) * pow(in[1], 3)
                + in[4] * pow(in[0], 2) * pow(in[1], 2)
                + in[3] * in[0] * in[1]
                + in[2]);
        }, {x, z, a0, a1, a2, a3});
        EXPECT_EQ(got.value, expect);
        EXPECT_EQ(got.rounds, 2);
        EXPECT_EQ(got.prep, 12u);
        EXPECT_EQ(got.monos, 6u);
        EXPECT_EQ(got.wires, 8u);
        EXPECT_EQ(got.draws, 18);
    }
}

TEST(Beaver, ElementaryQuadraticAndQuarticFactor)
{
    auto quadratic = run_poly(5, [](session64 & s, const std::vector<wire64> & in) {
        auto x = in[0], z = in[1], a0 = in[2], a1 = in[3], a2 = in[4];
        return s(a2 * pow(x, 2) * pow(z, 2) + a1 * x * z + a0);
    }, {3, 4, 5, 6, 7});
    const u64 mq = 3u * 4u;
    EXPECT_EQ(quadratic.value, 5u + 6u * mq + 7u * mq * mq);
    EXPECT_EQ(quadratic.rounds, 2);
    EXPECT_EQ(quadratic.wires, 7u);

    auto quartic = run_poly(7, [](session64 & s, const std::vector<wire64> & in) {
        auto x = in[0], z = in[1];
        return s(in[6] * pow(x, 4) * pow(z, 4) + in[5] * pow(x, 3) * pow(z, 3)
            + in[4] * pow(x, 2) * pow(z, 2) + in[3] * x * z + in[2]);
    }, {2, 3, 1, 2, 3, 4, 5});
    const u64 m = 2u * 3u;
    const u64 expect = 1u + 2u * m + 3u * mpow(m, 2) + 4u * mpow(m, 3) + 5u * mpow(m, 4);
    EXPECT_EQ(quartic.value, expect);
    EXPECT_EQ(quartic.rounds, 2);
    EXPECT_EQ(quartic.wires, 9u);
    EXPECT_LT(quartic.prep, quadratic.prep + 40u);
}

TEST(Beaver, ElementaryLeadingTermKeepsAnExtraFactor)
{
    // x and y share exponents. z appears only on the leading term.
    auto got = run_poly(6, [](session64 & s, const std::vector<wire64> & in) {
        auto x = in[0], y = in[1], z = in[2], a0 = in[3], a1 = in[4], a2 = in[5];
        return s(a2 * pow(x, 2) * pow(y, 2) * z + a1 * x * y + a0);
    }, {2, 3, 4, 5, 6, 7});
    const u64 m = 2u * 3u;
    EXPECT_EQ(got.value, 5u + 6u * m + 7u * m * m * 4u);
    EXPECT_EQ(got.rounds, 2);
}

TEST(Beaver, ElementaryThreeWireCluster)
{
    auto got = run_poly(5, [](session64 & s, const std::vector<wire64> & in) {
        auto x = in[0], y = in[1], z = in[2], a0 = in[3], a1 = in[4];
        return s(u64{2} * pow(x, 2) * pow(y, 2) * pow(z, 2) + a1 * x * y * z + a0);
    }, {2, 3, 4, 5, 6});
    const u64 m = 2u * 3u * 4u;
    EXPECT_EQ(got.value, 5u + 6u * m + 2u * m * m);
    EXPECT_EQ(got.rounds, 2);
}

TEST(Beaver, MismatchedPowersPeelTheSharedVariable)
{
    // x is in every term and z is in every term, but their exponents do not
    // match, so the polynomial is not rewritten as a polynomial in x*z.
    auto got = run_poly(4, [](session64 & s, const std::vector<wire64> & in) {
        auto a = in[0], b = in[1], x = in[2], z = in[3];
        return s(a * pow(x, 2) * z + b * x * pow(z, 2));
    }, {2, 5, 3, 4});
    EXPECT_EQ(got.value, 2u * 9u * 4u + 5u * 3u * 16u);
    EXPECT_EQ(got.rounds, 2);
    EXPECT_EQ(got.wires, 6u);
}

TEST(Beaver, SecretScaleOfALowPublicPolynomialStaysOneRound)
{
    auto got = run_poly(2, [](session64 & s, const std::vector<wire64> & in) {
        return s(in[0] * (u64{1} + u64{2} * in[1] + u64{3} * pow(in[1], 2)));
    }, {4, 5});
    EXPECT_EQ(got.value, 4u * (1u + 2u * 5u + 3u * 25u));
    EXPECT_EQ(got.rounds, 1);
    EXPECT_EQ(got.wires, 3u);
}

TEST(Beaver, PublicCoefficientsStayOneRound)
{
    auto got = run_poly(1, [](session64 & s, const std::vector<wire64> & in) {
        return s.horner(in[0], {u64{1}, u64{2}, u64{3}, u64{4}, u64{5}});
    }, {3});
    const u64 x = 3;
    EXPECT_EQ(got.value, 1u + 2u * x + 3u * x * x + 4u * x * x * x + 5u * mpow(x, 4));
    EXPECT_EQ(got.rounds, 1);
    EXPECT_EQ(got.wires, 2u);

    session64 compared;
    auto manual_x = compared.input();
    auto manual = compared(u64{1} + u64{2} * manual_x + u64{3} * pow(manual_x, 2)
        + u64{4} * pow(manual_x, 3) + u64{5} * pow(manual_x, 4));
    EXPECT_EQ(compared.round_of(manual), 1);
    EXPECT_EQ(compared.preprocessing_count(), got.prep);
    EXPECT_EQ(compared.monomial_count(), got.monos);
}

TEST(Beaver, PublicSignSplitsOnceThePowerCrossesGrow)
{
    session64 quadratic;
    auto qsgn = quadratic.input();
    auto qx = quadratic.input();
    auto qy = quadratic(qsgn * (u64{1} + qx + pow(qx, 2)));
    session64 quadratic_hand;
    auto qhs = quadratic_hand.input();
    auto qhx = quadratic_hand.input();
    auto qinner = quadratic_hand(u64{1} + qhx + pow(qhx, 2));
    auto qouter = quadratic_hand(qhs * qinner);
    // Equal preprocessing. The one-round form wins the tie.
    EXPECT_EQ(quadratic.round_of(qy), 1);
    EXPECT_EQ(quadratic_hand.round_of(qouter), 2);
    EXPECT_EQ(quadratic.preprocessing_count(), quadratic_hand.preprocessing_count());

    session64 cubic;
    auto sgn = cubic.input();
    auto x = cubic.input();
    auto y = cubic(sgn * (u64{1} + x + pow(x, 2) + pow(x, 3)));
    session64 cubic_hand;
    auto ms = cubic_hand.input();
    auto mx = cubic_hand.input();
    auto inner = cubic_hand(u64{1} + mx + pow(mx, 2) + pow(mx, 3));
    auto outer = cubic_hand(ms * inner);
    EXPECT_EQ(cubic.round_of(y), 2);
    EXPECT_EQ(cubic_hand.round_of(outer), 2);
    EXPECT_EQ(cubic.preprocessing_count(), cubic_hand.preprocessing_count());

    Counter rng;
    cubic.sample(rng);
    cubic.bind(sgn, u64{3}, rng);
    cubic.bind(x, u64{4}, rng);
    cubic.evaluate();
    EXPECT_EQ(cubic.open(y), 3u * (1u + 4u + 16u + 64u));
    (void)qouter;
}

TEST(Beaver, UnivariateSecretCubicMatchesTheFactoredBudget)
{
    session64 uni;
    auto x = uni.input();
    auto a0 = uni.input();
    auto a1 = uni.input();
    auto a2 = uni.input();
    auto a3 = uni.input();
    auto cubic = uni(a3 * pow(x, 3) + a2 * pow(x, 2) + a1 * x + a0);
    EXPECT_EQ(uni.round_of(cubic), 1);
    EXPECT_EQ(uni.wire_count(), 6u);

    session64 factored;
    auto fx = factored.input();
    auto fz = factored.input();
    auto fa0 = factored.input();
    auto fa1 = factored.input();
    auto fa2 = factored.input();
    auto fa3 = factored.input();
    auto logp = factored(fa3 * pow(fx, 3) * pow(fz, 3) + fa2 * pow(fx, 2) * pow(fz, 2)
        + fa1 * fx * fz + fa0);
    // Two extra input blinds and one product share, with m standing in for x.
    EXPECT_EQ(factored.preprocessing_count(), uni.preprocessing_count() + 3u);
    EXPECT_EQ(factored.round_of(logp), 2);

    session64 scaled;
    auto sx = scaled.input();
    auto ss = scaled.input();
    auto s0 = scaled.input();
    auto s1 = scaled.input();
    auto s2 = scaled.input();
    auto s3 = scaled.input();
    auto signed_cubic = scaled(ss * (s3 * pow(sx, 3) + s2 * pow(sx, 2) + s1 * sx + s0));
    EXPECT_EQ(scaled.preprocessing_count(), uni.preprocessing_count() + 3u);
    EXPECT_EQ(scaled.round_of(signed_cubic), 2);
}

TEST(Beaver, ConstantCoefficientIsNotMasked)
{
    session64 secret_const;
    auto x = secret_const.input();
    auto sgn = secret_const.input();
    auto a0 = secret_const.input();
    auto a1 = secret_const.input();
    auto secret = secret_const(sgn * (a1 * x + a0));

    session64 public_const;
    auto px = public_const.input();
    auto ps = public_const.input();
    auto pa1 = public_const.input();
    auto pub = public_const(ps * (pa1 * px + u64{4}));

    EXPECT_EQ(secret_const.preprocessing_count(), public_const.preprocessing_count());
    EXPECT_EQ(secret_const.wire_count(), public_const.wire_count() + 1u);
    EXPECT_EQ(secret_const.round_of(secret), 2);
    EXPECT_EQ(public_const.round_of(pub), 2);
}

TEST(Beaver, SumsAndCancellationsNeedNoProducts)
{
    auto sum = run_poly(3, [](session64 & s, const std::vector<wire64> & in) {
        return s(in[0] + in[1] + in[2]);
    }, {4, 5, 6});
    EXPECT_EQ(sum.value, 15u);
    EXPECT_EQ(sum.prep, 0u);
    EXPECT_EQ(sum.monos, 0u);
    EXPECT_EQ(sum.rounds, 1);

    auto cancelled = run_poly(2, [](session64 & s, const std::vector<wire64> & in) {
        return s(in[0] * in[1] - in[1] * in[0] + u64{4});
    }, {8, 9});
    EXPECT_EQ(cancelled.value, 4u);
    EXPECT_EQ(cancelled.prep, 0u);
    EXPECT_EQ(cancelled.monos, 0u);

    auto wiped = run_poly(2, [](session64 & s, const std::vector<wire64> & in) {
        return s(in[0] * in[1] - in[0] * in[1]);
    }, {8, 9});
    EXPECT_EQ(wiped.value, 0u);
    EXPECT_EQ(wiped.prep, 0u);

    session64 empty;
    auto dummy = empty.input();
    auto zero = empty.horner(dummy, {});
    EXPECT_EQ(empty.preprocessing_count(), 0u);
    Counter rng;
    empty.sample(rng);
    EXPECT_EQ(rng.draws, 0);
    empty.bind(dummy, u64{12}, rng);
    empty.evaluate();
    EXPECT_EQ(empty.open(zero), 0u);

    auto constant = run_poly(1, [](session64 & s, const std::vector<wire64> & in) {
        return s(pow(in[0], 0) + u64{10});
    }, {99});
    EXPECT_EQ(constant.value, 11u);
    EXPECT_EQ(constant.prep, 0u);
}

TEST(Beaver, SubtractionAndNegativeCoefficients)
{
    auto got = run_poly(4, [](session64 & s, const std::vector<wire64> & in) {
        auto sgn = in[0], x = in[1], a1 = in[2], a0 = in[3];
        return s(sgn * (pow(x, 2) - a1 * x - a0));
    }, {2, 5, 4, 6});
    const u64 inner = u64{25} - u64{4} * u64{5} - u64{6};
    EXPECT_EQ(got.value, u64{2} * inner);
    EXPECT_EQ(got.rounds, 2);

    auto wrapped = run_poly(3, [](session64 & s, const std::vector<wire64> & in) {
        return s(in[0] - in[1] * in[2]);
    }, {3, 10, 4});
    EXPECT_EQ(wrapped.value, u64{3} - u64{10} * u64{4});

    auto neg = run_poly(2, [](session64 & s, const std::vector<wire64> & in) {
        return s(-(in[0] * in[1] + in[0]));
    }, {6, 7});
    EXPECT_EQ(neg.value, u64{0} - (u64{6} * u64{7} + u64{6}));
}

TEST(Beaver, PowAndMonomialSpellingsAgree)
{
    auto as_pow = run_poly(4, [](session64 & s, const std::vector<wire64> & in) {
        return s(in[2] * pow(in[0], 2) * pow(in[1], 2) + in[3] * in[0] * in[1]);
    }, {3, 5, 2, 4});
    auto as_mono = run_poly(4, [](session64 & s, const std::vector<wire64> & in) {
        return s(in[2] * in[0] * in[0] * in[1] * in[1] + in[3] * in[0] * in[1]);
    }, {3, 5, 2, 4});
    const u64 m = 3u * 5u;
    EXPECT_EQ(as_pow.value, 2u * m * m + 4u * m);
    EXPECT_EQ(as_mono.value, as_pow.value);
    EXPECT_EQ(as_mono.prep, as_pow.prep);
    EXPECT_EQ(as_mono.rounds, as_pow.rounds);
    EXPECT_EQ(as_mono.monos, as_pow.monos);
}

TEST(Beaver, FactoredProductIsVisibleAndTheSexticTermIsNot)
{
    session64 s;
    auto x = s.input();
    auto z = s.input();
    auto a0 = s.input();
    auto a1 = s.input();
    auto y = s(a1 * pow(x, 2) * pow(z, 2) + a0);
    s.pin(y);
    const auto prep = s.preprocessing_count();
    Counter rng;
    s.sample(rng);
    EXPECT_EQ(s.preprocessing_count(), prep);
    EXPECT_EQ(s.monomial({{x, 1u}, {z, 1u}}).open(),
        s.lambda(x).open() * s.lambda(z).open());
    EXPECT_THROW((void)[&] { return s.monomial({{x, 2u}, {z, 2u}}); }(),
        std::logic_error);
    s.bind(x, u64{6}, rng);
    s.bind(z, u64{7}, rng);
    s.bind(a0, u64{8}, rng);
    s.bind(a1, u64{9}, rng);
    s.evaluate();
    s.evaluate();
    EXPECT_EQ(s.open(y), 8u + 9u * 42u * 42u);
    EXPECT_EQ(s.delta(y), s.open(y) + s.lambda(y).open());
    EXPECT_EQ(s.delta(x), u64{6} + s.lambda(x).open());
    EXPECT_EQ(dpf::reconstruct(s.value(y).party0(), s.value(y).party1()), s.open(y));
    int draws = rng.draws;
    s.sample(rng);
    EXPECT_EQ(rng.draws, draws);
    EXPECT_THROW(s.bind(x, u64{1}, rng), std::logic_error);
}

TEST(Beaver, OutputBlindIsSampledOnlyWhenPinned)
{
    session64 s;
    auto x = s.input();
    auto z = s.input();
    auto y = s(x * z + u64{3});
    EXPECT_THROW(s.lambda(y), std::logic_error);
    const auto before = s.preprocessing_count();
    s.pin(y);
    EXPECT_EQ(s.preprocessing_count(), before + 1u);
    Counter rng;
    s.sample(rng);
    auto lam = s.lambda(y);
    s.bind(x, u64{4}, rng);
    s.bind(z, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y), 23u);
    EXPECT_EQ(s.delta(y), 23u + lam.open());
}

TEST(Beaver, SecondPolynomialReusesTheFactoredPair)
{
    session64 s;
    auto x = s.input();
    auto z = s.input();
    auto a0 = s.input();
    auto a1 = s.input();
    auto y0 = s(a1 * pow(x, 2) * pow(z, 2) + a0);
    const auto monos0 = s.monomial_count();
    const auto prep0 = s.preprocessing_count();
    auto b0 = s.input();
    auto b1 = s.input();
    auto y1 = s(b1 * x * z + b0);
    EXPECT_LT(s.monomial_count() - monos0, monos0);
    EXPECT_LT(s.preprocessing_count() - prep0, prep0);
    EXPECT_EQ(s.round_of(y0), 2);
    // b1*x*z + b0 is already one product plus a value share.
    EXPECT_EQ(s.round_of(y1), 1);
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{3}, rng);
    s.bind(z, u64{4}, rng);
    s.bind(a0, u64{5}, rng);
    s.bind(a1, u64{6}, rng);
    s.bind(b0, u64{7}, rng);
    s.bind(b1, u64{8}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y0), 5u + 6u * 12u * 12u);
    EXPECT_EQ(s.open(y1), 7u + 8u * 12u);
}

TEST(Beaver, FactoredOutputFeedsALaterProduct)
{
    session64 s;
    auto x = s.input();
    auto z = s.input();
    auto a0 = s.input();
    auto a1 = s.input();
    auto y = s(a1 * x * z + a0);
    Counter rng;
    s.sample(rng);
    auto blind_x = s.lambda(x);
    s.bind(x, u64{5}, rng);
    s.bind(z, u64{6}, rng);
    s.bind(a0, u64{7}, rng);
    s.bind(a1, u64{8}, rng);
    auto w = s.input();
    auto prod = s(y * w);
    EXPECT_EQ(s.round_of(y), 1);
    EXPECT_EQ(s.round_of(prod), 2);
    s.sample(rng);
    EXPECT_EQ(s.lambda(x), blind_x);
    s.bind(w, u64{9}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y), 7u + 8u * 30u);
    EXPECT_EQ(s.open(prod), s.open(y) * 9u);
    EXPECT_EQ(s.lambda(x), blind_x);
}

TEST(Beaver, ScheduledPolynomialReplaysFromASeed)
{
    using block = dpf::prg::aes128::block_type;
    block seed = simde_mm_set_epi64x(0x1234, 0x5678);
    dpf::beavers::oracle<u64> left(seed, 4);
    dpf::beavers::oracle<u64> right(seed, 4);

    auto build = [](session64 & s) {
        auto x = s.input();
        auto z = s.input();
        auto a0 = s.input();
        auto a1 = s.input();
        auto a2 = s.input();
        auto y = s(a2 * pow(x, 2) * pow(z, 2) + a1 * x * z + a0);
        return std::tuple{x, z, a0, a1, a2, y};
    };
    session64 a;
    session64 b;
    auto [ax, az, aa0, aa1, aa2, ay] = build(a);
    auto [bx, bz, ba0, ba1, ba2, by] = build(b);
    a.sample_from(left, 4);
    b.sample_from(right, 4);
    auto am = a.material_at(left, 4);
    auto bm = b.material_at(right, 4);
    EXPECT_EQ(am.bundles, bm.bundles);
    EXPECT_EQ(am.lambda, bm.lambda);
    auto other = a.material_at(left, 5);
    EXPECT_NE(other.lambda, am.lambda);

    Counter ra, rb;
    a.bind(ax, u64{2}, ra);
    a.bind(az, u64{3}, ra);
    a.bind(aa0, u64{4}, ra);
    a.bind(aa1, u64{5}, ra);
    a.bind(aa2, u64{6}, ra);
    b.bind(bx, u64{2}, rb);
    b.bind(bz, u64{3}, rb);
    b.bind(ba0, u64{4}, rb);
    b.bind(ba1, u64{5}, rb);
    b.bind(ba2, u64{6}, rb);
    a.evaluate();
    b.evaluate();
    const u64 m = 6;
    EXPECT_EQ(a.open(ay), 4u + 5u * m + 6u * m * m);
    EXPECT_EQ(b.open(by), a.open(ay));
    EXPECT_EQ(a.lambda(ax), b.lambda(bx));
}

TEST(Beaver, BindSharesFeedsAScheduledPolynomial)
{
    session64 s;
    auto x = s.input();
    auto z = s.input();
    auto c = s.input();
    auto y = s(c * (pow(x, 2) * pow(z, 2) + x * z + u64{1}));
    s.sample();
    s.bind_shares(x, u64{2}, u64{5});
    s.bind_shares(z, u64{1}, u64{2});
    s.bind_shares(c, u64{10}, u64{4});
    s.evaluate();
    const u64 xv = 7, zv = 3, cv = 14, m = xv * zv;
    EXPECT_EQ(s.open(y), cv * (m * m + m + 1u));
    EXPECT_EQ(dpf::reconstruct(s.value(x).party0(), s.value(x).party1()), xv);
}

TEST(Beaver, GrottoPolynomialsAtSeveralPoints)
{
    const u64 xs[] = {0, 1, 2, 5};
    const u64 signs[] = {0, 1, 3};
    const u64 coeffs[] = {0, 4, 7};
    for (u64 x : xs)
    for (u64 sgn : signs)
    for (u64 a0 : coeffs)
    for (u64 a1 : coeffs)
    {
        const u64 a2 = a0 + 1;
        const u64 a3 = a1 + 2;
        auto linear = run_poly(4, [](session64 & s, const std::vector<wire64> & in) {
            return s(in[1] * (in[3] * in[0] + in[2]));
        }, {x, sgn, a0, a1});
        EXPECT_EQ(linear.value, sgn * (a1 * x + a0));
        EXPECT_EQ(linear.rounds, 2);
        EXPECT_EQ(linear.prep, 6u);

        auto quad = run_poly(5, [](session64 & s, const std::vector<wire64> & in) {
            return s(in[1] * (in[4] * pow(in[0], 2) + in[3] * in[0] + in[2]));
        }, {x, sgn, a0, a1, a2});
        EXPECT_EQ(quad.value, sgn * (a2 * x * x + a1 * x + a0));
        EXPECT_EQ(quad.rounds, 2);
        EXPECT_EQ(quad.prep, 9u);

        auto cube = run_poly(6, [](session64 & s, const std::vector<wire64> & in) {
            return s(in[1] * (in[5] * pow(in[0], 3) + in[4] * pow(in[0], 2)
                + in[3] * in[0] + in[2]));
        }, {x, sgn, a0, a1, a2, a3});
        EXPECT_EQ(cube.value, sgn * (a3 * x * x * x + a2 * x * x + a1 * x + a0));
        EXPECT_EQ(cube.rounds, 2);
        EXPECT_EQ(cube.prep, 12u);
    }
}

TEST(Beaver, ExpScaleAgreesWithAHandSplit)
{
    session64 automatic;
    auto r = automatic.input();
    auto c = automatic.input();
    auto a0 = automatic.input();
    auto a1 = automatic.input();
    auto a2 = automatic.input();
    auto a3 = automatic.input();
    auto y = automatic(c * (a3 * pow(r, 3) + a2 * pow(r, 2) + a1 * r + a0));

    session64 manual;
    auto mr = manual.input();
    auto mc = manual.input();
    auto m0 = manual.input();
    auto m1 = manual.input();
    auto m2 = manual.input();
    auto m3 = manual.input();
    auto inner = manual(m3 * pow(mr, 3) + m2 * pow(mr, 2) + m1 * mr + m0);
    auto outer = manual(mc * inner);
    EXPECT_EQ(automatic.round_of(y), manual.round_of(outer));
    EXPECT_EQ(automatic.preprocessing_count(), manual.preprocessing_count());
    EXPECT_EQ(automatic.wire_count(), manual.wire_count());

    const u64 rv = 3, cv = 4, c0 = 5, c1 = 0, c2 = 2, c3 = 1;
    Counter ra, rm;
    automatic.sample(ra);
    manual.sample(rm);
    automatic.bind(r, rv, ra);
    automatic.bind(c, cv, ra);
    automatic.bind(a0, c0, ra);
    automatic.bind(a1, c1, ra);
    automatic.bind(a2, c2, ra);
    automatic.bind(a3, c3, ra);
    manual.bind(mr, rv, rm);
    manual.bind(mc, cv, rm);
    manual.bind(m0, c0, rm);
    manual.bind(m1, c1, rm);
    manual.bind(m2, c2, rm);
    manual.bind(m3, c3, rm);
    automatic.evaluate();
    manual.evaluate();
    const u64 expect = cv * (c3 * 27u + c2 * 9u + c1 * rv + c0);
    EXPECT_EQ(automatic.open(y), expect);
    EXPECT_EQ(manual.open(outer), expect);
}

TEST(Beaver, ScaleAfterAPowerAgreesWithAHandSplit)
{
    session64 automatic;
    auto x = automatic.input();
    auto z = automatic.input();
    auto sgn = automatic.input();
    auto a0 = automatic.input();
    auto a1 = automatic.input();
    auto y = automatic(sgn * (a1 * pow(x, 2) * pow(z, 2) + a0));

    session64 manual;
    auto mx = manual.input();
    auto mz = manual.input();
    auto ms = manual.input();
    auto m0 = manual.input();
    auto m1 = manual.input();
    auto mm = manual(mx * mz);
    auto inner = manual(m1 * pow(mm, 2) + m0);
    auto outer = manual(ms * inner);
    EXPECT_EQ(automatic.round_of(y), 3);
    EXPECT_EQ(manual.round_of(outer), 3);
    EXPECT_EQ(automatic.preprocessing_count(), manual.preprocessing_count());
    EXPECT_EQ(automatic.wire_count(), manual.wire_count());

    Counter ra, rm;
    automatic.sample(ra);
    manual.sample(rm);
    automatic.bind(x, u64{2}, ra);
    automatic.bind(z, u64{5}, ra);
    automatic.bind(sgn, u64{3}, ra);
    automatic.bind(a0, u64{4}, ra);
    automatic.bind(a1, u64{6}, ra);
    manual.bind(mx, u64{2}, rm);
    manual.bind(mz, u64{5}, rm);
    manual.bind(ms, u64{3}, rm);
    manual.bind(m0, u64{4}, rm);
    manual.bind(m1, u64{6}, rm);
    automatic.evaluate();
    manual.evaluate();
    EXPECT_EQ(automatic.open(y), 3u * (6u * 100u + 4u));
    EXPECT_EQ(manual.open(outer), automatic.open(y));
}

TEST(Beaver, ModintScheduledPolynomialsWrap)
{
    using M = dpf::modint<17>;
    const M points[] = {M{0}, M{1}, M{8}, M{16}};
    for (M x : points)
    for (M z : points)
    for (M a3 : {M{0}, M{5}})
    {
        dpf::beavers::session<M> s;
        auto wx = s.input();
        auto wz = s.input();
        auto wa0 = s.input();
        auto wa1 = s.input();
        auto wa2 = s.input();
        auto wa3 = s.input();
        auto y = s(wa3 * pow(wx, 3) * pow(wz, 3) + wa2 * pow(wx, 2) * pow(wz, 2)
            + wa1 * wx * wz + wa0);
        EXPECT_EQ(s.round_of(y), 2);
        EXPECT_EQ(s.preprocessing_count(), 12u);
        Seq<M> rng;
        s.sample(rng);
        const M a0{3}, a1{4}, a2{9};
        s.bind(wx, x, rng);
        s.bind(wz, z, rng);
        s.bind(wa0, a0, rng);
        s.bind(wa1, a1, rng);
        s.bind(wa2, a2, rng);
        s.bind(wa3, a3, rng);
        s.evaluate();
        const M m = x * z;
        EXPECT_EQ(s.open(y), a0 + a1 * m + a2 * m * m + a3 * m * m * m);
    }

    dpf::beavers::session<M> scaled;
    auto r = scaled.input();
    auto c = scaled.input();
    auto y = scaled(c * (M{4} + M{5} * r + M{16} * pow(r, 2)));
    EXPECT_EQ(scaled.round_of(y), 1);
    Seq<M> rng;
    scaled.sample(rng);
    scaled.bind(r, M{6}, rng);
    scaled.bind(c, M{3}, rng);
    scaled.evaluate();
    EXPECT_EQ(scaled.open(y), M{3} * (M{4} + M{5} * M{6} + M{16} * M{6} * M{6}));
}

TEST(Beaver, XorWrapperScheduledPolynomials)
{
    using W = dpf::xor_wrapper<std::uint32_t>;
    dpf::beavers::session<W> s;
    auto x = s.input();
    auto z = s.input();
    auto a = s.input();
    auto b = s.input();
    auto y = s(a * x * z + b);
    auto square = s(x * x);
    auto binomial = s((x + z) * (x + z));
    EXPECT_EQ(s.round_of(y), 1);
    EXPECT_EQ(s.round_of(square), 1);
    XorSeq<W> rng;
    s.sample(rng);
    const W xv{0b11001100u};
    const W zv{0b10101010u};
    const W av{0b11110000u};
    const W bv{0b00001111u};
    s.bind(x, xv, rng);
    s.bind(z, zv, rng);
    s.bind(a, av, rng);
    s.bind(b, bv, rng);
    s.evaluate();
    EXPECT_EQ(s.open(y), av * (xv * zv) + bv);
    EXPECT_EQ(s.open(square), xv);
    EXPECT_EQ(s.open(binomial), xv + zv);
}

TEST(Beaver, HighPowersAndTheExpansionLimit)
{
    auto sixteenth = run_poly(1, [](session64 & s, const std::vector<wire64> & in) {
        return s(pow(in[0], 16));
    }, {2});
    EXPECT_EQ(sixteenth.value, 65536u);
    EXPECT_EQ(sixteenth.rounds, 1);

    auto combined = run_poly(1, [](session64 & s, const std::vector<wire64> & in) {
        return s(pow(in[0], 10) * pow(in[0], 6));
    }, {3});
    EXPECT_EQ(combined.value, mpow(3, 16));

    session64 s;
    auto x = s.input();
    EXPECT_THROW((void)[&] { return pow(x, 17u); }(), std::invalid_argument);
    EXPECT_THROW((void)[&] { return pow(x, 10) * pow(x, 7); }(), std::invalid_argument);

    std::vector<wire64> wide;
    wide.reserve(12);
    for (int i = 0; i < 12; ++i)
        wide.push_back(s.input());
    auto expr12 = wide[0] * wide[1];
    for (int i = 2; i < 12; ++i)
        expr12 = expr12 * wide[i];
    auto all = s(expr12);
    Counter rng;
    s.sample(rng);
    for (auto w : wide)
        s.bind(w, u64{1}, rng);
    s.bind(x, u64{2}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(all), 1u);

    auto thirteenth = s.input();
    EXPECT_THROW((void)[&] { return s(expr12 * thirteenth); }(), std::invalid_argument);
}

TEST(Beaver, ScheduledPolynomialRejectsEarlyUse)
{
    session64 a;
    session64 b;
    auto x = a.input();
    auto y = b.input();
    EXPECT_THROW((void)[&] { return x + y; }(), std::invalid_argument);
    EXPECT_THROW((void)[&] { return x * y; }(), std::invalid_argument);
    auto z = a(x + pow(x, 2));
    EXPECT_THROW(a.open(z), std::logic_error);
    EXPECT_THROW(a.evaluate(), std::logic_error);
    a.sample();
    EXPECT_THROW(a.evaluate(), std::logic_error);
    EXPECT_THROW(a.delta(x), std::logic_error);
    a.bind(x, u64{3});
    a.evaluate();
    EXPECT_EQ(a.open(z), 3u + 9u);
    EXPECT_THROW(a.bind(x, u64{4}), std::logic_error);
    EXPECT_THROW(a.bind_shares(z, u64{1}, u64{1}), std::invalid_argument);
}

TEST(Beaver, RejectsBadUse)
{
    dpf::beavers::session<u64> a;
    dpf::beavers::session<u64> b;
    auto x = a.input();
    auto y = b.input();
    EXPECT_THROW((void)[&] { return a.product(x, y); }(), std::invalid_argument);
    EXPECT_THROW((void)[&] { return x * y; }(), std::invalid_argument);
    auto bit = a.bit();
    EXPECT_THROW(a.bind(bit, u64{2}), std::invalid_argument);
    auto z = a.product(x, x);
    EXPECT_THROW(a.evaluate(), std::logic_error);
    a.sample();
    EXPECT_THROW(a.evaluate(), std::logic_error);
    EXPECT_THROW((void)[&] { return a.bit_mul(x, x); }(), std::invalid_argument);
    (void)z;
}

TEST(Beaver, ProductExtremes)
{
    const std::tuple<u64, u64, u64> cases[] = {
        {0ull, 5ull, 0ull},
        {7ull, 0ull, 0ull},
        {~u64{0}, ~u64{0}, 1ull},
        {1ull, ~u64{0}, ~u64{0}},
    };
    for (const auto & [a, b, want] : cases)
    {
        session64 s;
        auto x = s.input();
        auto y = s.input();
        auto z = s(x * y);
        Counter rng;
        s.sample(rng);
        s.bind(x, a, rng);
        s.bind(y, b, rng);
        s.evaluate();
        EXPECT_EQ(s.open(z), want) << a << " * " << b;
    }
}