src/limitedmap.h

   1 // Copyright (c) 2012-2014 The Bitcoin Core developers
   2 // Distributed under the MIT software license, see the accompanying
   3 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
   4
   5 #ifndef BITCOIN_LIMITEDMAP_H
   6 #define BITCOIN_LIMITEDMAP_H
   7
   8 #include <assert.h>
   9 #include <map>
  10
  11 /** STL-like map container that only keeps the N elements with the highest value. */
  12 template <typename K, typename V>
  13 class limitedmap
  14 {
  15 public:
  16     typedef K key_type;
  17     typedef V mapped_type;
  18     typedef std::pair<const key_type, mapped_type> value_type;
  19     typedef typename std::map<K, V>::const_iterator const_iterator;
  20     typedef typename std::map<K, V>::size_type size_type;
  21
  22 protected:
  23     std::map<K, V> map;
  24     typedef typename std::map<K, V>::iterator iterator;
  25     std::multimap<V, iterator> rmap;
  26     typedef typename std::multimap<V, iterator>::iterator rmap_iterator;
  27     size_type nMaxSize;
  28
  29 public:
  30     limitedmap(size_type nMaxSizeIn = 0) { nMaxSize = nMaxSizeIn; }
  31     const_iterator begin() const { return map.begin(); }
  32     const_iterator end() const { return map.end(); }
  33     size_type size() const { return map.size(); }
  34     bool empty() const { return map.empty(); }
  35     const_iterator find(const key_type& k) const { return map.find(k); }
  36     size_type count(const key_type& k) const { return map.count(k); }
  37     void insert(const value_type& x)
  38     {
  39         std::pair<iterator, bool> ret = map.insert(x);
  40         if (ret.second) {
  41             if (nMaxSize && map.size() == nMaxSize) {
  42                 map.erase(rmap.begin()->second);
  43                 rmap.erase(rmap.begin());
  44             }
  45             rmap.insert(make_pair(x.second, ret.first));
  46         }
  47         return;
  48     }
  49     void erase(const key_type& k)
  50     {
  51         iterator itTarget = map.find(k);
  52         if (itTarget == map.end())
  53             return;
  54         std::pair<rmap_iterator, rmap_iterator> itPair = rmap.equal_range(itTarget->second);
  55         for (rmap_iterator it = itPair.first; it != itPair.second; ++it)
  56             if (it->second == itTarget) {
  57                 rmap.erase(it);
  58                 map.erase(itTarget);
  59                 return;
  60             }
  61         // Shouldn't ever get here
  62         assert(0);
  63     }
  64     void update(const_iterator itIn, const mapped_type& v)
  65     {
  66         // TODO: When we switch to C++11, use map.erase(itIn, itIn) to get the non-const iterator.
  67         iterator itTarget = map.find(itIn->first);
  68         if (itTarget == map.end())
  69             return;
  70         std::pair<rmap_iterator, rmap_iterator> itPair = rmap.equal_range(itTarget->second);
  71         for (rmap_iterator it = itPair.first; it != itPair.second; ++it)
  72             if (it->second == itTarget) {
  73                 rmap.erase(it);
  74                 itTarget->second = v;
  75                 rmap.insert(make_pair(v, itTarget));
  76                 return;
  77             }
  78         // Shouldn't ever get here
  79         assert(0);
  80     }
  81     size_type max_size() const { return nMaxSize; }
  82     size_type max_size(size_type s)
  83     {
  84         if (s)
  85             while (map.size() > s) {
  86                 map.erase(rmap.begin()->second);
  87                 rmap.erase(rmap.begin());
  88             }
  89         nMaxSize = s;
  90         return nMaxSize;
  91     }
  92 };
  93
  94 #endif // BITCOIN_LIMITEDMAP_H