xref: /aosp_15_r20/external/pytorch/c10/test/util/Bitset_test.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <gtest/gtest.h>
2 
3 #include <c10/util/Bitset.h>
4 #include <c10/util/irange.h>
5 
6 using c10::utils::bitset;
7 
TEST(BitsetTest,givenEmptyBitset_whenGettingBit_thenIsZero)8 TEST(BitsetTest, givenEmptyBitset_whenGettingBit_thenIsZero) {
9   bitset b;
10   for (size_t i = 0; i < bitset::NUM_BITS(); ++i) {
11     EXPECT_FALSE(b.get(i));
12   }
13 }
14 
TEST(BitsetTest,givenEmptyBitset_whenUnsettingBit_thenIsZero)15 TEST(BitsetTest, givenEmptyBitset_whenUnsettingBit_thenIsZero) {
16   bitset b;
17   b.unset(4);
18   for (size_t i = 0; i < bitset::NUM_BITS(); ++i) {
19     EXPECT_FALSE(b.get(i));
20   }
21 }
22 
TEST(BitsetTest,givenEmptyBitset_whenSettingAndUnsettingBit_thenIsZero)23 TEST(BitsetTest, givenEmptyBitset_whenSettingAndUnsettingBit_thenIsZero) {
24   bitset b;
25   b.set(4);
26   b.unset(4);
27   for (size_t i = 0; i < bitset::NUM_BITS(); ++i) {
28     EXPECT_FALSE(b.get(i));
29   }
30 }
31 
TEST(BitsetTest,givenEmptyBitset_whenSettingBit_thenIsSet)32 TEST(BitsetTest, givenEmptyBitset_whenSettingBit_thenIsSet) {
33   bitset b;
34   b.set(6);
35   EXPECT_TRUE(b.get(6));
36 }
37 
TEST(BitsetTest,givenEmptyBitset_whenSettingBit_thenOthersStayUnset)38 TEST(BitsetTest, givenEmptyBitset_whenSettingBit_thenOthersStayUnset) {
39   bitset b;
40   b.set(6);
41   for (const auto i : c10::irange(6)) {
42     EXPECT_FALSE(b.get(i));
43   }
44   for (size_t i = 7; i < bitset::NUM_BITS(); ++i) {
45     EXPECT_FALSE(b.get(i));
46   }
47 }
48 
TEST(BitsetTest,givenNonemptyBitset_whenSettingBit_thenIsSet)49 TEST(BitsetTest, givenNonemptyBitset_whenSettingBit_thenIsSet) {
50   bitset b;
51   b.set(6);
52   b.set(30);
53   EXPECT_TRUE(b.get(30));
54 }
55 
TEST(BitsetTest,givenNonemptyBitset_whenSettingBit_thenOthersStayAtOldValue)56 TEST(BitsetTest, givenNonemptyBitset_whenSettingBit_thenOthersStayAtOldValue) {
57   bitset b;
58   b.set(6);
59   b.set(30);
60   for (const auto i : c10::irange(6)) {
61     EXPECT_FALSE(b.get(i));
62   }
63   for (const auto i : c10::irange(7, 30)) {
64     EXPECT_FALSE(b.get(i));
65   }
66   for (size_t i = 31; i < bitset::NUM_BITS(); ++i) {
67     EXPECT_FALSE(b.get(i));
68   }
69 }
70 
TEST(BitsetTest,givenNonemptyBitset_whenUnsettingBit_thenIsUnset)71 TEST(BitsetTest, givenNonemptyBitset_whenUnsettingBit_thenIsUnset) {
72   bitset b;
73   b.set(6);
74   b.set(30);
75   b.unset(6);
76   EXPECT_FALSE(b.get(6));
77 }
78 
TEST(BitsetTest,givenNonemptyBitset_whenUnsettingBit_thenOthersStayAtOldValue)79 TEST(
80     BitsetTest,
81     givenNonemptyBitset_whenUnsettingBit_thenOthersStayAtOldValue) {
82   bitset b;
83   b.set(6);
84   b.set(30);
85   b.unset(6);
86   for (const auto i : c10::irange(30)) {
87     EXPECT_FALSE(b.get(i));
88   }
89   EXPECT_TRUE(b.get(30));
90   for (size_t i = 31; i < bitset::NUM_BITS(); ++i) {
91     EXPECT_FALSE(b.get(i));
92   }
93 }
94 
95 struct IndexCallbackMock final {
96   std::vector<size_t> called_for_indices;
97 
operator ()IndexCallbackMock98   void operator()(size_t index) {
99     called_for_indices.push_back(index);
100   }
101 
expect_was_called_for_indicesIndexCallbackMock102   void expect_was_called_for_indices(std::vector<size_t> expected_indices) {
103     EXPECT_EQ(expected_indices.size(), called_for_indices.size());
104     for (const auto i : c10::irange(expected_indices.size())) {
105       EXPECT_EQ(expected_indices[i], called_for_indices[i]);
106     }
107   }
108 };
109 
TEST(BitsetTest,givenEmptyBitset_whenCallingForEachBit_thenDoesntCall)110 TEST(BitsetTest, givenEmptyBitset_whenCallingForEachBit_thenDoesntCall) {
111   IndexCallbackMock callback;
112   bitset b;
113   b.for_each_set_bit(callback);
114   callback.expect_was_called_for_indices({});
115 }
116 
TEST(BitsetTest,givenBitsetWithOneBitSet_whenCallingForEachBit_thenCallsForEachBit)117 TEST(
118     BitsetTest,
119     givenBitsetWithOneBitSet_whenCallingForEachBit_thenCallsForEachBit) {
120   IndexCallbackMock callback;
121   bitset b;
122   b.set(5);
123   b.for_each_set_bit(callback);
124   callback.expect_was_called_for_indices({5});
125 }
126 
TEST(BitsetTest,givenBitsetWithMultipleBitsSet_whenCallingForEachBit_thenCallsForEachBit)127 TEST(
128     BitsetTest,
129     givenBitsetWithMultipleBitsSet_whenCallingForEachBit_thenCallsForEachBit) {
130   IndexCallbackMock callback;
131   bitset b;
132   b.set(5);
133   b.set(2);
134   b.set(25);
135   b.set(32);
136   b.set(50);
137   b.set(0);
138   b.unset(25);
139   b.set(10);
140   b.for_each_set_bit(callback);
141   callback.expect_was_called_for_indices({0, 2, 5, 10, 32, 50});
142 }
143