ジェネリック型制約なしに WeakDictionary を実装

パフォーマンスのために，データの生成にコストがかかるオブジェクトを辞書にキャッシュしておく，ということをよく行います。生成するオブジェクトが少ないならあまり気にする必要はないのですが，規模が大きくなるとメモリーリークの問題が出てきます。 Java には WeakHashMap のような弱い参照でキーを保持することにより，ガベージコレクションの恩恵を受けることができます。しかし .NET Framework には弱い参照を扱うクラスはおそらく存在しません。そのため，たとえば Dictionary<string, WeakReference> のような辞書を作ったりするわけですが，これは少し扱いづらいです。

検索するといくつか実装が見つかったのですが (下記)，いずれもキー・値いずれにもクラス制約がついています。確かに弱い参照はクラスにしか使えないものですが，キーも値もクラスである必然性はありません。キーのみが弱い参照，あるいは値のみが弱い参照でも良いはずです。

ということで，自分でも WeakDictionary を作ってみました。

参照の抽象化

基本的な方針は『SilverlightやWPFやってっと、弱い参照のキーと値を管理するWeakDictionary<TKey, TValue>が欲しくなるよね。』にあるように，参照を抽象化することで，強い参照も弱い参照も null 参照も同一視してしまうことです。

public abstract class Reference<T>
{
   private readonly int hashCode;

   protected Reference(int hashCode)
   {
      this.hashCode = hashCode;
   }

   public abstract T Target { get; }
   public abstract bool IsAlive { get; }

   public override int GetHashCode()
   {
      return this.hashCode;
   }
}

ハッシュ値を生成時に決定することで， WeakDictionary 内での IEqualityComparer によるキーの比較を容易にできます。

そして弱い参照，強い参照， null 参照を実装します。

internal sealed class WeakReference<T> : Reference<T>
{
   private readonly WeakReference reference;

   internal WeakReference(T target, int hashCode)
      : base(hashCode)
   {
      this.reference = new WeakReference(target, false);
   }

   public override T Target
   {
      get { return (T)this.reference.Target; }
   }

   public override bool IsAlive
   {
      get { return this.reference.IsAlive; }
   }
}

internal sealed class StrongReference<T> : Reference<T>
{
   private readonly T target;

   internal StrongReference(T target, int hashCode)
      : base(hashCode)
   {
      this.target = target;
   }

   public override T Target
   {
      get { return this.target; }
   }

   public override bool IsAlive
   {
      get { return true; }
   }
}

internal sealed class NullReference<T> : Reference<T>
{
   private readonly isAlive;

   internal NullReference(bool isAlive)
      : base(typeof(T).GetHashCode())
   {
      this.isAlive = isAlive;
   }

   public override T Target
   {
      get { return default(T); }
   }

   public override bool IsAlive
   {
      get { return this.isAlive; }
   }
}

NullReference.IsAlive は null がキーや値として意味を持つかどうかを示します。 WeakDictionary の中では，弱い参照の場合は False，強い参照の場合は True にすることを想定しています。

参照の生成

参照は事前にハッシュ値を決定するので，外部から好き勝手に生成されては整合性を失う恐れがあります^[A]。そこで Reference を生成するためのクラスを作成します。 Object の比較に値の比較と参照の比較があるように， Reference のインスタンスが保持しているオブジェクトの比較にも値の比較が必要です。そのために IEqualityComparer を実装します。

public class ReferenceFactory<T> : IEqualityComparer<Reference<T>>
{
   private readonly IEqualityComparer<T> comparer;

   public ReferenceFactory(IEqualityComparer<T> comparer = null)
   {
      this.comparer = comparer ?? EqualityComparer<T>.Default;
   }

   public bool Equals(Reference<T> x, Reference<T> y)
   {
      if (x.IsAlive && y.IsAlive)
      {
         return this.comparer.Equals(x.Target, y.Target);
      }
      return ReferenceEquals(x, y);
   }

   public int GetHashCode(Reference<T> obj)
   {
      return obj.GetHashCode();
   }

   public Reference<T> Create(T target, bool isNullAvailable = false, bool isWeak = false)
   {
      if (target == null)
      {
         return new NullReference<T>(isNullAvailable);
      }
      var hash = this.comparer.GetHashCode(target);
      if (typeof(T).IsClass && isWeak)
      {
         return new WeakReference<T>(target, hash);
      }
      return new StrongReference<T>(target, hash);
   }
}

ReferenceFactory クラスのコンストラクターに Reference がラップするオブジェクトの比較のための IEqualityComparer を与えることで，参照同士の等値比較を実現しています。 Equals メソッドでは，比較する参照がいずれも生きていれば値の比較を行いますが，死んでいればハッシュ値のみ比較します。ハッシュ値は参照が作成されたときに生成されているので，回収前の値が等しければ True を返すことが保証されます^[B]。

Create メソッドで参照を作成します。値型は強い参照と同一視することで，これから実装する WeakDictionary のキーと値のジェネリック型に制約を加えなくて済むようにしています。

