From 346ee5cd81409e1ce931e5b4a5a78f3220373b10 Mon Sep 17 00:00:00 2001
From: Charles Hall <charles@computer.surgery>
Date: Sat, 2 Nov 2024 12:46:22 -0700
Subject: [PATCH] actually, just copy dump-state

---
 src/cli/repair.rs | 137 +++++++++++++++++++---------------------------
 1 file changed, 55 insertions(+), 82 deletions(-)

diff --git a/src/cli/repair.rs b/src/cli/repair.rs
index 88c3fb69..abf4be0f 100644
--- a/src/cli/repair.rs
+++ b/src/cli/repair.rs
@@ -3,6 +3,7 @@
 //! Implementation of the `repair` subcommand
 
 use std::{
+    cmp::Ordering as CmpOrdering,
     collections::{HashMap, HashSet},
     error::Error,
     sync::{
@@ -18,23 +19,43 @@ use petgraph::{
 };
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use ruma::{
-    events::{
-        room::member::MembershipState, StateEventType, TimelineEventType,
-    },
+    events::{StateEventType, TimelineEventType},
     state_res::StateMap,
     EventId, RoomId, RoomVersionId, UserId,
 };
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use tracing as t;
 
 use super::RepairArgs;
 use crate::{config, database::KeyValueDatabase, utils, PduEvent};
 
-/// Extractor to read just the membership state
-#[derive(Deserialize)]
-struct ExtractMembership {
-    /// Membership state
-    membership: MembershipState,
+/// Serializable information about a state event
+#[derive(Serialize, PartialEq, Eq)]
+struct StateEvent {
+    /// The kind of state event
+    kind: StateEventType,
+
+    /// The `state_key` of the event
+    key: String,
+
+    /// The event itself
+    event: Arc<PduEvent>,
+}
+
+impl Ord for StateEvent {
+    fn cmp(&self, other: &Self) -> CmpOrdering {
+        CmpOrdering::Equal
+            .then_with(|| self.event.room_id.cmp(&other.event.room_id))
+            .then_with(|| self.kind.cmp(&other.kind))
+            .then_with(|| self.key.cmp(&other.key))
+            .then_with(|| self.event.event_id.cmp(&other.event.event_id))
+    }
+}
+
+impl PartialOrd for StateEvent {
+    fn partial_cmp(&self, other: &Self) -> Option<CmpOrdering> {
+        Some(self.cmp(other))
+    }
 }
 
 /// Extractor to read just the `room_version` field
@@ -69,91 +90,43 @@ pub(crate) async fn run(args: RepairArgs) -> Result<(), Box<dyn Error>> {
 
     let states = resolve_room_states(&db, graphs);
 
-    repair_roomuserid_joined(&db, &states)?;
+    let mut serializable_state = Vec::new();
 
-    t::info!("Done");
-
-    Ok(())
-}
-
-/// Repair the `roomuserid_joined` map
-#[t::instrument(skip(db, states))]
-fn repair_roomuserid_joined(
-    db: &KeyValueDatabase,
-    states: &HashMap<Arc<RoomId>, StateMap<Arc<EventId>>>,
-) -> Result<(), Box<dyn Error>> {
-    t::info!(
-        length = db.roomuserid_joined.iter().count(),
-        "Original length of map",
-    );
-
-    db.roomuserid_joined.clear()?;
-
-    for (room_id, state) in states {
-        let states =
-            state.iter().filter_map(|((kind, key), id)| {
-                if *kind != StateEventType::RoomMember {
-                    return None;
-                }
-
-                let user_id = UserId::parse(key).inspect_err(|error| {
-                    t::error!(
-                        %error,
-                        %key,
-                        event_id = %id,
-                        "Failed to parse m.room.member state key as user ID",
-                    );
-                }).ok()?;
-
-                let pdu = get_accepted_pdu_by_event_id(db, id).ok()??;
-
-                let membership = match serde_json::from_str(pdu.content.get()) {
-                    Ok(ExtractMembership {
-                        membership,
-                    }) => membership,
+    for (_, state) in states {
+        serializable_state.extend(state.into_iter().filter_map(
+            |((kind, key), event)| {
+                let event = match get_either_pdu_by_event_id(&db, &event) {
+                    Ok(Some(x)) => Arc::new(x),
+                    Ok(None) => {
+                        t::warn!(event_id = %event, "Unknown event, omitting");
+                        return None;
+                    }
                     Err(error) => {
                         t::warn!(
                             %error,
-                            event_id = %id,
-                            "Failed to deserialize PDU content",
+                            event_id = %event,
+                            "Failed to get event, omitting",
                         );
                         return None;
                     }
                 };
 
-                Some((user_id, membership))
-            });
+                let x = StateEvent {
+                    kind,
+                    key,
+                    event,
+                };
 
-        for (user_id, membership) in states {
-            let is_joined = match membership {
-                MembershipState::Invite
-                | MembershipState::Ban
-                | MembershipState::Leave => false,
-
-                MembershipState::Join => true,
-
-                state => {
-                    t::warn!(
-                        %room_id,
-                        %user_id,
-                        %state,
-                        "Unsupported membership state",
-                    );
-                    false
-                }
-            };
-
-            if !is_joined {
-                continue;
-            }
-
-            let key = [room_id.as_bytes(), user_id.as_bytes()].join(&0xFF);
-
-            db.roomuserid_joined.insert(&key, &[])?;
-        }
+                Some(x)
+            },
+        ));
     }
 
-    t::info!(length = db.roomuserid_joined.iter().count(), "New length of map");
+    serializable_state.sort_unstable();
+
+    serde_json::to_writer(std::io::stdout(), &serializable_state)?;
+
+    t::info!("Done");
 
     Ok(())
 }