1/*
2 * Copyright (C) 2014 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#include "suggest/policyimpl/dictionary/utils/trie_map.h"
18
19#include <gtest/gtest.h>
20
21#include <algorithm>
22#include <cstdlib>
23#include <functional>
24#include <map>
25#include <random>
26#include <unordered_map>
27
28namespace latinime {
29namespace {
30
31TEST(TrieMapTest, TestSetAndGet) {
32    TrieMap trieMap;
33    trieMap.putRoot(10, 10);
34    EXPECT_EQ(10ull, trieMap.getRoot(10).mValue);
35    trieMap.putRoot(0x10A, 10);
36    EXPECT_EQ(10ull, trieMap.getRoot(10).mValue);
37    EXPECT_EQ(10ull, trieMap.getRoot(0x10A).mValue);
38    trieMap.putRoot(10, 1000);
39    EXPECT_EQ(1000ull, trieMap.getRoot(10).mValue);
40    trieMap.putRoot(11, 1000);
41    EXPECT_EQ(1000ull, trieMap.getRoot(11).mValue);
42    const int next = trieMap.getNextLevelBitmapEntryIndex(10);
43    trieMap.put(9, 9, next);
44    EXPECT_EQ(9ull, trieMap.get(9, next).mValue);
45    EXPECT_FALSE(trieMap.get(11, next).mIsValid);
46    trieMap.putRoot(0, 0xFFFFFFFFFull);
47    EXPECT_EQ(0xFFFFFFFFFull, trieMap.getRoot(0).mValue);
48}
49
50TEST(TrieMapTest, TestSetAndGetLarge) {
51    static const int ELEMENT_COUNT = 200000;
52    TrieMap trieMap;
53    for (int i = 0; i < ELEMENT_COUNT; ++i) {
54        EXPECT_TRUE(trieMap.putRoot(i, i));
55    }
56    for (int i = 0; i < ELEMENT_COUNT; ++i) {
57        EXPECT_EQ(static_cast<uint64_t>(i), trieMap.getRoot(i).mValue);
58    }
59}
60
61TEST(TrieMapTest, TestRandSetAndGetLarge) {
62    static const int ELEMENT_COUNT = 100000;
63    TrieMap trieMap;
64    std::unordered_map<int, uint64_t> testKeyValuePairs;
65
66    // Use the uniform integer distribution [S_INT_MIN, S_INT_MAX].
67    std::uniform_int_distribution<int> keyDistribution(S_INT_MIN, S_INT_MAX);
68    auto keyRandomNumberGenerator = std::bind(keyDistribution, std::mt19937());
69
70    // Use the uniform distribution [0, TrieMap::MAX_VALUE].
71    std::uniform_int_distribution<uint64_t> valueDistribution(0, TrieMap::MAX_VALUE);
72    auto valueRandomNumberGenerator = std::bind(valueDistribution, std::mt19937());
73
74    for (int i = 0; i < ELEMENT_COUNT; ++i) {
75        const int key = keyRandomNumberGenerator();
76        const uint64_t value = valueRandomNumberGenerator();
77        EXPECT_TRUE(trieMap.putRoot(key, value)) << key << " " << value;
78        testKeyValuePairs[key] = value;
79    }
80    for (const auto &v : testKeyValuePairs) {
81        EXPECT_EQ(v.second, trieMap.getRoot(v.first).mValue);
82    }
83}
84
85TEST(TrieMapTest, TestMultiLevel) {
86    static const int FIRST_LEVEL_ENTRY_COUNT = 10000;
87    static const int SECOND_LEVEL_ENTRY_COUNT = 20000;
88    static const int THIRD_LEVEL_ENTRY_COUNT = 40000;
89
90    TrieMap trieMap;
91    std::vector<int> firstLevelKeys;
92    std::map<int, uint64_t> firstLevelEntries;
93    std::vector<std::pair<int, int>> secondLevelKeys;
94    std::map<int, std::map<int, uint64_t>> twoLevelMap;
95    std::map<int, std::map<int, std::map<int, uint64_t>>> threeLevelMap;
96
97    // Use the uniform integer distribution [0, S_INT_MAX].
98    std::uniform_int_distribution<int> distribution(0, S_INT_MAX);
99    auto keyRandomNumberGenerator = std::bind(distribution, std::mt19937());
100    auto randomNumberGeneratorForKeySelection = std::bind(distribution, std::mt19937());
101
102    // Use the uniform distribution [0, TrieMap::MAX_VALUE].
103    std::uniform_int_distribution<uint64_t> valueDistribution(0, TrieMap::MAX_VALUE);
104    auto valueRandomNumberGenerator = std::bind(valueDistribution, std::mt19937());
105
106    for (int i = 0; i < FIRST_LEVEL_ENTRY_COUNT; ++i) {
107        const int key = keyRandomNumberGenerator();
108        const uint64_t value = valueRandomNumberGenerator();
109        EXPECT_TRUE(trieMap.putRoot(key, value));
110        firstLevelKeys.push_back(key);
111        firstLevelEntries[key] = value;
112    }
113
114    for (int i = 0; i < SECOND_LEVEL_ENTRY_COUNT; ++i) {
115        const int key = keyRandomNumberGenerator();
116        const uint64_t value = valueRandomNumberGenerator();
117        const int firstLevelKey =
118                firstLevelKeys[randomNumberGeneratorForKeySelection() % FIRST_LEVEL_ENTRY_COUNT];
119        const int nextLevelBitmapEntryIndex = trieMap.getNextLevelBitmapEntryIndex(firstLevelKey);
120        EXPECT_NE(TrieMap::INVALID_INDEX, nextLevelBitmapEntryIndex);
121        EXPECT_TRUE(trieMap.put(key, value, nextLevelBitmapEntryIndex));
122        secondLevelKeys.push_back(std::make_pair(firstLevelKey, key));
123        twoLevelMap[firstLevelKey][key] = value;
124    }
125
126    for (int i = 0; i < THIRD_LEVEL_ENTRY_COUNT; ++i) {
127        const int key = keyRandomNumberGenerator();
128        const uint64_t value = valueRandomNumberGenerator();
129        const std::pair<int, int> secondLevelKey =
130                secondLevelKeys[randomNumberGeneratorForKeySelection() % SECOND_LEVEL_ENTRY_COUNT];
131        const int secondLevel = trieMap.getNextLevelBitmapEntryIndex(secondLevelKey.first);
132        EXPECT_NE(TrieMap::INVALID_INDEX, secondLevel);
133        const int thirdLevel = trieMap.getNextLevelBitmapEntryIndex(
134                secondLevelKey.second, secondLevel);
135        EXPECT_NE(TrieMap::INVALID_INDEX, thirdLevel);
136        EXPECT_TRUE(trieMap.put(key, value, thirdLevel));
137        threeLevelMap[secondLevelKey.first][secondLevelKey.second][key] = value;
138    }
139
140    for (const auto &firstLevelEntry : firstLevelEntries) {
141        EXPECT_EQ(firstLevelEntry.second, trieMap.getRoot(firstLevelEntry.first).mValue);
142    }
143
144    for (const auto &firstLevelEntry : twoLevelMap) {
145        const int secondLevel = trieMap.getNextLevelBitmapEntryIndex(firstLevelEntry.first);
146        EXPECT_NE(TrieMap::INVALID_INDEX, secondLevel);
147        for (const auto &secondLevelEntry : firstLevelEntry.second) {
148            EXPECT_EQ(secondLevelEntry.second,
149                    trieMap.get(secondLevelEntry.first, secondLevel).mValue);
150        }
151    }
152
153    for (const auto &firstLevelEntry : threeLevelMap) {
154        const int secondLevel = trieMap.getNextLevelBitmapEntryIndex(firstLevelEntry.first);
155        EXPECT_NE(TrieMap::INVALID_INDEX, secondLevel);
156        for (const auto &secondLevelEntry : firstLevelEntry.second) {
157            const int thirdLevel =
158                    trieMap.getNextLevelBitmapEntryIndex(secondLevelEntry.first, secondLevel);
159            EXPECT_NE(TrieMap::INVALID_INDEX, thirdLevel);
160            for (const auto &thirdLevelEntry : secondLevelEntry.second) {
161                EXPECT_EQ(thirdLevelEntry.second,
162                        trieMap.get(thirdLevelEntry.first, thirdLevel).mValue);
163            }
164        }
165    }
166
167    // Iteration
168    for (const auto &firstLevelEntry : trieMap.getEntriesInRootLevel()) {
169        EXPECT_EQ(trieMap.getRoot(firstLevelEntry.key()).mValue, firstLevelEntry.value());
170        EXPECT_EQ(firstLevelEntries[firstLevelEntry.key()], firstLevelEntry.value());
171        firstLevelEntries.erase(firstLevelEntry.key());
172        for (const auto &secondLevelEntry : firstLevelEntry.getEntriesInNextLevel()) {
173            EXPECT_EQ(twoLevelMap[firstLevelEntry.key()][secondLevelEntry.key()],
174                    secondLevelEntry.value());
175            twoLevelMap[firstLevelEntry.key()].erase(secondLevelEntry.key());
176            for (const auto &thirdLevelEntry : secondLevelEntry.getEntriesInNextLevel()) {
177                EXPECT_EQ(threeLevelMap[firstLevelEntry.key()][secondLevelEntry.key()]
178                        [thirdLevelEntry.key()], thirdLevelEntry.value());
179                threeLevelMap[firstLevelEntry.key()][secondLevelEntry.key()].erase(
180                        thirdLevelEntry.key());
181            }
182        }
183    }
184
185    // Ensure all entries have been traversed.
186    EXPECT_TRUE(firstLevelEntries.empty());
187    for (const auto &secondLevelEntry : twoLevelMap) {
188        EXPECT_TRUE(secondLevelEntry.second.empty());
189    }
190    for (const auto &secondLevelEntry : threeLevelMap) {
191        for (const auto &thirdLevelEntry : secondLevelEntry.second) {
192            EXPECT_TRUE(thirdLevelEntry.second.empty());
193        }
194    }
195}
196
197TEST(TrieMapTest, TestIteration) {
198    static const int ELEMENT_COUNT = 200000;
199    TrieMap trieMap;
200    std::unordered_map<int, uint64_t> testKeyValuePairs;
201
202    // Use the uniform integer distribution [S_INT_MIN, S_INT_MAX].
203    std::uniform_int_distribution<int> keyDistribution(S_INT_MIN, S_INT_MAX);
204    auto keyRandomNumberGenerator = std::bind(keyDistribution, std::mt19937());
205
206    // Use the uniform distribution [0, TrieMap::MAX_VALUE].
207    std::uniform_int_distribution<uint64_t> valueDistribution(0, TrieMap::MAX_VALUE);
208    auto valueRandomNumberGenerator = std::bind(valueDistribution, std::mt19937());
209    for (int i = 0; i < ELEMENT_COUNT; ++i) {
210        const int key = keyRandomNumberGenerator();
211        const uint64_t value = valueRandomNumberGenerator();
212        EXPECT_TRUE(trieMap.putRoot(key, value));
213        testKeyValuePairs[key] = value;
214    }
215    for (const auto &entry : trieMap.getEntriesInRootLevel()) {
216        EXPECT_EQ(trieMap.getRoot(entry.key()).mValue, entry.value());
217        EXPECT_EQ(testKeyValuePairs[entry.key()], entry.value());
218        testKeyValuePairs.erase(entry.key());
219    }
220    EXPECT_TRUE(testKeyValuePairs.empty());
221}
222
223}  // namespace
224}  // namespace latinime
225