libdpf/test/tests/cohort_test.cpp
Ryan Henry 0d22946a0e Checkpoint the party/runtime stack before share-program and malicious-mode work.
Ship the TLS mesh, composer, Beaver/Yao/leaf MPC, prep/online paths, apps, and docs so the tree is pushable before elevating share_expr, security_mode, and prep resume.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-28 05:59:19 -06:00

252 lines
8.5 KiB
C++

#include <gtest/gtest.h>
#include "dpf.hpp"
#include <cstdint>
#include <vector>
namespace
{
using in_type = std::uint8_t;
using out_type = std::uint64_t;
template <typename T>
out_type raw_of(const T & v)
{
if constexpr (dpf::is_secret_share_v<std::decay_t<T>>)
return static_cast<out_type>(v.raw());
else
return static_cast<out_type>(v);
}
out_type open(out_type a, out_type b)
{
return static_cast<out_type>(a - b);
}
template <typename Key>
std::vector<out_type> interval_alone(const Key & key, in_type from, in_type to)
{
auto buf = dpf::make_output_buffer_for_interval(key, from, to);
dpf::eval_interval(key, from, to, buf);
std::vector<out_type> v;
v.reserve(buf.size());
for (std::size_t i = 0; i < buf.size(); ++i)
v.push_back(raw_of(buf[i]));
return v;
}
} // namespace
TEST(Cohort, SamePointGenAndPointEval)
{
const in_type alpha = 7;
const std::vector<out_type> beta{1, 2, 3, 4, 5};
auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
ASSERT_EQ(c0.size(), beta.size());
std::vector<out_type> y0, y1;
c0.eval_point(alpha, y0);
c1.eval_point(alpha, y1);
ASSERT_EQ(y0.size(), beta.size());
for (std::size_t k = 0; k < beta.size(); ++k)
{
EXPECT_EQ(open(y0[k], y1[k]), beta[k]);
EXPECT_EQ(y0[k], raw_of(*dpf::eval_point(c0.keys()[k], alpha)));
EXPECT_EQ(y1[k], raw_of(*dpf::eval_point(c1.keys()[k], alpha)));
}
for (std::size_t k = 0; k < beta.size(); ++k)
{
const in_type other = static_cast<in_type>(alpha + 1 + k);
std::vector<out_type> z0, z1;
c0.eval_point(other, z0);
c1.eval_point(other, z1);
EXPECT_EQ(open(z0[k], z1[k]), 0u);
EXPECT_EQ(z0[k], raw_of(*dpf::eval_point(c0.keys()[k], other)));
}
}
TEST(Cohort, IntervalInterleavesLeaves)
{
const in_type alpha = 9;
const std::vector<out_type> beta{4, 8, 15, 16};
auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
const in_type from = 0, to = 15;
std::vector<out_type> y0, y1;
c0.eval_interval(from, to, y0);
c1.eval_interval(from, to, y1);
using key0 = std::decay_t<decltype(c0.keys()[0])>;
constexpr std::size_t opl = key0::outputs_per_leaf;
const std::size_t n = beta.size();
const auto alone0 = interval_alone(c0.keys()[0], from, to);
ASSERT_EQ(y0.size(), alone0.size() * n);
ASSERT_EQ(alone0.size() % opl, 0u);
for (std::size_t k = 0; k < n; ++k)
{
const auto a0 = interval_alone(c0.keys()[k], from, to);
const auto a1 = interval_alone(c1.keys()[k], from, to);
ASSERT_EQ(a0.size(), alone0.size());
for (std::size_t i = 0; i < a0.size(); ++i)
{
const std::size_t node = i / opl;
const std::size_t lane = i % opl;
const std::size_t slot = (node * n + k) * opl + lane;
EXPECT_EQ(y0[slot], a0[i]);
EXPECT_EQ(y1[slot], a1[i]);
const in_type x = static_cast<in_type>(from + i);
const out_type opened = open(a0[i], a1[i]);
if (x == alpha)
EXPECT_EQ(opened, beta[k]);
else
EXPECT_EQ(opened, 0u);
}
}
}
TEST(Cohort, SequenceRecipeAndInnerProduct)
{
const in_type alpha = 5;
const std::vector<out_type> beta{3, 9, 1};
auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
const std::vector<in_type> points{1, 5, 5, 20, 40};
const std::size_t n = beta.size();
std::vector<out_type> y0, y1;
c0.eval_sequence(points.begin(), points.end(), y0);
c1.eval_sequence(points.begin(), points.end(), y1);
ASSERT_EQ(y0.size(), points.size() * n);
using key0 = std::decay_t<decltype(c0.keys()[0])>;
auto recipe = dpf::make_sequence_recipe<key0>(
points.begin(), points.end());
std::vector<out_type> r0, r1;
c0.eval_sequence(recipe, r0);
c1.eval_sequence(recipe, r1);
for (std::size_t q = 0; q < points.size(); ++q)
{
for (std::size_t k = 0; k < n; ++k)
{
const std::size_t slot = dpf::cohort_index(q, k, n);
const auto e0 = raw_of(*dpf::eval_point(c0.keys()[k], points[q]));
const auto e1 = raw_of(*dpf::eval_point(c1.keys()[k], points[q]));
EXPECT_EQ(y0[slot], e0);
EXPECT_EQ(y1[slot], e1);
EXPECT_EQ(r0[slot], e0);
EXPECT_EQ(r1[slot], e1);
if (points[q] == alpha)
EXPECT_EQ(open(e0, e1), beta[k]);
else
EXPECT_EQ(open(e0, e1), 0u);
}
}
std::vector<out_type> w(points.size());
for (std::size_t q = 0; q < w.size(); ++q)
w[q] = q + 3;
const auto s0 = c0.eval_sequence_inner_product(recipe, w);
const auto s1 = c1.eval_sequence_inner_product(points.begin(), points.end(), w);
ASSERT_EQ(s0.size(), n);
for (std::size_t k = 0; k < n; ++k)
{
out_type acc0 = 0, acc1 = 0;
for (std::size_t q = 0; q < points.size(); ++q)
{
acc0 = static_cast<out_type>(acc0
+ y0[dpf::cohort_index(q, k, n)] * w[q]);
acc1 = static_cast<out_type>(acc1
+ y1[dpf::cohort_index(q, k, n)] * w[q]);
}
EXPECT_EQ(s0[k], acc0);
EXPECT_EQ(s1[k], acc1);
EXPECT_EQ(open(s0[k], s1[k]), beta[k] * (w[1] + w[2]));
}
}
TEST(Cohort, IntervalInnerProduct)
{
const in_type alpha = 4;
const std::vector<out_type> beta{6, 7};
auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
const in_type from = 0, to = 15;
const auto alone = interval_alone(c0.keys()[0], from, to);
std::vector<out_type> w(alone.size());
for (std::size_t i = 0; i < w.size(); ++i)
w[i] = (i * 5u) + 1u;
const auto a0 = c0.eval_interval_inner_product(from, to, w);
const auto a1 = c1.eval_interval_inner_product(from, to, w);
for (std::size_t k = 0; k < beta.size(); ++k)
{
const auto b0 = interval_alone(c0.keys()[k], from, to);
const auto b1 = interval_alone(c1.keys()[k], from, to);
out_type e0 = 0;
out_type e1 = 0;
for (std::size_t i = 0; i < w.size(); ++i)
{
e0 = static_cast<out_type>(e0 + b0[i] * w[i]);
e1 = static_cast<out_type>(e1 + b1[i] * w[i]);
}
EXPECT_EQ(a0[k], e0);
EXPECT_EQ(a1[k], e1);
// Same weights against one-key batched leaf inner product.
auto m0 = dpf::make_basic_interval_memoizer(c0.keys()[k], from, to);
auto m1 = dpf::make_basic_interval_memoizer(c1.keys()[k], from, to);
const auto ip0 = dpf::eval_inner_product(c0.keys()[k], from, to, w, m0);
const auto ip1 = dpf::eval_inner_product(c1.keys()[k], from, to, w, m1);
EXPECT_EQ(raw_of(ip0), a0[k]);
EXPECT_EQ(raw_of(ip1), a1[k]);
EXPECT_EQ(open(a0[k], a1[k]), beta[k] * w[alpha - from]);
}
}
TEST(Cohort, InterleaveLeavesMatchesSequenceLayout)
{
// Sequence cohort layout is `out[q * n + k]`, the same order
// `interleave_leaves` writes for whole leaf values.
const in_type alpha = 6;
const std::vector<out_type> beta{2, 3, 5};
auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
(void)c1;
const std::vector<in_type> pts{1, 6, 10, 20};
const std::size_t n = beta.size();
std::vector<out_type> cohort;
c0.eval_sequence(pts.begin(), pts.end(), cohort);
std::vector<std::vector<out_type>> per_key(n);
std::vector<const out_type *> ptrs(n);
for (std::size_t k = 0; k < n; ++k)
{
per_key[k].resize(pts.size());
for (std::size_t q = 0; q < pts.size(); ++q)
per_key[k][q] = raw_of(*dpf::eval_point(c0.keys()[k], pts[q]));
ptrs[k] = per_key[k].data();
}
std::vector<out_type> interleaved(pts.size() * n);
dpf::interleave_leaves<out_type>(
interleaved.data(), ptrs.data(), n, pts.size());
ASSERT_EQ(interleaved.size(), cohort.size());
for (std::size_t i = 0; i < interleaved.size(); ++i)
EXPECT_EQ(interleaved[i], cohort[i]) << "slot " << i;
}
TEST(Cohort, HalfTreeSamePoint)
{
using ht = dpf::prg::aes128_ccr;
const in_type alpha = 3;
const std::vector<out_type> beta{11, 13, 17, 19};
auto [c0, c1] = dpf::make_dpf_cohort<ht, dpf::prg::aes128>(
alpha, beta.begin(), beta.end());
std::vector<out_type> y0, y1;
c0.eval_point(alpha, y0);
c1.eval_point(alpha, y1);
for (std::size_t k = 0; k < beta.size(); ++k)
{
EXPECT_EQ(open(y0[k], y1[k]), beta[k]);
EXPECT_EQ(y0[k], raw_of(*dpf::eval_point(c0.keys()[k], alpha)));
}
}