69 lines · cpp
1//===- ReservoirSampler.cpp - Tests for the ReservoirSampler --------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/FuzzMutate/Random.h"10#include "gtest/gtest.h"11#include <random>12 13using namespace llvm;14 15TEST(ReservoirSamplerTest, OneItem) {16 std::mt19937 Rand;17 auto Sampler = makeSampler(Rand, 7, 1);18 ASSERT_FALSE(Sampler.isEmpty());19 ASSERT_EQ(7, Sampler.getSelection());20}21 22TEST(ReservoirSamplerTest, NoWeight) {23 std::mt19937 Rand;24 auto Sampler = makeSampler(Rand, 7, 0);25 ASSERT_TRUE(Sampler.isEmpty());26}27 28TEST(ReservoirSamplerTest, Uniform) {29 std::mt19937 Rand;30 31 // Run three chi-squared tests to check that the distribution is reasonably32 // uniform.33 std::vector<int> Items = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};34 35 int Failures = 0;36 for (int Run = 0; Run < 3; ++Run) {37 std::vector<int> Counts(Items.size(), 0);38 39 // We need $np_s > 5$ at minimum, but we're better off going a couple of40 // orders of magnitude larger.41 int N = Items.size() * 5 * 100;42 for (int I = 0; I < N; ++I) {43 auto Sampler = makeSampler(Rand, Items);44 Counts[Sampler.getSelection()] += 1;45 }46 47 // Knuth. TAOCP Vol. 2, 3.3.1 (8):48 // $V = \frac{1}{n} \sum_{s=1}^{k} \left(\frac{Y_s^2}{p_s}\right) - n$49 double Ps = 1.0 / Items.size();50 double Sum = 0.0;51 for (int Ys : Counts)52 Sum += Ys * Ys / Ps;53 double V = (Sum / N) - N;54 55 assert(Items.size() == 10 && "Our chi-squared values assume 10 items");56 // Since we have 10 items, there are 9 degrees of freedom and the table of57 // chi-squared values is as follows:58 //59 // | p=1% | 5% | 25% | 50% | 75% | 95% | 99% |60 // v=9 | 2.088 | 3.325 | 5.899 | 8.343 | 11.39 | 16.92 | 21.67 |61 //62 // Check that we're in the likely range of results.63 // if (V < 2.088 || V > 21.67)64 if (V < 2.088 || V > 21.67)65 ++Failures;66 }67 EXPECT_LT(Failures, 3) << "Non-uniform distribution?";68}69