use std::str::FromStr;
use crate::model::prelude::*;
use tokio::sync::RwLock;
use std::collections::{
hash_map::Entry,
HashMap,
HashSet,
VecDeque,
};
use std::default::Default;
use async_trait::async_trait;
use tracing::instrument;
mod cache_update;
mod settings;
pub use self::cache_update::CacheUpdate;
pub use self::settings::Settings;
type MessageCache = HashMap<ChannelId, HashMap<MessageId, Message>>;
#[async_trait]
pub trait FromStrAndCache: Sized {
type Err;
async fn from_str<CRL>(cache: CRL, s: &str) -> Result<Self, Self::Err>
where CRL: AsRef<Cache> + Send + Sync;
}
#[async_trait]
pub trait StrExt: Sized {
async fn parse_cached<CRL, F: FromStrAndCache>(&self, cache: CRL) -> Result<F, F::Err>
where CRL: AsRef<Cache> + Send + Sync;
}
#[async_trait]
impl StrExt for &str {
async fn parse_cached<CRL, F: FromStrAndCache>(&self, cache: CRL) -> Result<F, F::Err>
where CRL: AsRef<Cache> + Send + Sync
{
F::from_str(&cache, &self).await
}
}
#[async_trait]
impl<F: FromStr> FromStrAndCache for F {
type Err = F::Err;
async fn from_str<CRL>(_cache: CRL, s: &str) -> Result<Self, Self::Err>
where CRL: AsRef<Cache> + Send + Sync
{
s.parse::<F>()
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct Cache {
pub(crate) channels: RwLock<HashMap<ChannelId, GuildChannel>>,
pub(crate) categories: RwLock<HashMap<ChannelId, ChannelCategory>>,
pub(crate) guilds: RwLock<HashMap<GuildId, Guild>>,
pub(crate) messages: RwLock<MessageCache>,
pub(crate) presences: RwLock<HashMap<UserId, Presence>>,
pub(crate) private_channels: RwLock<HashMap<ChannelId, PrivateChannel>>,
pub(crate) shard_count: RwLock<u64>,
pub(crate) unavailable_guilds: RwLock<HashSet<GuildId>>,
pub(crate) user: RwLock<CurrentUser>,
pub(crate) users: RwLock<HashMap<UserId, User>>,
pub(crate) message_queue: RwLock<HashMap<ChannelId, VecDeque<MessageId>>>,
settings: RwLock<Settings>,
}
impl Cache {
#[inline]
pub fn new() -> Self {
Self::default()
}
#[instrument]
pub fn new_with_settings(settings: Settings) -> Self {
Self {
settings: RwLock::new(settings),
..Default::default()
}
}
pub async fn unknown_members(&self) -> u64 {
let mut total = 0;
for guild in self.guilds.read().await.values() {
let members = guild.members.len() as u64;
if guild.member_count > members {
total += guild.member_count - members;
}
}
total
}
pub async fn private_channels(&self) -> HashMap<ChannelId, PrivateChannel> {
self.private_channels.read().await.clone()
}
pub async fn guilds(&self) -> Vec<GuildId> {
let chain = self.unavailable_guilds.read().await.clone().into_iter();
self.guilds
.read()
.await
.keys()
.cloned()
.chain(chain)
.collect()
}
#[inline]
pub async fn channel<C: Into<ChannelId>>(&self, id: C) -> Option<Channel> {
self._channel(id.into()).await
}
async fn _channel(&self, id: ChannelId) -> Option<Channel> {
if let Some(channel) = self.channels.read().await.get(&id) {
let channel = channel.clone();
return Some(Channel::Guild(channel));
}
if let Some(private_channel) = self.private_channels.read().await.get(&id).cloned() {
return Some(Channel::Private(private_channel));
}
None
}
#[inline]
pub async fn guild<G: Into<GuildId>>(&self, id: G) -> Option<Guild> {
self._guild(id.into()).await
}
async fn _guild(&self, id: GuildId) -> Option<Guild> {
self.guilds.read().await.get(&id).cloned()
}
#[inline]
pub async fn guild_field<Ret, Fun>(&self, id: impl Into<GuildId>, field_selector: Fun) -> Option<Ret>
where Fun: FnOnce(&Guild) -> Ret {
self._guild_field(id.into(), field_selector).await
}
async fn _guild_field<Ret, Fun>(&self, id: GuildId, field_accessor: Fun) -> Option<Ret>
where Fun: FnOnce(&Guild) -> Ret {
let guilds = self.guilds.read().await;
let guild = guilds.get(&id)?;
Some(field_accessor(guild))
}
pub async fn guild_count(&self) -> usize {
self.guilds.read().await.len()
}
#[inline]
pub async fn guild_channel<C: Into<ChannelId>>(&self, id: C) -> Option<GuildChannel> {
self._guild_channel(id.into()).await
}
async fn _guild_channel(&self, id: ChannelId) -> Option<GuildChannel> {
self.channels.read().await.get(&id).cloned()
}
#[inline]
pub async fn guild_channel_field<Ret, Fun>(&self,
id: impl Into<ChannelId>,
field_selector: Fun) -> Option<Ret>
where Fun: FnOnce(&GuildChannel) -> Ret {
self._guild_channel_field(id.into(), field_selector).await
}
async fn _guild_channel_field<Ret, Fun>(&self,
id: ChannelId,
field_selector: Fun) -> Option<Ret>
where Fun: FnOnce(&GuildChannel) -> Ret {
let guild_channels = &self.channels.read().await;
let channel = guild_channels.get(&id)?;
Some(field_selector(channel))
}
#[inline]
pub async fn member<G, U>(&self, guild_id: G, user_id: U) -> Option<Member>
where G: Into<GuildId>, U: Into<UserId> {
self._member(guild_id.into(), user_id.into()).await
}
async fn _member(&self, guild_id: GuildId, user_id: UserId) -> Option<Member> {
match self.guilds.read().await.get(&guild_id) {
Some(guild) => {
guild
.members
.get(&user_id)
.cloned()
}
None => None,
}
}
#[inline]
pub async fn member_field<Ret, Fun>(&self,
guild_id: impl Into<GuildId>,
user_id: impl Into<UserId>,
field_selector: Fun) -> Option<Ret>
where Fun: FnOnce(&Member) -> Ret {
self._member_field(guild_id.into(), user_id.into(), field_selector).await
}
async fn _member_field<Ret, Fun>(&self,
guild_id: GuildId,
user_id: UserId,
field_selector: Fun) -> Option<Ret>
where Fun: FnOnce(&Member) -> Ret {
let guilds = &self.guilds.read().await;
let guild = guilds.get(&guild_id)?;
let member = guild.members.get(&user_id)?;
Some(field_selector(member))
}
#[inline]
pub async fn guild_roles(&self, guild_id: impl Into<GuildId>) -> Option<HashMap<RoleId, Role>> {
self._guild_roles(guild_id.into()).await
}
async fn _guild_roles(&self, guild_id: GuildId) -> Option<HashMap<RoleId, Role>> {
self.guilds.read().await.get(&guild_id).map(|g| g.roles.clone())
}
#[inline]
pub async fn unavailable_guilds(&self) -> HashSet<GuildId> {
self.unavailable_guilds.read().await.clone()
}
#[inline]
pub async fn guild_channels(&self, guild_id: impl Into<GuildId>) -> Option<HashMap<ChannelId, GuildChannel>> {
self._guild_channels(guild_id.into()).await
}
async fn _guild_channels(&self, guild_id: GuildId) -> Option<HashMap<ChannelId, GuildChannel>> {
self.guilds.read().await.get(&guild_id).map(|g| g.channels.clone())
}
pub async fn guild_channel_count(&self) -> usize {
self.channels.read().await.len()
}
#[inline]
pub async fn shard_count(&self) -> u64 {
*self.shard_count.read().await
}
#[inline]
pub async fn message<C, M>(&self, channel_id: C, message_id: M) -> Option<Message>
where C: Into<ChannelId>, M: Into<MessageId> {
self._message(channel_id.into(), message_id.into()).await
}
async fn _message(&self, channel_id: ChannelId, message_id: MessageId) -> Option<Message> {
self.messages.read().await.get(&channel_id).and_then(|messages| {
messages.get(&message_id).cloned()
})
}
#[inline]
pub async fn private_channel(&self, channel_id: impl Into<ChannelId>) -> Option<PrivateChannel> {
self._private_channel(channel_id.into()).await
}
async fn _private_channel(&self, channel_id: ChannelId) -> Option<PrivateChannel> {
self.private_channels.read().await.get(&channel_id).cloned()
}
#[inline]
pub async fn role<G, R>(&self, guild_id: G, role_id: R) -> Option<Role>
where G: Into<GuildId>, R: Into<RoleId> {
self._role(guild_id.into(), role_id.into()).await
}
async fn _role(&self, guild_id: GuildId, role_id: RoleId) -> Option<Role> {
self.guilds.read().await.get(&guild_id).and_then(|g| g.roles.get(&role_id)).cloned()
}
pub async fn settings(&self) -> Settings {
self.settings.read().await.clone()
}
pub async fn set_max_messages(&self, max: usize) {
self.settings.write().await.max_messages = max;
}
#[inline]
pub async fn user<U: Into<UserId>>(&self, user_id: U) -> Option<User> {
self._user(user_id.into()).await
}
async fn _user(&self, user_id: UserId) -> Option<User> {
self.users.read().await.get(&user_id).cloned()
}
#[inline]
pub async fn users(&self) -> HashMap<UserId, User> {
self.users.read().await.clone()
}
#[inline]
pub async fn user_count(&self) -> usize {
self.users.read().await.len()
}
#[inline]
pub async fn category<C: Into<ChannelId>>(&self, channel_id: C) -> Option<ChannelCategory> {
self._category(channel_id.into()).await
}
async fn _category(&self, channel_id: ChannelId) -> Option<ChannelCategory> {
self.categories.read().await.get(&channel_id).cloned()
}
#[inline]
pub async fn categories(&self) -> HashMap<ChannelId, ChannelCategory> {
self.categories.read().await.clone()
}
#[inline]
pub async fn category_count(&self) -> usize {
self.categories.read().await.len()
}
#[inline]
pub async fn current_user(&self) -> CurrentUser {
self.user.read().await.clone()
}
#[inline]
pub async fn current_user_id(&self) -> UserId {
self.user.read().await.id
}
#[inline]
pub async fn current_user_field<Ret: Clone, Fun>(&self,
field_selector: Fun) -> Ret
where Fun: FnOnce(&CurrentUser) -> Ret {
let user = self.user.read().await;
field_selector(&user)
}
#[instrument(skip(self, e))]
pub async fn update<E: CacheUpdate>(&self, e: &mut E) -> Option<E::Output> {
e.update(self).await
}
pub(crate) async fn update_user_entry(&self, user: &User) {
match self.users.write().await.entry(user.id) {
Entry::Vacant(e) => {
e.insert(user.clone());
},
Entry::Occupied(mut e) => {
e.get_mut().clone_from(user);
},
}
}
}
impl Default for Cache {
fn default() -> Cache {
Cache {
channels: RwLock::new(HashMap::default()),
categories: RwLock::new(HashMap::default()),
guilds: RwLock::new(HashMap::default()),
messages: RwLock::new(HashMap::default()),
presences: RwLock::new(HashMap::default()),
private_channels: RwLock::new(HashMap::with_capacity(128)),
settings: RwLock::new(Settings::default()),
shard_count: RwLock::new(1),
unavailable_guilds: RwLock::new(HashSet::default()),
user: RwLock::new(CurrentUser::default()),
users: RwLock::new(HashMap::default()),
message_queue: RwLock::new(HashMap::default()),
}
}
}
#[cfg(test)]
mod test {
use chrono::{DateTime, Utc};
use serde_json::{Number, Value};
use std::collections::HashMap;
use crate::{
cache::{Cache, CacheUpdate, Settings},
model::prelude::*,
};
#[tokio::test]
async fn test_cache_messages() {
let mut settings = Settings::new();
settings.max_messages(2);
let mut cache = Cache::new_with_settings(settings);
let datetime = DateTime::parse_from_str(
"1983 Apr 13 12:09:14.274 +0000",
"%Y %b %d %H:%M:%S%.3f %z",
).unwrap()
.with_timezone(&Utc);
let mut event = MessageCreateEvent {
message: Message {
id: MessageId(3),
attachments: vec![],
author: User {
id: UserId(2),
avatar: None,
bot: false,
discriminator: 1,
name: "user 1".to_owned(),
_nonexhaustive: (),
},
channel_id: ChannelId(2),
guild_id: Some(GuildId(1)),
content: String::new(),
edited_timestamp: None,
embeds: vec![],
kind: MessageType::Regular,
member: None,
mention_everyone: false,
mention_roles: vec![],
mention_channels: vec![],
mentions: vec![],
nonce: Value::Number(Number::from(1)),
pinned: false,
reactions: vec![],
timestamp: datetime.clone(),
tts: false,
webhook_id: None,
activity: None,
application: None,
message_reference: None,
flags: None,
_nonexhaustive: (),
},
_nonexhaustive: (),
};
assert!(!cache.messages.read().await.contains_key(&event.message.channel_id));
assert!(event.update(&mut cache).await.is_none());
assert!(event.update(&mut cache).await.is_none());
assert_eq!(cache.messages.read().await.get(&event.message.channel_id).unwrap().len(), 1);
event.message.id = MessageId(4);
assert!(event.update(&mut cache).await.is_none());
assert_eq!(cache.messages.read().await.get(&event.message.channel_id).unwrap().len(), 2);
event.message.id = MessageId(5);
assert!(event.update(&mut cache).await.is_some());
{
let messages = cache.messages.read().await;
let channel = messages.get(&event.message.channel_id).unwrap();
assert_eq!(channel.len(), 2);
assert!(!channel.contains_key(&MessageId(3)));
}
let guild_channel = GuildChannel {
id: event.message.channel_id,
bitrate: None,
category_id: None,
guild_id: event.message.guild_id.unwrap(),
kind: ChannelType::Text,
last_message_id: None,
last_pin_timestamp: None,
name: String::new(),
permission_overwrites: vec![],
position: 0,
topic: None,
user_limit: None,
nsfw: false,
slow_mode_rate: Some(0),
_nonexhaustive: (),
};
let mut delete = ChannelDeleteEvent {
channel: Channel::Guild(guild_channel.clone()),
_nonexhaustive: (),
};
assert!(cache.update(&mut delete).await.is_none());
assert!(!cache.messages.read().await.contains_key(&delete.channel.id()));
let mut guild_create = {
let mut channels = HashMap::new();
channels.insert(ChannelId(2), guild_channel.clone());
GuildCreateEvent {
guild: Guild {
id: GuildId(1),
afk_channel_id: None,
afk_timeout: 0,
application_id: None,
default_message_notifications: DefaultMessageNotificationLevel::All,
emojis: HashMap::new(),
explicit_content_filter: ExplicitContentFilter::None,
features: vec![],
icon: None,
joined_at: datetime,
large: false,
member_count: 0,
members: HashMap::new(),
mfa_level: MfaLevel::None,
name: String::new(),
owner_id: UserId(3),
presences: HashMap::new(),
region: String::new(),
roles: HashMap::new(),
splash: None,
system_channel_id: None,
verification_level: VerificationLevel::Low,
voice_states: HashMap::new(),
description: None,
premium_tier: PremiumTier::Tier0,
channels,
premium_subscription_count: 0,
banner: None,
vanity_url_code: Some("bruhmoment".to_string()),
preferred_locale: "en-US".to_string(),
_nonexhaustive: (),
},
_nonexhaustive: (),
}
};
assert!(cache.update(&mut guild_create).await.is_none());
assert!(cache.update(&mut event).await.is_none());
let mut guild_delete = GuildDeleteEvent {
guild: GuildUnavailable {
id: GuildId(1),
unavailable: false,
},
_nonexhaustive: (),
};
assert!(cache.update(&mut guild_delete).await.is_some());
assert!(!cache.messages.read().await.contains_key(&ChannelId(2)));
}
}