`WeakDictionary` の実装

まず核となる重要な部分です。

public class WeakDictionary<TKey, TValue> : IDictionary<TKey, TValue>
{
   private readonly Dictionary<Reference<TKey>, Reference<TValue>> internalDictionary;
   private readonly ReferenceFactory<TKey> keyFactory;
   private readonly ReferenceFactory<TValue> valueFactory;
   private readonly bool isWeakKey;
   private readonly bool isWeakValue;

   public WeakDictionary(bool isWeakKey, bool isWeakValue, IEqualityComparer<TKey> comparer = null)
   {
      this.keyFactory = new ReferenceFactory<TKey>(comparer);
      this.valueFactory = new ReferenceFactory<TValue>();
      this.internalDictionary = new Dictionary<Reference<TKey>, Reference<TValue>>(this.keyFactory);
      this.isWeakKey = isWeakKey;
      this.isWeakValue = isWeakValue;
   }

   public bool IsWeakKey
   {
      get { return this.isWeakKey; }
   }

   public bool IsWeakValue
   {
      get { return this.isWeakValue; }
   }

   private Reference<TKey> CreateReferenceKey(TKey key)
   {
      return this.keyFactory.Create(key, false, IsWeakKey);
   }

   private Reference<TValue> CreateReferenceValue(TValue value)
   {
      return this.valueFactory.Create(value, !IsWeakValue, IsWeakValue);
   }

   private void RemoveDead()
   {
      var removing = (from pair in this.internalDictionary
                      where !pair.Key.IsAlive || !pair.Value.IsAlive
                      select pair.Key).ToList();
      foreach (var key in removing)
      {
         this.internalDictionary.Remove(key);
      }
   }
}

ReferenceFactory.Create の実装で値型を渡したときは強い参照として扱うことにしているので，キーや値に値型を使う場合は isWeakKey や isWeakValue は無視されます。ここで不正なケースのチェックを行い，例外を投げるのも良いでしょう。

この実装では CreateReferenceKey の 2 番目の引数が常に False になっているので， null キーは参照の強弱によらず使用できません。 CreateReferenceValue の 2 番目の引数が !IsWeakValue なので，値に弱い参照を用いているときは null が値として使用できませんが，強い参照のときは null が許容されます。特にこうしなければいけないという理由はなく，この引数の指定次第で null の扱いが色々と変えられます。ちなみに IDictionary(Of TKey, TValue).Add メソッド (TKey, TValue) ではキーが null のときに例外を投げることになっています。

以降の操作で内部辞書である internalDictionary にアクセスする際に，キーや値を CreateReferenceKey および CreateReferenceValue を呼んで参照を作ってやれば良いです。適当なタイミングで RemoveDead を呼べば，内部辞書の要素でガベージコレクターに回収された弱い参照がなくなるという仕組みです。ただし RemoveDead は内部辞書のすべての要素をチェックする O(n) 操作です。頻繁に呼んでいてはコストがかかるので，適切なタイミングで呼んでやる必要があります。

たとえばCount プロパティ， Add メソッド， TryGetValue メソッドの実装は次のようになるでしょう。

public int Count
{
   get
   {
      RemoveDead();
      return this.internalDictionary.Count;
   }
}

public void Add(TKey key, TValue value)
{
   if (key == null)
   {
      throw new ArgumentNullException("key");
   }
   var referenceKey = CreateReferenceKey(key);
   var referenceValue = CreateReferenceValue(value);
   this.internalDictionary.Add(referenceKey, referenceValue);
}

public bool TryGetValue(TKey key, out TValue value)
{
   try
   {
      var referenceKey = CreateReferenceKey(key);
      Reference<TValue> referenceValue;
      if (this.internalDictionary.TryGetValue(referenceKey, out referenceValue))
      {
         value = referenceValue.Target;
         return referenceValue.IsAlive;
      }
      value = default(TValue);
      return false;
   }
   finally
   {
      RemoveDead();
   }
}

Count は事前に RemoveDead を呼んでガベージコレクターに回収された要素を除くことで，生きている要素のみ internalDictionary に残るので，その数が生きている要素数となります^[C]。 Add や TryGetValue は単純で，参照を生成してから内部辞書にアクセスしています。

他のプロパティやメソッドの実装は省略しますが，大体似たような感じです。

まとめ

where TKey : class where TValue : class なんていらんかったんや！

脚注

各クラスのコンストラクターが internal になっているのはそのためです。 [↩]
もちろん正しく GetHashCode を実装すればの話ですが。 [↩]
もちろん RemoveDead を呼んでから internalDictionary.Count を呼ぶまでの間にガベージコレクションが働けばその限りではありません。内部辞書の要素の Reference がラップしているオブジェクトすべてを lock すればガベージコレクションから守れる気もしますが (未確認)，そこまでして正確な要素数が必要な状況で WeakDictionary は使わないでしょう。 [↩]

捨てられたブログ

参照の抽象化

参照の生成

WeakDictionary の実装

まとめ

脚注

`WeakDictionary` の実装