1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 // This file provides utility functions for use with STL map-like data
17 // structures, such as std::map and hash_map. Some functions will also work with
18 // sets, such as ContainsKey().
19 
20 #ifndef TENSORFLOW_LIB_GTL_MAP_UTIL_H_
21 #define TENSORFLOW_LIB_GTL_MAP_UTIL_H_
22 
23 #include <stddef.h>
24 #include <iterator>
25 #include <memory>
26 #include <string>
27 #include <utility>
28 
29 namespace tensorflow {
30 namespace gtl {
31 
32 // Returns a pointer to the const value associated with the given key if it
33 // exists, or NULL otherwise.
34 template <class Collection>
FindOrNull(const Collection & collection,const typename Collection::value_type::first_type & key)35 const typename Collection::value_type::second_type* FindOrNull(
36     const Collection& collection,
37     const typename Collection::value_type::first_type& key) {
38   typename Collection::const_iterator it = collection.find(key);
39   if (it == collection.end()) {
40     return 0;
41   }
42   return &it->second;
43 }
44 
45 // Same as above but returns a pointer to the non-const value.
46 template <class Collection>
FindOrNull(Collection & collection,const typename Collection::value_type::first_type & key)47 typename Collection::value_type::second_type* FindOrNull(
48     Collection& collection,  // NOLINT
49     const typename Collection::value_type::first_type& key) {
50   typename Collection::iterator it = collection.find(key);
51   if (it == collection.end()) {
52     return 0;
53   }
54   return &it->second;
55 }
56 
57 // Returns the pointer value associated with the given key. If none is found,
58 // NULL is returned. The function is designed to be used with a map of keys to
59 // pointers.
60 //
61 // This function does not distinguish between a missing key and a key mapped
62 // to a NULL value.
63 template <class Collection>
FindPtrOrNull(const Collection & collection,const typename Collection::value_type::first_type & key)64 typename Collection::value_type::second_type FindPtrOrNull(
65     const Collection& collection,
66     const typename Collection::value_type::first_type& key) {
67   typename Collection::const_iterator it = collection.find(key);
68   if (it == collection.end()) {
69     return typename Collection::value_type::second_type();
70   }
71   return it->second;
72 }
73 
74 // Returns a const reference to the value associated with the given key if it
75 // exists, otherwise returns a const reference to the provided default value.
76 //
77 // WARNING: If a temporary object is passed as the default "value,"
78 // this function will return a reference to that temporary object,
79 // which will be destroyed at the end of the statement. A common
80 // example: if you have a map with string values, and you pass a char*
81 // as the default "value," either use the returned value immediately
82 // or store it in a string (not string&).
83 template <class Collection>
FindWithDefault(const Collection & collection,const typename Collection::value_type::first_type & key,const typename Collection::value_type::second_type & value)84 const typename Collection::value_type::second_type& FindWithDefault(
85     const Collection& collection,
86     const typename Collection::value_type::first_type& key,
87     const typename Collection::value_type::second_type& value) {
88   typename Collection::const_iterator it = collection.find(key);
89   if (it == collection.end()) {
90     return value;
91   }
92   return it->second;
93 }
94 
95 // Inserts the given key-value pair into the collection. Returns true if and
96 // only if the key from the given pair didn't previously exist. Otherwise, the
97 // value in the map is replaced with the value from the given pair.
98 template <class Collection>
InsertOrUpdate(Collection * const collection,const typename Collection::value_type & vt)99 bool InsertOrUpdate(Collection* const collection,
100                     const typename Collection::value_type& vt) {
101   std::pair<typename Collection::iterator, bool> ret = collection->insert(vt);
102   if (!ret.second) {
103     // update
104     ret.first->second = vt.second;
105     return false;
106   }
107   return true;
108 }
109 
110 // Same as above, except that the key and value are passed separately.
111 template <class Collection>
InsertOrUpdate(Collection * const collection,const typename Collection::value_type::first_type & key,const typename Collection::value_type::second_type & value)112 bool InsertOrUpdate(Collection* const collection,
113                     const typename Collection::value_type::first_type& key,
114                     const typename Collection::value_type::second_type& value) {
115   return InsertOrUpdate(collection,
116                         typename Collection::value_type(key, value));
117 }
118 
119 // Inserts the given key and value into the given collection if and only if the
120 // given key did NOT already exist in the collection. If the key previously
121 // existed in the collection, the value is not changed. Returns true if the
122 // key-value pair was inserted; returns false if the key was already present.
123 template <class Collection>
InsertIfNotPresent(Collection * const collection,const typename Collection::value_type & vt)124 bool InsertIfNotPresent(Collection* const collection,
125                         const typename Collection::value_type& vt) {
126   return collection->insert(vt).second;
127 }
128 
129 // Same as above except the key and value are passed separately.
130 template <class Collection>
InsertIfNotPresent(Collection * const collection,const typename Collection::value_type::first_type & key,const typename Collection::value_type::second_type & value)131 bool InsertIfNotPresent(
132     Collection* const collection,
133     const typename Collection::value_type::first_type& key,
134     const typename Collection::value_type::second_type& value) {
135   return InsertIfNotPresent(collection,
136                             typename Collection::value_type(key, value));
137 }
138 
139 // Looks up a given key and value pair in a collection and inserts the key-value
140 // pair if it's not already present. Returns a reference to the value associated
141 // with the key.
142 template <class Collection>
LookupOrInsert(Collection * const collection,const typename Collection::value_type & vt)143 typename Collection::value_type::second_type& LookupOrInsert(
144     Collection* const collection, const typename Collection::value_type& vt) {
145   return collection->insert(vt).first->second;
146 }
147 
148 // Same as above except the key-value are passed separately.
149 template <class Collection>
LookupOrInsert(Collection * const collection,const typename Collection::value_type::first_type & key,const typename Collection::value_type::second_type & value)150 typename Collection::value_type::second_type& LookupOrInsert(
151     Collection* const collection,
152     const typename Collection::value_type::first_type& key,
153     const typename Collection::value_type::second_type& value) {
154   return LookupOrInsert(collection,
155                         typename Collection::value_type(key, value));
156 }
157 
158 }  // namespace gtl
159 }  // namespace tensorflow
160 
161 #endif  // TENSORFLOW_LIB_GTL_MAP_UTIL_H_
162