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)
  31     {
  32         assert(nMaxSizeIn > 0);
  33         nMaxSize = nMaxSizeIn;
  34     }
  35     const_iterator begin() const { return map.begin(); }
  36     const_iterator end() const { return map.end(); }
  37     size_type size() const { return map.size(); }
  38     bool empty() const { return map.empty(); }
  39     const_iterator find(const key_type& k) const { return map.find(k); }
  40     size_type count(const key_type& k) const { return map.count(k); }
  41     void insert(const value_type& x)
  42     {
  43         std::pair<iterator, bool> ret = map.insert(x);
  44         if (ret.second) {
  45             if (map.size() > nMaxSize) {
  46                 map.erase(rmap.begin()->second);
  47                 rmap.erase(rmap.begin());
  48             }
  49             rmap.insert(make_pair(x.second, ret.first));
  50         }
  51     }
  52     void erase(const key_type& k)
  53     {
  54         iterator itTarget = map.find(k);
  55         if (itTarget == map.end())
  56             return;
  57         std::pair<rmap_iterator, rmap_iterator> itPair = rmap.equal_range(itTarget->second);
  58         for (rmap_iterator it = itPair.first; it != itPair.second; ++it)
  59             if (it->second == itTarget) {
  60                 rmap.erase(it);
  61                 map.erase(itTarget);
  62                 return;
  63             }
  64         // Shouldn't ever get here
  65         assert(0);
  66     }
  67     void update(const_iterator itIn, const mapped_type& v)
  68     {
  69         // TODO: When we switch to C++11, use map.erase(itIn, itIn) to get the non-const iterator.
  70         iterator itTarget = map.find(itIn->first);
  71         if (itTarget == map.end())
  72             return;
  73         std::pair<rmap_iterator, rmap_iterator> itPair = rmap.equal_range(itTarget->second);
  74         for (rmap_iterator it = itPair.first; it != itPair.second; ++it)
  75             if (it->second == itTarget) {
  76                 rmap.erase(it);
  77                 itTarget->second = v;
  78                 rmap.insert(make_pair(v, itTarget));
  79                 return;
  80             }
  81         // Shouldn't ever get here
  82         assert(0);
  83     }
  84     size_type max_size() const { return nMaxSize; }
  85     size_type max_size(size_type s)
  86     {
  87         assert(s > 0);
  88         while (map.size() > s) {
  89             map.erase(rmap.begin()->second);
  90             rmap.erase(rmap.begin());
  91         }
  92         nMaxSize = s;
  93         return nMaxSize;
  94     }
  95 };
  96
  97 #endif // BITCOIN_LIMITEDMAP_H