1 // Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
2 // This source code is licensed under both the GPLv2 (found in the
3 // COPYING file in the root directory) and Apache 2.0 License
4 // (found in the LICENSE.Apache file in the root directory).
5
6 #include <gtest/gtest.h>
7
8 #include <climits>
9
10 #include <queue>
11 #include <random>
12 #include <utility>
13
14 #include "util/heap.h"
15
16 #ifndef GFLAGS
17 const int64_t FLAGS_iters = 100000;
18 #else
19 #include "util/gflags_compat.h"
20 DEFINE_int64(iters, 100000, "number of pseudo-random operations in each test");
21 #endif // GFLAGS
22
23 /*
24 * Compares the custom heap implementation in util/heap.h against
25 * std::priority_queue on a pseudo-random sequence of operations.
26 */
27
28 namespace ROCKSDB_NAMESPACE {
29
30 using HeapTestValue = uint64_t;
31 using Params = std::tuple<size_t, HeapTestValue, int64_t>;
32
33 class HeapTest : public ::testing::TestWithParam<Params> {
34 };
35
TEST_P(HeapTest,Test)36 TEST_P(HeapTest, Test) {
37 // This test performs the same pseudorandom sequence of operations on a
38 // BinaryHeap and an std::priority_queue, comparing output. The three
39 // possible operations are insert, replace top and pop.
40 //
41 // Insert is chosen slightly more often than the others so that the size of
42 // the heap slowly grows. Once the size heats the MAX_HEAP_SIZE limit, we
43 // disallow inserting until the heap becomes empty, testing the "draining"
44 // scenario.
45
46 const auto MAX_HEAP_SIZE = std::get<0>(GetParam());
47 const auto MAX_VALUE = std::get<1>(GetParam());
48 const auto RNG_SEED = std::get<2>(GetParam());
49
50 BinaryHeap<HeapTestValue> heap;
51 std::priority_queue<HeapTestValue> ref;
52
53 std::mt19937 rng(static_cast<unsigned int>(RNG_SEED));
54 std::uniform_int_distribution<HeapTestValue> value_dist(0, MAX_VALUE);
55 int ndrains = 0;
56 bool draining = false; // hit max size, draining until we empty the heap
57 size_t size = 0;
58 for (int64_t i = 0; i < FLAGS_iters; ++i) {
59 if (size == 0) {
60 draining = false;
61 }
62
63 if (!draining &&
64 (size == 0 || std::bernoulli_distribution(0.4)(rng))) {
65 // insert
66 HeapTestValue val = value_dist(rng);
67 heap.push(val);
68 ref.push(val);
69 ++size;
70 if (size == MAX_HEAP_SIZE) {
71 draining = true;
72 ++ndrains;
73 }
74 } else if (std::bernoulli_distribution(0.5)(rng)) {
75 // replace top
76 HeapTestValue val = value_dist(rng);
77 heap.replace_top(val);
78 ref.pop();
79 ref.push(val);
80 } else {
81 // pop
82 assert(size > 0);
83 heap.pop();
84 ref.pop();
85 --size;
86 }
87
88 // After every operation, check that the public methods give the same
89 // results
90 assert((size == 0) == ref.empty());
91 ASSERT_EQ(size == 0, heap.empty());
92 if (size > 0) {
93 ASSERT_EQ(ref.top(), heap.top());
94 }
95 }
96
97 // Probabilities should be set up to occasionally hit the max heap size and
98 // drain it
99 assert(ndrains > 0);
100
101 heap.clear();
102 ASSERT_TRUE(heap.empty());
103 }
104
105 // Basic test, MAX_VALUE = 3*MAX_HEAP_SIZE (occasional duplicates)
106 INSTANTIATE_TEST_CASE_P(
107 Basic, HeapTest,
108 ::testing::Values(Params(1000, 3000, 0x1b575cf05b708945))
109 );
110 // Mid-size heap with small values (many duplicates)
111 INSTANTIATE_TEST_CASE_P(
112 SmallValues, HeapTest,
113 ::testing::Values(Params(100, 10, 0x5ae213f7bd5dccd0))
114 );
115 // Small heap, large value range (no duplicates)
116 INSTANTIATE_TEST_CASE_P(
117 SmallHeap, HeapTest,
118 ::testing::Values(Params(10, ULLONG_MAX, 0x3e1fa8f4d01707cf))
119 );
120 // Two-element heap
121 INSTANTIATE_TEST_CASE_P(
122 TwoElementHeap, HeapTest,
123 ::testing::Values(Params(2, 5, 0x4b5e13ea988c6abc))
124 );
125 // One-element heap
126 INSTANTIATE_TEST_CASE_P(
127 OneElementHeap, HeapTest,
128 ::testing::Values(Params(1, 3, 0x176a1019ab0b612e))
129 );
130
131 } // namespace ROCKSDB_NAMESPACE
132
main(int argc,char ** argv)133 int main(int argc, char** argv) {
134 ::testing::InitGoogleTest(&argc, argv);
135 #ifdef GFLAGS
136 GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, true);
137 #endif // GFLAGS
138 return RUN_ALL_TESTS();
139 }
140