Use flat_hash_map in ResourceCache

This is the recommended hashmap in most cases.

PiperOrigin-RevId: 488772031
This commit is contained in:
Camillo Lugaresi 2022-11-15 15:16:11 -08:00 committed by Copybara-Service
parent 38b636f7ee
commit a67069156e
2 changed files with 8 additions and 6 deletions

View File

@ -228,6 +228,7 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework/port:logging",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/functional:function_ref",
],
)

View File

@ -17,6 +17,7 @@
#include <unordered_map>
#include "absl/container/flat_hash_map.h"
#include "absl/functional/function_ref.h"
#include "mediapipe/framework/port/logging.h"
@ -26,7 +27,8 @@ namespace mediapipe {
// resource (e.g., image dimension for an image pool) is described bye the `Key`
// type. The `Value` type must include an unset value, with implicit conversion
// to bool reflecting set/unset state.
template <typename Key, typename Value, typename KeyHash>
template <typename Key, typename Value,
typename KeyHash = typename absl::flat_hash_map<Key, int>::hasher>
class ResourceCache {
public:
Value Lookup(
@ -36,15 +38,14 @@ class ResourceCache {
Entry* entry;
if (map_it == map_.end()) {
std::tie(map_it, std::ignore) =
map_.emplace(std::piecewise_construct, std::forward_as_tuple(key),
std::forward_as_tuple(key));
entry = &map_it->second;
map_.try_emplace(key, std::make_unique<Entry>(key));
entry = map_it->second.get();
CHECK_EQ(entry->request_count, 0);
entry->request_count = 1;
entry_list_.Append(entry);
if (entry->prev != nullptr) CHECK_GE(entry->prev->request_count, 1);
} else {
entry = &map_it->second;
entry = map_it->second.get();
++entry->request_count;
Entry* larger = entry->prev;
while (larger != nullptr &&
@ -171,7 +172,7 @@ class ResourceCache {
size_t size_ = 0;
};
std::unordered_map<Key, Entry, KeyHash> map_;
absl::flat_hash_map<Key, std::unique_ptr<Entry>, KeyHash> map_;
EntryList entry_list_;
int total_request_count_ = 0;
};