From ad1ba0ab508453db5a74b0d5a5e57028669ce481 Mon Sep 17 00:00:00 2001
From: RogueException <rogueexception@gmail.com>
Date: Mon, 20 Jun 2016 21:14:19 -0300
Subject: [PATCH] Added DMChannels to DataStore

---
 src/Discord.Net/Data/DefaultDataStore.cs | 36 +++++++++++++++++++++++++++-----
 src/Discord.Net/Data/IDataStore.cs       |  4 ++++
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/Discord.Net/Data/DefaultDataStore.cs b/src/Discord.Net/Data/DefaultDataStore.cs
index 14dbcae31..c1d768ed9 100644
--- a/src/Discord.Net/Data/DefaultDataStore.cs
+++ b/src/Discord.Net/Data/DefaultDataStore.cs
@@ -8,12 +8,13 @@ namespace Discord.Data
 {
     public class DefaultDataStore : DataStore
     {
+        private const int CollectionConcurrencyLevel = 1; //WebSocket updater/event handler. //TODO: Needs profiling, increase to 2?
         private const double AverageChannelsPerGuild = 10.22; //Source: Googie2149
         private const double AverageUsersPerGuild = 47.78; //Source: Googie2149
-        private const double CollectionMultiplier = 1.05; //Add buffer to handle growth
-        private const double CollectionConcurrencyLevel = 1; //WebSocket updater/event handler. //TODO: Needs profiling, increase to 2?
+        private const double CollectionMultiplier = 1.05; //Add 5% buffer to handle growth
 
         private readonly ConcurrentDictionary<ulong, ICachedChannel> _channels;
+        private readonly ConcurrentDictionary<ulong, CachedDMChannel> _dmChannels;
         private readonly ConcurrentDictionary<ulong, CachedGuild> _guilds;
         private readonly ConcurrentDictionary<ulong, CachedPublicUser> _users;
 
@@ -25,9 +26,10 @@ namespace Discord.Data
         {
             double estimatedChannelCount = guildCount * AverageChannelsPerGuild + dmChannelCount;
             double estimatedUsersCount = guildCount * AverageUsersPerGuild;
-            _channels = new ConcurrentDictionary<ulong, ICachedChannel>(1, (int)(estimatedChannelCount * CollectionMultiplier));
-            _guilds = new ConcurrentDictionary<ulong, CachedGuild>(1, (int)(guildCount * CollectionMultiplier));
-            _users = new ConcurrentDictionary<ulong, CachedPublicUser>(1, (int)(estimatedUsersCount * CollectionMultiplier));
+            _channels = new ConcurrentDictionary<ulong, ICachedChannel>(CollectionConcurrencyLevel, (int)(estimatedChannelCount * CollectionMultiplier));
+            _dmChannels = new ConcurrentDictionary<ulong, CachedDMChannel>(CollectionConcurrencyLevel, (int)(dmChannelCount * CollectionMultiplier));
+            _guilds = new ConcurrentDictionary<ulong, CachedGuild>(CollectionConcurrencyLevel, (int)(guildCount * CollectionMultiplier));
+            _users = new ConcurrentDictionary<ulong, CachedPublicUser>(CollectionConcurrencyLevel, (int)(estimatedUsersCount * CollectionMultiplier));
         }
 
         internal override ICachedChannel GetChannel(ulong id)
@@ -49,6 +51,30 @@ namespace Discord.Data
             return null;
         }
 
+        internal override CachedDMChannel GetDMChannel(ulong userId)
+        {
+            CachedDMChannel channel;
+            if (_dmChannels.TryGetValue(userId, out channel))
+                return channel;
+            return null;
+        }
+        internal override void AddDMChannel(CachedDMChannel channel)
+        {
+            _channels[channel.Id] = channel;
+            _dmChannels[channel.Recipient.Id] = channel;
+        }
+        internal override CachedDMChannel RemoveDMChannel(ulong userId)
+        {
+            CachedDMChannel channel;
+            ICachedChannel ignored;
+            if (_dmChannels.TryRemove(userId, out channel))
+            {
+                if (_channels.TryRemove(channel.Id, out ignored))
+                    return channel;
+            }
+            return null;
+        }
+
         internal override CachedGuild GetGuild(ulong id)
         {
             CachedGuild guild;
diff --git a/src/Discord.Net/Data/IDataStore.cs b/src/Discord.Net/Data/IDataStore.cs
index cc849cd94..9d7525a35 100644
--- a/src/Discord.Net/Data/IDataStore.cs
+++ b/src/Discord.Net/Data/IDataStore.cs
@@ -13,6 +13,10 @@ namespace Discord.Data
         internal abstract void AddChannel(ICachedChannel channel);
         internal abstract ICachedChannel RemoveChannel(ulong id);
 
+        internal abstract CachedDMChannel GetDMChannel(ulong userId);
+        internal abstract void AddDMChannel(CachedDMChannel channel);
+        internal abstract CachedDMChannel RemoveDMChannel(ulong userId);
+
         internal abstract CachedGuild GetGuild(ulong id);
         internal abstract void AddGuild(CachedGuild guild);
         internal abstract CachedGuild RemoveGuild(ulong id);