1 #include <atomic>
2 #include <thread>
3 #include <vector>
4
5 #include <c10/util/Lazy.h>
6 #include <gtest/gtest.h>
7
8 namespace c10_test {
9
10 // Long enough not to fit in typical SSO.
11 const std::string kLongString = "I am a long enough string";
12
TEST(LazyTest,OptimisticLazy)13 TEST(LazyTest, OptimisticLazy) {
14 std::atomic<size_t> invocations = 0;
15 auto factory = [&] {
16 ++invocations;
17 return kLongString;
18 };
19
20 c10::OptimisticLazy<std::string> s;
21
22 constexpr size_t kNumThreads = 16;
23 std::vector<std::thread> threads;
24 std::atomic<std::string*> address = nullptr;
25
26 for (size_t i = 0; i < kNumThreads; ++i) {
27 threads.emplace_back([&] {
28 auto* p = &s.ensure(factory);
29 auto old = address.exchange(p);
30 if (old != nullptr) {
31 // Even racing ensure()s should return a stable reference.
32 EXPECT_EQ(old, p);
33 }
34 });
35 }
36
37 for (auto& t : threads) {
38 t.join();
39 }
40
41 EXPECT_GE(invocations.load(), 1);
42 EXPECT_EQ(*address.load(), kLongString);
43
44 invocations = 0;
45 s.reset();
46 s.ensure(factory);
47 EXPECT_EQ(invocations.load(), 1);
48
49 invocations = 0;
50
51 auto sCopy = s;
52 EXPECT_EQ(sCopy.ensure(factory), kLongString);
53 EXPECT_EQ(invocations.load(), 0);
54
55 auto sMove = std::move(s);
56 EXPECT_EQ(sMove.ensure(factory), kLongString);
57 EXPECT_EQ(invocations.load(), 0);
58 // NOLINTNEXTLINE(bugprone-use-after-move)
59 EXPECT_EQ(s.ensure(factory), kLongString);
60 EXPECT_EQ(invocations.load(), 1);
61
62 invocations = 0;
63
64 s = sCopy;
65 EXPECT_EQ(s.ensure(factory), kLongString);
66 EXPECT_EQ(invocations.load(), 0);
67
68 s = std::move(sCopy);
69 EXPECT_EQ(s.ensure(factory), kLongString);
70 EXPECT_EQ(invocations.load(), 0);
71 }
72
TEST(LazyTest,PrecomputedLazyValue)73 TEST(LazyTest, PrecomputedLazyValue) {
74 static const std::string kLongString = "I am a string";
75 EXPECT_EQ(
76 std::make_shared<c10::PrecomputedLazyValue<std::string>>(kLongString)
77 ->get(),
78 kLongString);
79 }
80
TEST(LazyTest,OptimisticLazyValue)81 TEST(LazyTest, OptimisticLazyValue) {
82 static const std::string kLongString = "I am a string";
83
84 class LazyString : public c10::OptimisticLazyValue<std::string> {
85 std::string compute() const override {
86 return kLongString;
87 }
88 };
89
90 auto ls = std::make_shared<LazyString>();
91 EXPECT_EQ(ls->get(), kLongString);
92
93 // Returned reference should be stable.
94 EXPECT_EQ(&ls->get(), &ls->get());
95 }
96
97 } // namespace c10_test
98