std::(unordered_)map でメモリ使用量を見積もる

以前，トライと STL コンテナの比較をした際，std::map, std::tr1::unordered_map についてはメモリ使用量をちゃんと測っていなかったが，都合により，コンテナ本体のメモリ使用量を見積もる必要が出てきたので gcc 4.7 の実装を眺めてみた．
結論から言うと，gcc では std::map <_Key, _Tp> のメモリ使用量は，

template<typename _Val>
struct _Rb_tree_node {
  _Rb_tree_color; // ノードの色 (enum; int)
  _Base_ptr;      // 親ノードへのポインタ
  _Base_ptr;      // 左ノードへのポインタ
  _Base_ptr;      // 右ノードへのポインタ
  _Val;           // value_type (std::pair <const _Key, _Tp>)
};
-> sizeof (std::_Rb_tree_node <std::pair <const _Key, _Tp> >) * size ()

で，std::unordered_map <_Key, _Tp> のメモリ使用量は，

class _Hashtable {
 _Node**;  // bucket (__detail::_Hash_node<_Value, __cache_hash_code>)
           //                              _Value: std::pair <const _Key, _Tp>
};
template<typename _Value, bool __cache_hash_code>
struct _Hash_node
{
  _Value;         // value_type (std::pair <const _key, _Tp>)
  // std::size_t; // hash 値の cache (__cache_hash_code = true 時)
  _Hash_node*;    // 同じ bucket 中の次のノードへのポインタ
};
-> sizeof (_Node*) * bucket_count () + sizeof (_Node) * size ()

でそれぞれ見積もることができるようだ（要素数に依存しないオーバーヘッドは無視; _Key や _Tp にポインタ（データメンバ）が含まれる場合は再帰的に辿ってメモリ使用量を加算する必要がある）．
コンテナの構築時には上記の見積もりより大きいメモリが必要となる場合があり，特に

• （アロケータを設計しない限り）std::(unordered_)map では要素ごとにメモリ確保が行われるので，byte-aligned なメモリ確保だと一要素辺りのオーバーヘッドが無視できなくなる
• std::unordered_map ではハッシュ表 (_Hash_table) に指数的な拡張戦略*1が用いられているので，rehash () で予め十分な bucket を確保しない限り拡張時にメモリのオーバーヘッドが生じる

という二点には注意する必要がある．
確認に使った gcc 4.3- 依存のコードは以下．

#include <cstdio>
#include <cstdlib>
#include <map>
#include <tr1/unordered_map>

typedef size_t  _Key;
typedef size_t  _Tp;

static const size_t ALIGN = 16;

template <typename T>
void put_to_n (T & m, const size_t N) {
  for (size_t i = 0; i < N; ++i)
    m.insert (typename T::value_type (i, 0));
}

int main (int argc, char ** argv) {
  if (argc < 3) {
    std::fprintf (stderr, "Usage: %s N ALIGN 0|1|2 (0=MAP, 1=HASH, 2=REHASH)\n",
                  argv[0]);
    std::exit (1);
  }
  enum type_t { MAP = 0, HASH = 1, REHASH = 2 };
  const type_t TYPE  = static_cast <type_t> (std::strtol (argv[1], NULL, 10));
  const size_t N     = std::strtol (argv[2], NULL, 10);
  if (TYPE == MAP) {
    typedef std::map <_Key, _Tp>  container_t;
    typedef std::_Rb_tree_node <container_t::value_type> _Node;
    container_t m;
    size_t node_size = sizeof (_Node);
    node_size = (node_size / ALIGN + (node_size % ALIGN ? 1 : 0)) * ALIGN;
    std::fprintf (stderr, "%ld bytes\n", node_size * N);
    put_to_n (m, N);
  } else {
    typedef std::tr1::unordered_map <_Key, _Tp>  container_t; // std::tr1
    typedef std::tr1::__detail::_Hash_node <container_t::value_type, false> _Node;
    container_t m;
    if (TYPE == REHASH)
      // m.reserve (N); // gcc 4.7 (c++11) only
      m.rehash (N / m.max_load_factor ()); // for gcc 4.3 - 4.6
    const size_t bucket_size = TYPE == REHASH ?
                               m.bucket_count () : N / m.max_load_factor ();
    size_t node_size = sizeof (_Node);
    node_size = (node_size / ALIGN + (node_size % ALIGN ? 1 : 0)) * ALIGN;
    std::fprintf (stderr, "%ld bytes\n",
                  sizeof (_Node *) * bucket_size + node_size * N);
    put_to_n (m, N);
  }
  return 0;
}

見れば分かるが，map <_Key, _Tp> -> set <_Key> に置換すれば std::(unordered_)set のメモリ使用量も見積もれる．MacBook Air (Mid 2011) 上の gcc 4.7.0 20111123 でコンパイルした時の結果は以下の通り．

> g++ -O2 -m32 test.cc -o test32
> g++ -O2 -m64 test.cc -o test64

> cont 1c run 'test[32|64]' '[0|1|2]' 1000000                             
> run test32 0 1000000
32000000 bytes
elapsed (real): 0.396s; RSS=31.5M

> run test32 1 1000000
20000000 bytes
elapsed (real): 0.176s; RSS=24.0M

> run test32 2 1000000
20225292 bytes
elapsed (real): 0.165s; RSS=20.0M

> run test64 0 1000000
48000000 bytes
elapsed (real): 0.361s; RSS=47.0M

> run test64 1 1000000
40000000 bytes
elapsed (real): 0.208s; RSS=47.4M

> run test64 2 1000000
40450584 bytes
elapsed (real): 0.186s; RSS=39.6M

std::map と予め rehash () (gcc 4.7 以降では reserve () がある) した std::unordered_map では，静的な情報のみで正確にメモリ使用量を見積もることができている．要素数が未知の場合など，予め rehash () できない場合，std::unordered_map では拡張時にオーバヘッドがかかる*2．
例によって，日本語もコードも適当なので，後でいろいろ修正する予定．
[追記] トライ（ダブル配列，簡潔データ構造）と STL コンテナの比較（最新版） - ny23の日記で STL コンテナ本体のメモリ使用量を算出してみたら，64-bit 環境では圧縮無しのダブル配列と同程度（キーが短い郵便番号辞書ではそれ以上の）のメモリを消費するという結果になった．

*1:実際はもう少し複雑で，hashtable_policy.h の _Prime_rehash_policy::_M_need_rehash () を眺めると分かるが，bucket 辺りの平均要素数 load_factor () = size () / bucket_count () が max_load_factor () (デフォルトで 1.0) を越えると，ハッシュ表のサイズが現在のサイズの _M_growth_factor (=2.f) 倍（以上の素数倍）まで拡張される．

*2:最後の拡張時の要素数を N' とした場合，最終的にコンテナが必要とするメモリより sizeof (_Node *) * N' / max_load_factor () - sizeof (_Node) * (N - N') bytes 程度余分にメモリが必要となる．sizeof (_Node) が小さい時には注意が必要．