diff options
Diffstat (limited to 'tensorflow/compiler/xla/map_util.h')
-rw-r--r-- | tensorflow/compiler/xla/map_util.h | 65 |
1 files changed, 65 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/map_util.h b/tensorflow/compiler/xla/map_util.h new file mode 100644 index 0000000000..51d0d5f86f --- /dev/null +++ b/tensorflow/compiler/xla/map_util.h @@ -0,0 +1,65 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ + +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +// FindOrDie returns a const reference to the value associated with +// the given key if it exists. Crashes otherwise. +// +// This is intended as a replacement for operator[] as an rvalue (for reading) +// when the key is guaranteed to exist. +template <class Collection> +const typename Collection::value_type::second_type& FindOrDie( + const Collection& collection, + const typename Collection::value_type::first_type& key) { + typename Collection::const_iterator it = collection.find(key); + CHECK(it != collection.end()) << "Map key not found: " << key; + return it->second; +} + +// Same as above, but returns a non-const reference. +template <class Collection> +typename Collection::value_type::second_type& FindOrDie( + Collection& collection, // NOLINT + const typename Collection::value_type::first_type& key) { + typename Collection::iterator it = collection.find(key); + CHECK(it != collection.end()) << "Map key not found: " << key; + return it->second; +} + +// Inserts the key-value pair into the collection. Dies if key was already +// present. +template <class Collection> +void InsertOrDie(Collection* const collection, + const typename Collection::value_type::first_type& key, + const typename Collection::value_type::second_type& data) { + auto p = collection->insert(std::make_pair(key, data)); + CHECK(p.second) << "duplicate key: " << key; +} + +// Returns true if and only if the given collection contains the given key. +template <class Collection, class Key> +bool ContainsKey(const Collection& collection, const Key& key) { + return collection.find(key) != collection.end(); +} + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ |