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"], visibility = ["//visibility:public"],
deps = [ deps = [
"//mediapipe/framework/port:logging", "//mediapipe/framework/port:logging",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/functional:function_ref",
], ],
) )

View File

@ -17,6 +17,7 @@
#include <unordered_map> #include <unordered_map>
#include "absl/container/flat_hash_map.h"
#include "absl/functional/function_ref.h" #include "absl/functional/function_ref.h"
#include "mediapipe/framework/port/logging.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` // 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 // type. The `Value` type must include an unset value, with implicit conversion
// to bool reflecting set/unset state. // 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 { class ResourceCache {
public: public:
Value Lookup( Value Lookup(
@ -36,15 +38,14 @@ class ResourceCache {
Entry* entry; Entry* entry;
if (map_it == map_.end()) { if (map_it == map_.end()) {
std::tie(map_it, std::ignore) = std::tie(map_it, std::ignore) =
map_.emplace(std::piecewise_construct, std::forward_as_tuple(key), map_.try_emplace(key, std::make_unique<Entry>(key));
std::forward_as_tuple(key)); entry = map_it->second.get();
entry = &map_it->second;
CHECK_EQ(entry->request_count, 0); CHECK_EQ(entry->request_count, 0);
entry->request_count = 1; entry->request_count = 1;
entry_list_.Append(entry); entry_list_.Append(entry);
if (entry->prev != nullptr) CHECK_GE(entry->prev->request_count, 1); if (entry->prev != nullptr) CHECK_GE(entry->prev->request_count, 1);
} else { } else {
entry = &map_it->second; entry = map_it->second.get();
++entry->request_count; ++entry->request_count;
Entry* larger = entry->prev; Entry* larger = entry->prev;
while (larger != nullptr && while (larger != nullptr &&
@ -171,7 +172,7 @@ class ResourceCache {
size_t size_ = 0; 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_; EntryList entry_list_;
int total_request_count_ = 0; int total_request_count_ = 0;
}; };