Add relations endpoints, edits and threads work now

2 files changed, 55 insertions, 8 deletions

diff --git a/src/database/key_value/rooms/pdu_metadata.rs b/src/database/key_value/rooms/pdu_metadata.rs
index 4b3f810..0641f9d 100644
--- a/src/database/key_value/rooms/pdu_metadata.rs
+++ b/src/database/key_value/rooms/pdu_metadata.rs
@@ -1,17 +1,64 @@
-use std::sync::Arc;
+use std::{mem, sync::Arc};
 
-use ruma::{EventId, RoomId};
+use ruma::{EventId, RoomId, UserId};
 
-use crate::{database::KeyValueDatabase, service, Result};
+use crate::{
+    database::KeyValueDatabase,
+    service::{self, rooms::timeline::PduCount},
+    services, utils, Error, PduEvent, Result,
+};
 
 impl service::rooms::pdu_metadata::Data for KeyValueDatabase {
     fn add_relation(&self, from: u64, to: u64) -> Result<()> {
-        let mut key = from.to_be_bytes().to_vec();
-        key.extend_from_slice(&to.to_be_bytes());
-        self.fromto_relation.insert(&key, &[])?;
+        let mut key = to.to_be_bytes().to_vec();
+        key.extend_from_slice(&from.to_be_bytes());
+        self.tofrom_relation.insert(&key, &[])?;
         Ok(())
     }
 
+    fn relations_until<'a>(
+        &'a self,
+        user_id: &'a UserId,
+        shortroomid: u64,
+        target: u64,
+        until: PduCount,
+    ) -> Result<Box<dyn Iterator<Item = Result<(PduCount, PduEvent)>> + 'a>> {
+        let prefix = target.to_be_bytes().to_vec();
+        let mut current = prefix.clone();
+
+        let count_raw = match until {
+            PduCount::Normal(x) => x - 1,
+            PduCount::Backfilled(x) => {
+                current.extend_from_slice(&0_u64.to_be_bytes());
+                u64::MAX - x - 1
+            }
+        };
+        current.extend_from_slice(&count_raw.to_be_bytes());
+
+        Ok(Box::new(
+            self.tofrom_relation
+                .iter_from(&current, true)
+                .take_while(move |(k, _)| k.starts_with(&prefix))
+                .map(move |(tofrom, _data)| {
+                    let from = utils::u64_from_bytes(&tofrom[(mem::size_of::<u64>())..])
+                        .map_err(|_| Error::bad_database("Invalid count in tofrom_relation."))?;
+
+                    let mut pduid = shortroomid.to_be_bytes().to_vec();
+                    pduid.extend_from_slice(&from.to_be_bytes());
+
+                    let mut pdu = services()
+                        .rooms
+                        .timeline
+                        .get_pdu_from_id(&pduid)?
+                        .ok_or_else(|| Error::bad_database("Pdu in tofrom_relation is invalid."))?;
+                    if pdu.sender != user_id {
+                        pdu.remove_transaction_id()?;
+                    }
+                    Ok((PduCount::Normal(from), pdu))
+                }),
+        ))
+    }
+
     fn mark_as_referenced(&self, room_id: &RoomId, event_ids: &[Arc<EventId>]) -> Result<()> {
         for prev in event_ids {
             let mut key = room_id.as_bytes().to_vec();
diff --git a/src/database/mod.rs b/src/database/mod.rs
index b864ceb..5d89d4a 100644
--- a/src/database/mod.rs
+++ b/src/database/mod.rs
@@ -131,7 +131,7 @@ pub struct KeyValueDatabase {
     pub(super) softfailedeventids: Arc<dyn KvTree>,
 
     /// ShortEventId + ShortEventId -> ().
-    pub(super) fromto_relation: Arc<dyn KvTree>,
+    pub(super) tofrom_relation: Arc<dyn KvTree>,
     /// RoomId + EventId -> Parent PDU EventId.
     pub(super) referencedevents: Arc<dyn KvTree>,
 
@@ -348,7 +348,7 @@ impl KeyValueDatabase {
             eventid_outlierpdu: builder.open_tree("eventid_outlierpdu")?,
             softfailedeventids: builder.open_tree("softfailedeventids")?,
 
-            fromto_relation: builder.open_tree("fromto_relation")?,
+            tofrom_relation: builder.open_tree("tofrom_relation")?,
             referencedevents: builder.open_tree("referencedevents")?,
             roomuserdataid_accountdata: builder.open_tree("roomuserdataid_accountdata")?,
             roomusertype_roomuserdataid: builder.open_tree("roomusertype_roomuserdataid")?,