src/limitedmap.h

   1 // Copyright (c) 2012 The Bitcoin developers
   2 // Distributed under the MIT/X11 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> // TODO: remove
   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> class limitedmap
  13 {
  14 public:
  15     typedef K key_type;
  16     typedef V mapped_type;
  17     typedef std::pair<const key_type, mapped_type> value_type;
  18     typedef typename std::map<K, V>::const_iterator const_iterator;
  19     typedef typename std::map<K, V>::size_type size_type;
  20
  21 protected:
  22     std::map<K, V> map;
  23     typedef typename std::map<K, V>::iterator iterator;
  24     std::multimap<V, iterator> rmap;
  25     typedef typename std::multimap<V, iterator>::iterator rmap_iterator;
  26     size_type nMaxSize;
  27
  28 public:
  29     limitedmap(size_type nMaxSizeIn = 0) { nMaxSize = nMaxSizeIn; }
  30     const_iterator begin() const { return map.begin(); }
  31     const_iterator end() const { return map.end(); }
  32     size_type size() const { return map.size(); }
  33     bool empty() const { return map.empty(); }
  34     const_iterator find(const key_type& k) const { return map.find(k); }
  35     size_type count(const key_type& k) const { return map.count(k); }
  36     void insert(const value_type& x)
  37     {
  38         std::pair<iterator, bool> ret = map.insert(x);
  39         if (ret.second)
  40         {
  41             if (nMaxSize && map.size() == nMaxSize)
  42             {
  43                 map.erase(rmap.begin()->second);
  44                 rmap.erase(rmap.begin());
  45             }
  46             rmap.insert(make_pair(x.second, ret.first));
  47         }
  48         return;
  49     }
  50     void erase(const key_type& k)
  51     {
  52         iterator itTarget = map.find(k);
  53         if (itTarget == map.end())
  54             return;
  55         std::pair<rmap_iterator, rmap_iterator> itPair = rmap.equal_range(itTarget->second);
  56         for (rmap_iterator it = itPair.first; it != itPair.second; ++it)
  57             if (it->second == itTarget)
  58             {
  59                 rmap.erase(it);
  60                 map.erase(itTarget);
  61                 return;
  62             }
  63         // Shouldn't ever get here
  64         assert(0); //TODO remove me
  65         map.erase(itTarget);
  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             {
  77                 rmap.erase(it);
  78                 itTarget->second = v;
  79                 rmap.insert(make_pair(v, itTarget));
  80                 return;
  81             }
  82         // Shouldn't ever get here
  83         assert(0); //TODO remove me
  84         itTarget->second = v;
  85         rmap.insert(make_pair(v, itTarget));
  86     }
  87     size_type max_size() const { return nMaxSize; }
  88     size_type max_size(size_type s)
  89     {
  90         if (s)
  91             while (map.size() > s)
  92             {
  93                 map.erase(rmap.begin()->second);
  94                 rmap.erase(rmap.begin());
  95             }
  96         nMaxSize = s;
  97         return nMaxSize;
  98     }
  99 };
 100
 101 #endif // BITCOIN_LIMITEDMAP_H