You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
CrowdControl/Assets/Plugins/UniTask/Runtime/Internal/WeakDictionary.cs

335 lines
9.4 KiB
C#

#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
using System;
using System.Collections.Generic;
using System.Threading;
namespace Cysharp.Threading.Tasks.Internal
{
// Add, Remove, Enumerate with sweep. All operations are thread safe(in spinlock).
internal class WeakDictionary<TKey, TValue>
where TKey : class
{
Entry[] buckets;
int size;
SpinLock gate; // mutable struct(not readonly)
readonly float loadFactor;
readonly IEqualityComparer<TKey> keyEqualityComparer;
public WeakDictionary(int capacity = 4, float loadFactor = 0.75f, IEqualityComparer<TKey> keyComparer = null)
{
var tableSize = CalculateCapacity(capacity, loadFactor);
this.buckets = new Entry[tableSize];
this.loadFactor = loadFactor;
this.gate = new SpinLock(false);
this.keyEqualityComparer = keyComparer ?? EqualityComparer<TKey>.Default;
}
public bool TryAdd(TKey key, TValue value)
{
bool lockTaken = false;
try
{
gate.Enter(ref lockTaken);
return TryAddInternal(key, value);
}
finally
{
if (lockTaken) gate.Exit(false);
}
}
public bool TryGetValue(TKey key, out TValue value)
{
bool lockTaken = false;
try
{
gate.Enter(ref lockTaken);
if (TryGetEntry(key, out _, out var entry))
{
value = entry.Value;
return true;
}
value = default(TValue);
return false;
}
finally
{
if (lockTaken) gate.Exit(false);
}
}
public bool TryRemove(TKey key)
{
bool lockTaken = false;
try
{
gate.Enter(ref lockTaken);
if (TryGetEntry(key, out var hashIndex, out var entry))
{
Remove(hashIndex, entry);
return true;
}
return false;
}
finally
{
if (lockTaken) gate.Exit(false);
}
}
bool TryAddInternal(TKey key, TValue value)
{
var nextCapacity = CalculateCapacity(size + 1, loadFactor);
TRY_ADD_AGAIN:
if (buckets.Length < nextCapacity)
{
// rehash
var nextBucket = new Entry[nextCapacity];
for (int i = 0; i < buckets.Length; i++)
{
var e = buckets[i];
while (e != null)
{
AddToBuckets(nextBucket, key, e.Value, e.Hash);
e = e.Next;
}
}
buckets = nextBucket;
goto TRY_ADD_AGAIN;
}
else
{
// add entry
var successAdd = AddToBuckets(buckets, key, value, keyEqualityComparer.GetHashCode(key));
if (successAdd) size++;
return successAdd;
}
}
bool AddToBuckets(Entry[] targetBuckets, TKey newKey, TValue value, int keyHash)
{
var h = keyHash;
var hashIndex = h & (targetBuckets.Length - 1);
TRY_ADD_AGAIN:
if (targetBuckets[hashIndex] == null)
{
targetBuckets[hashIndex] = new Entry
{
Key = new WeakReference<TKey>(newKey, false),
Value = value,
Hash = h
};
return true;
}
else
{
// add to last.
var entry = targetBuckets[hashIndex];
while (entry != null)
{
if (entry.Key.TryGetTarget(out var target))
{
if (keyEqualityComparer.Equals(newKey, target))
{
return false; // duplicate
}
}
else
{
Remove(hashIndex, entry);
if (targetBuckets[hashIndex] == null) goto TRY_ADD_AGAIN; // add new entry
}
if (entry.Next != null)
{
entry = entry.Next;
}
else
{
// found last
entry.Next = new Entry
{
Key = new WeakReference<TKey>(newKey, false),
Value = value,
Hash = h
};
entry.Next.Prev = entry;
}
}
return false;
}
}
bool TryGetEntry(TKey key, out int hashIndex, out Entry entry)
{
var table = buckets;
var hash = keyEqualityComparer.GetHashCode(key);
hashIndex = hash & table.Length - 1;
entry = table[hashIndex];
while (entry != null)
{
if (entry.Key.TryGetTarget(out var target))
{
if (keyEqualityComparer.Equals(key, target))
{
return true;
}
}
else
{
// sweap
Remove(hashIndex, entry);
}
entry = entry.Next;
}
return false;
}
void Remove(int hashIndex, Entry entry)
{
if (entry.Prev == null && entry.Next == null)
{
buckets[hashIndex] = null;
}
else
{
if (entry.Prev == null)
{
buckets[hashIndex] = entry.Next;
}
if (entry.Prev != null)
{
entry.Prev.Next = entry.Next;
}
if (entry.Next != null)
{
entry.Next.Prev = entry.Prev;
}
}
size--;
}
public List<KeyValuePair<TKey, TValue>> ToList()
{
var list = new List<KeyValuePair<TKey, TValue>>(size);
ToList(ref list, false);
return list;
}
// avoid allocate everytime.
public int ToList(ref List<KeyValuePair<TKey, TValue>> list, bool clear = true)
{
if (clear)
{
list.Clear();
}
var listIndex = 0;
bool lockTaken = false;
try
{
for (int i = 0; i < buckets.Length; i++)
{
var entry = buckets[i];
while (entry != null)
{
if (entry.Key.TryGetTarget(out var target))
{
var item = new KeyValuePair<TKey, TValue>(target, entry.Value);
if (listIndex < list.Count)
{
list[listIndex++] = item;
}
else
{
list.Add(item);
listIndex++;
}
}
else
{
// sweap
Remove(i, entry);
}
entry = entry.Next;
}
}
}
finally
{
if (lockTaken) gate.Exit(false);
}
return listIndex;
}
static int CalculateCapacity(int collectionSize, float loadFactor)
{
var size = (int)(((float)collectionSize) / loadFactor);
size--;
size |= size >> 1;
size |= size >> 2;
size |= size >> 4;
size |= size >> 8;
size |= size >> 16;
size += 1;
if (size < 8)
{
size = 8;
}
return size;
}
class Entry
{
public WeakReference<TKey> Key;
public TValue Value;
public int Hash;
public Entry Prev;
public Entry Next;
// debug only
public override string ToString()
{
if (Key.TryGetTarget(out var target))
{
return target + "(" + Count() + ")";
}
else
{
return "(Dead)";
}
}
int Count()
{
var count = 1;
var n = this;
while (n.Next != null)
{
count++;
n = n.Next;
}
return count;
}
}
}
}