1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3Licensed under the Apache License, Version 2.0 (the "License"); 4you may not use this file except in compliance with the License. 5You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9Unless required by applicable law or agreed to in writing, software 10distributed under the License is distributed on an "AS IS" BASIS, 11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12See the License for the specific language governing permissions and 13limitations under the License. 14==============================================================================*/ 15 16#ifndef TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ 17#define TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ 18 19#include <functional> 20#include <sstream> 21 22#include "tensorflow/compiler/xla/statusor.h" 23#include "tensorflow/compiler/xla/util.h" 24#include "tensorflow/core/platform/logging.h" 25 26namespace xla { 27 28// FindOrDie returns a const reference to the value associated with 29// the given key if it exists. Crashes otherwise. 30// 31// This is intended as a replacement for operator[] as an rvalue (for reading) 32// when the key is guaranteed to exist. 33template <class Collection> 34const typename Collection::value_type::second_type& FindOrDie( 35 const Collection& collection, 36 const typename Collection::value_type::first_type& key) { 37 typename Collection::const_iterator it = collection.find(key); 38 CHECK(it != collection.end()) << "Map key not found: " << key; 39 return it->second; 40} 41 42// Same as above, but returns a non-const reference. 43template <class Collection> 44typename Collection::value_type::second_type& FindOrDie( 45 Collection& collection, // NOLINT 46 const typename Collection::value_type::first_type& key) { 47 typename Collection::iterator it = collection.find(key); 48 CHECK(it != collection.end()) << "Map key not found: " << key; 49 return it->second; 50} 51 52// Like FindOrDie but returns an error instead of dying if `key` is not in 53// `container`. 54template <class Collection> 55StatusOr< 56 std::reference_wrapper<const typename Collection::value_type::second_type>> 57MaybeFind(const Collection& collection, 58 const typename Collection::value_type::first_type& key) { 59 typename Collection::const_iterator it = collection.find(key); 60 if (it == collection.end()) { 61 std::ostringstream os; 62 os << key; 63 return NotFound("key not found: %s", os.str().c_str()); 64 } 65 return {it->second}; 66} 67 68// Returns a const reference to the value associated with the given key if it 69// exists, otherwise returns a const reference to the provided default value. 70// 71// WARNING: If a temporary object is passed as the default "value," 72// this function will return a reference to that temporary object, 73// which will be destroyed at the end of the statement. A common 74// example: if you have a map with string values, and you pass a char* 75// as the default "value," either use the returned value immediately 76// or store it in a string (not string&). 77template <class Collection> 78const typename Collection::value_type::second_type& FindOrDefault( 79 const Collection& collection, 80 const typename Collection::value_type::first_type& key, 81 const typename Collection::value_type::second_type& value) { 82 auto it = collection.find(key); 83 if (it != collection.end()) return it->second; 84 return value; 85} 86 87// Inserts the key-value pair into the collection. Dies if key was already 88// present. 89template <class Collection> 90void InsertOrDie(Collection* const collection, 91 const typename Collection::value_type::first_type& key, 92 const typename Collection::value_type::second_type& data) { 93 auto p = collection->insert(std::make_pair(key, data)); 94 CHECK(p.second) << "duplicate key: " << key; 95} 96 97// Returns true if and only if the given collection contains the given key. 98template <class Collection, class Key> 99bool ContainsKey(const Collection& collection, const Key& key) { 100 return collection.find(key) != collection.end(); 101} 102 103// Inserts `value` into `set`. Dies if it was already present. 104template <class Set> 105void InsertOrDie(Set* const set, const typename Set::value_type& value) { 106 CHECK(set->insert(value).second) << "duplicate value: " << value; 107} 108 109} // namespace xla 110 111#endif // TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ 112