track reposts and add repost count

Author: Sychic

src/config.rs

@@ -7,6 +7,7 @@ use crate::pass::Pass;
 pub struct Config {
     pub token: String,
     pub reply_cache_size: usize,
+    pub seen_cache_size: Option<usize>,
     #[serde(default)]
     pub ignored_users: Vec<Id<UserMarker>>,
     #[serde(default)]

src/main.rs

@@ -3,6 +3,9 @@ use std::future::IntoFuture;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
+use anyhow::Ok;
+use cache::SeenCache;
+use repost::add_repost_counts;
 use twilight_gateway::{Event, EventTypeFlags, Intents, Shard, ShardId, StreamExt as _};
 use twilight_http::Client;
 use twilight_model::channel::message::{AllowedMentions, MessageFlags};
@@ -18,11 +21,13 @@ use crate::{cache::ReplyCache, config::Config};
 mod cache;
 mod config;
 mod pass;
+mod repost;
 
-struct State {
+pub struct State {
     config: Config,
     rest: Client,
     replies: RwLock<ReplyCache>,
+    seen: RwLock<SeenCache>,
 }
 
 #[tokio::main]
@@ -38,8 +43,12 @@ async fn main() -> Result<(), anyhow::Error> {
         Intents::GUILD_MESSAGES | Intents::MESSAGE_CONTENT,
     );
 
+    // Use config size if it exists, otherwise default to reply cache size
+    let seen_size = config.seen_cache_size.unwrap_or(config.reply_cache_size);
+
     let state = Arc::new(State {
         replies: RwLock::new(ReplyCache::with_capacity(config.reply_cache_size)),
+        seen: RwLock::new(SeenCache::with_capacity(seen_size)),
         config,
         rest,
     });
@@ -118,12 +127,23 @@ async fn dispatch_event(state: Arc<State>, event: Event) -> Result<(), anyhow::E
                         .await?;
 
                     state.replies.write().unwrap().insert(token, reply.id);
+                    if let Some(new_content) = add_repost_counts(&state, reply.id, &reply.content) {
+                        state.rest
+                            .update_message(reply.channel_id, reply.id)
+                            .content(Some(new_content).as_deref())
+                            .allowed_mentions(Some(&AllowedMentions::default()))
+                            .await?;
+                    }
                 }
             }
         }
 
         // UPDATE: Edit our reply when someone edits a link in/out
         Event::MessageUpdate(message) => {
+                        // Ignore non-content edits like removing embeds
+                        if message.edited_timestamp == None {
+                            return Ok(());
+                        }
             let entry = state.replies.read().unwrap().get_entry(message.id);
             let Some(entry) = entry else {
                 return Ok(());
@@ -143,11 +163,12 @@ async fn dispatch_event(state: Arc<State>, event: Event) -> Result<(), anyhow::E
             if let CacheEntry::Filled(reply_id) = entry {
                 if !message.content.is_empty() {
                     if let Some(content) = Pass::apply_all(&state.config.passes, &message.content) {
+                        let content = add_repost_counts(&state, reply_id, &content);
                         state
                             .rest
                             .update_message(message.channel_id, reply_id)
                             .allowed_mentions(Some(&AllowedMentions::default()))
-                            .content(Some(&content))
+                            .content(content.as_deref())
                             .await?;
                     } else {
                         state

src/repost.rs

@@ -0,0 +1,57 @@
+use std::sync::Arc;
+
+use regex::{Captures, Regex};
+use twilight_model::id::{marker::MessageMarker, Id};
+
+use crate::State;
+
+/// Pattern that matches urls which have been transformed by [pass]
+const URL_REGEX: &str = "\\[`\\w+`\\]\\((?P<url>.+)\\)";
+
+/// Counts the number of times that a url has been seen
+fn check_repost(state: &Arc<State>, embed_url: &str) -> usize {
+    let existing_posts = state.seen.read().unwrap().search_by_value(embed_url);
+    existing_posts.len()
+}
+
+/// Takes a string and inserts the number of times it has been seen
+pub fn add_repost_counts(state: &Arc<State>, reply_id: Id<MessageMarker>, content: &str) -> Option<String> {
+    let seen = state.seen.read().unwrap().get_entry(reply_id);
+    if let Some(_seen) = seen {
+        return Some(content.to_owned());
+    }
+    let mut new_content = content.to_owned();
+    for url in find_urls(content) {
+        let times = check_repost(state, url);
+        let token = state.seen.write().unwrap().file_pending(reply_id);
+        if let Some(token) = token {
+            if times == 0 {
+                state.seen.write().unwrap().insert(token, url.to_owned());
+                continue;
+            }
+            new_content = add_repost_count(content, times);
+            state.seen.write().unwrap().insert(token, url.to_owned());
+        }
+    }
+    if new_content == content {
+        return None
+    }
+    Some(new_content)
+}
+
+/// Returns a vector of URLs that exist in a string
+fn find_urls(content: &str) -> Vec<&str> {
+    Regex::new(URL_REGEX)
+        .unwrap()
+        .captures_iter(content)
+        .map(|e| e.name("url").unwrap().as_str())
+        .collect()
+}
+
+/// Adds the number of times a link has been reposted to a string
+fn add_repost_count(content: &str, repost_count: usize) -> String {
+    let regex = Regex::new(URL_REGEX).unwrap();
+    regex.replace(content, |caps: &Captures|
+        format!("{} Posted {} time(s) ", &caps[0], repost_count)
+    ).into_owned()
+}