1 // Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
2 // All rights reserved.
3 //
4 // Redistribution and use in source and binary forms, with or without
5 // modification, are permitted provided that the following conditions are met:
6 //
7 // * Redistributions of source code must retain the above copyright
8 // notice, this list of conditions and the following disclaimer.
9 //
10 // * Redistributions in binary form must reproduce the above copyright
11 // notice, this list of conditions and the following disclaimer in the
12 // documentation and/or other materials provided with the distribution.
13 //
14 // * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
15 // its contributors may be used to endorse or promote products derived
16 // from this software without specific prior written permission.
17 //
18 // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
19 // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
20 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
21 // ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
22 // LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
23 // CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
24 // SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
25 // INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
26 // CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
27 // ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
28 // POSSIBILITY OF SUCH DAMAGE.
29 //
30 // Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)
31
32 #ifndef COLMAP_SRC_UTIL_CACHE_H_
33 #define COLMAP_SRC_UTIL_CACHE_H_
34
35 #include <iostream>
36 #include <list>
37 #include <unordered_map>
38
39 #include "util/logging.h"
40
41 namespace colmap {
42
43 // Least Recently Used cache implementation. Whenever the cache size is
44 // exceeded, the least recently used (by Get and GetMutable) is deleted.
45 template <typename key_t, typename value_t>
46 class LRUCache {
47 public:
48 LRUCache(const size_t max_num_elems,
49 const std::function<value_t(const key_t&)>& getter_func);
50 virtual ~LRUCache() = default;
51
52 // The number of elements in the cache.
53 size_t NumElems() const;
54 size_t MaxNumElems() const;
55
56 // Check whether the element with the given key exists.
57 bool Exists(const key_t& key) const;
58
59 // Get the value of an element either from the cache or compute the new value.
60 const value_t& Get(const key_t& key);
61 value_t& GetMutable(const key_t& key);
62
63 // Manually set the value of an element. Note that the ownership of the value
64 // is moved to the cache, which invalidates the object on the caller side.
65 virtual void Set(const key_t& key, value_t&& value);
66
67 // Pop least recently used element from cache.
68 virtual void Pop();
69
70 // Clear all elements from cache.
71 virtual void Clear();
72
73 protected:
74 typedef typename std::pair<key_t, value_t> key_value_pair_t;
75 typedef typename std::list<key_value_pair_t>::iterator list_iterator_t;
76
77 // Maximum number of least-recently-used elements the cache remembers.
78 const size_t max_num_elems_;
79
80 // List to keep track of the least-recently-used elements.
81 std::list<key_value_pair_t> elems_list_;
82
83 // Mapping from key to location in the list.
84 std::unordered_map<key_t, list_iterator_t> elems_map_;
85
86 // Function to compute new values if not in the cache.
87 const std::function<value_t(const key_t&)> getter_func_;
88 };
89
90 // Least Recently Used cache implementation that is constrained by a maximum
91 // memory limitation of its elements. Whenever the memory limit is exceeded, the
92 // least recently used (by Get and GetMutable) is deleted. Each element must
93 // implement a `size_t NumBytes()` method that returns its size in memory.
94 template <typename key_t, typename value_t>
95 class MemoryConstrainedLRUCache : public LRUCache<key_t, value_t> {
96 public:
97 MemoryConstrainedLRUCache(
98 const size_t max_num_bytes,
99 const std::function<value_t(const key_t&)>& getter_func);
100
101 size_t NumBytes() const;
102 size_t MaxNumBytes() const;
103 void UpdateNumBytes(const key_t& key);
104
105 void Set(const key_t& key, value_t&& value) override;
106 void Pop() override;
107 void Clear() override;
108
109 private:
110 using typename LRUCache<key_t, value_t>::key_value_pair_t;
111 using typename LRUCache<key_t, value_t>::list_iterator_t;
112 using LRUCache<key_t, value_t>::max_num_elems_;
113 using LRUCache<key_t, value_t>::elems_list_;
114 using LRUCache<key_t, value_t>::elems_map_;
115 using LRUCache<key_t, value_t>::getter_func_;
116
117 const size_t max_num_bytes_;
118 size_t num_bytes_;
119 std::unordered_map<key_t, size_t> elems_num_bytes_;
120 };
121
122 ////////////////////////////////////////////////////////////////////////////////
123 // Implementation
124 ////////////////////////////////////////////////////////////////////////////////
125
126 template <typename key_t, typename value_t>
LRUCache(const size_t max_num_elems,const std::function<value_t (const key_t &)> & getter_func)127 LRUCache<key_t, value_t>::LRUCache(
128 const size_t max_num_elems,
129 const std::function<value_t(const key_t&)>& getter_func)
130 : max_num_elems_(max_num_elems), getter_func_(getter_func) {
131 CHECK(getter_func);
132 CHECK_GT(max_num_elems, 0);
133 }
134
135 template <typename key_t, typename value_t>
NumElems()136 size_t LRUCache<key_t, value_t>::NumElems() const {
137 return elems_map_.size();
138 }
139
140 template <typename key_t, typename value_t>
MaxNumElems()141 size_t LRUCache<key_t, value_t>::MaxNumElems() const {
142 return max_num_elems_;
143 }
144
145 template <typename key_t, typename value_t>
Exists(const key_t & key)146 bool LRUCache<key_t, value_t>::Exists(const key_t& key) const {
147 return elems_map_.find(key) != elems_map_.end();
148 }
149
150 template <typename key_t, typename value_t>
Get(const key_t & key)151 const value_t& LRUCache<key_t, value_t>::Get(const key_t& key) {
152 return GetMutable(key);
153 }
154
155 template <typename key_t, typename value_t>
GetMutable(const key_t & key)156 value_t& LRUCache<key_t, value_t>::GetMutable(const key_t& key) {
157 const auto it = elems_map_.find(key);
158 if (it == elems_map_.end()) {
159 Set(key, std::move(getter_func_(key)));
160 return elems_map_[key]->second;
161 } else {
162 elems_list_.splice(elems_list_.begin(), elems_list_, it->second);
163 return it->second->second;
164 }
165 }
166
167 template <typename key_t, typename value_t>
Set(const key_t & key,value_t && value)168 void LRUCache<key_t, value_t>::Set(const key_t& key, value_t&& value) {
169 auto it = elems_map_.find(key);
170 elems_list_.push_front(key_value_pair_t(key, std::move(value)));
171 if (it != elems_map_.end()) {
172 elems_list_.erase(it->second);
173 elems_map_.erase(it);
174 }
175 elems_map_[key] = elems_list_.begin();
176 if (elems_map_.size() > max_num_elems_) {
177 Pop();
178 }
179 }
180
181 template <typename key_t, typename value_t>
Pop()182 void LRUCache<key_t, value_t>::Pop() {
183 if (!elems_list_.empty()) {
184 auto last = elems_list_.end();
185 --last;
186 elems_map_.erase(last->first);
187 elems_list_.pop_back();
188 }
189 }
190
191 template <typename key_t, typename value_t>
Clear()192 void LRUCache<key_t, value_t>::Clear() {
193 elems_list_.clear();
194 elems_map_.clear();
195 }
196
197 template <typename key_t, typename value_t>
MemoryConstrainedLRUCache(const size_t max_num_bytes,const std::function<value_t (const key_t &)> & getter_func)198 MemoryConstrainedLRUCache<key_t, value_t>::MemoryConstrainedLRUCache(
199 const size_t max_num_bytes,
200 const std::function<value_t(const key_t&)>& getter_func)
201 : LRUCache<key_t, value_t>(std::numeric_limits<size_t>::max(), getter_func),
202 max_num_bytes_(max_num_bytes),
203 num_bytes_(0) {
204 CHECK_GT(max_num_bytes, 0);
205 }
206
207 template <typename key_t, typename value_t>
NumBytes()208 size_t MemoryConstrainedLRUCache<key_t, value_t>::NumBytes() const {
209 return num_bytes_;
210 }
211
212 template <typename key_t, typename value_t>
MaxNumBytes()213 size_t MemoryConstrainedLRUCache<key_t, value_t>::MaxNumBytes() const {
214 return max_num_bytes_;
215 }
216
217 template <typename key_t, typename value_t>
Set(const key_t & key,value_t && value)218 void MemoryConstrainedLRUCache<key_t, value_t>::Set(const key_t& key,
219 value_t&& value) {
220 auto it = elems_map_.find(key);
221 elems_list_.push_front(key_value_pair_t(key, std::move(value)));
222 if (it != elems_map_.end()) {
223 elems_list_.erase(it->second);
224 elems_map_.erase(it);
225 }
226 elems_map_[key] = elems_list_.begin();
227
228 const size_t num_bytes = value.NumBytes();
229 num_bytes_ += num_bytes;
230 elems_num_bytes_.emplace(key, num_bytes);
231
232 while (num_bytes_ > max_num_bytes_ && elems_map_.size() > 1) {
233 Pop();
234 }
235 }
236
237 template <typename key_t, typename value_t>
Pop()238 void MemoryConstrainedLRUCache<key_t, value_t>::Pop() {
239 if (!elems_list_.empty()) {
240 auto last = elems_list_.end();
241 --last;
242 num_bytes_ -= elems_num_bytes_.at(last->first);
243 CHECK_GE(num_bytes_, 0);
244 elems_num_bytes_.erase(last->first);
245 elems_map_.erase(last->first);
246 elems_list_.pop_back();
247 }
248 }
249
250 template <typename key_t, typename value_t>
UpdateNumBytes(const key_t & key)251 void MemoryConstrainedLRUCache<key_t, value_t>::UpdateNumBytes(
252 const key_t& key) {
253 auto& num_bytes = elems_num_bytes_.at(key);
254 num_bytes_ -= num_bytes;
255 CHECK_GE(num_bytes_, 0);
256 num_bytes = LRUCache<key_t, value_t>::Get(key).NumBytes();
257 num_bytes_ += num_bytes;
258
259 while (num_bytes_ > max_num_bytes_ && elems_map_.size() > 1) {
260 Pop();
261 }
262 }
263
264 template <typename key_t, typename value_t>
Clear()265 void MemoryConstrainedLRUCache<key_t, value_t>::Clear() {
266 LRUCache<key_t, value_t>::Clear();
267 num_bytes_ = 0;
268 elems_num_bytes_.clear();
269 }
270
271 } // namespace colmap
272
273 #endif // COLMAP_SRC_UTIL_CACHE_H_
274