fix: Correct relations pagination and recursion

Ended up rewriting a fair bit, but should work better.

src/service/rooms/pdu_metadata/bundled_aggregations.rs

@@ -1,14 +1,16 @@
-use conduwuit::{Event, PduEvent, Result, err};
+use conduwuit::{
+	Event, PduEvent, Result, err,
+	matrix::{PduCount, event::relation_type},
+	utils::ReadyExt,
+};
 use ruma::{
 	UserId,
 	api::Direction,
-	events::relation::{BundledMessageLikeRelations, BundledReference, ReferenceChunk},
+	events::relation::{
+		BundledMessageLikeRelations, BundledReference, ReferenceChunk, RelationType,
+	},
 };
 
-use crate::rooms::timeline::PdusIterItem;
-
-const MAX_BUNDLED_RELATIONS: usize = 50;
-
 impl super::Service {
 	/// Gets bundled aggregations for an event according to the Matrix
 	/// specification.
@@ -26,52 +28,24 @@ impl super::Service {
 			return Ok(None);
 		}
 
-		let relations = self
-			.get_relations(
-				user_id,
-				&pdu.room_id_or_hash(),
-				pdu.event_id(),
-				conduwuit::PduCount::max(),
-				MAX_BUNDLED_RELATIONS,
-				0,
-				Direction::Backward,
-			)
+		let Some((room_id, target)) = self
+			.relation_target(&pdu.room_id_or_hash(), pdu.event_id())
+			.await
+		else {
+			return Ok(None);
+		};
+
+		let mut replace_events = Vec::new();
+		let mut reference_events = Vec::new();
+		self.db
+			.get_relations(room_id, target, PduCount::max(), Direction::Backward)
+			.ready_for_each(|(_, relation)| match relation_type(&relation) {
+				| Some(RelationType::Replacement) => replace_events.push(relation),
+				| Some(RelationType::Reference) => reference_events.push(relation),
+				| _ => {},
+			})
 			.await;
 
-		// The relations database code still handles the basic unsigned data
-		// We don't want to recursively fetch relations
-
-		if relations.is_empty() {
-			return Ok(None);
-		}
-
-		// Partition relations by type
-		let (replace_events, reference_events): (Vec<_>, Vec<_>) = relations
-			.iter()
-			.filter_map(|relation| {
-				let pdu = &relation.1;
-				let content = pdu.get_content_as_value();
-
-				content
-					.get("m.relates_to")
-					.and_then(|relates_to| relates_to.get("rel_type"))
-					.and_then(|rel_type| rel_type.as_str())
-					.and_then(|rel_type_str| match rel_type_str {
-						| "m.replace" => Some(RelationType::Replace(relation)),
-						| "m.reference" => Some(RelationType::Reference(relation)),
-						| _ => None, /* Ignore other relation types (threads are in DB but not
-						              * handled here) */
-					})
-			})
-			.fold((Vec::new(), Vec::new()), |(mut replaces, mut references), rel_type| {
-				match rel_type {
-					| RelationType::Replace(r) => replaces.push(r),
-					| RelationType::Reference(r) => references.push(r),
-				}
-				(replaces, references)
-			});
-
-		// If no relations to bundle, return None
 		if replace_events.is_empty() && reference_events.is_empty() {
 			return Ok(None);
 		}
@@ -85,7 +59,9 @@ impl super::Service {
 				.find_most_recent_valid_replacement(user_id, pdu, &replace_events)
 				.await?
 			{
-				bundled.replace = Some(Self::serialize_replacement(replacement)?);
+				bundled.replace = Some(Self::serialize_replacement(
+					&self.with_unsigned(user_id, replacement.clone()).await,
+				)?);
 			}
 		}
 
@@ -93,7 +69,7 @@ impl super::Service {
 		if !reference_events.is_empty() {
 			let reference_chunk: Vec<_> = reference_events
 				.into_iter()
-				.map(|relation| BundledReference::new(relation.1.event_id().to_owned()))
+				.map(|relation| BundledReference::new(relation.event_id().to_owned()))
 				.collect();
 
 			if !reference_chunk.is_empty() {
@@ -123,14 +99,12 @@ impl super::Service {
 		&self,
 		user_id: &UserId,
 		original_event: &PduEvent,
-		replacement_events: &[&'a PdusIterItem],
+		replacement_events: &'a [PduEvent],
 	) -> Result<Option<&'a PduEvent>> {
 		// Filter valid replacements and find the maximum in a single pass
 		let mut result: Option<&PduEvent> = None;
 
-		for relation in replacement_events {
-			let pdu = &relation.1;
-
+		for pdu in replacement_events {
 			// Validate replacement
 			if !Self::is_valid_replacement_event(original_event, pdu).await? {
 				continue;
@@ -245,12 +219,7 @@ impl super::Service {
 	/// Check if an event is itself a replacement
 	#[inline]
 	fn is_replacement_event(event: &PduEvent) -> bool {
-		event
-			.get_content_as_value()
-			.get("m.relates_to")
-			.and_then(|relates_to| relates_to.get("rel_type"))
-			.and_then(|rel_type| rel_type.as_str())
-			.is_some_and(|rel_type| rel_type == "m.replace")
+		relation_type(event) == Some(RelationType::Replacement)
 	}
 
 	/// Check if event has m.new_content or is encrypted (where m.new_content
@@ -262,12 +231,6 @@ impl super::Service {
 	}
 }
 
-/// Helper enum for partitioning relations
-enum RelationType<'a> {
-	Replace(&'a PdusIterItem),
-	Reference(&'a PdusIterItem),
-}
-
 #[cfg(test)]
 mod tests {
 	use conduwuit_core::pdu::{EventHash, PduEvent};

src/service/rooms/pdu_metadata/data.rs

@@ -2,8 +2,7 @@ use std::{mem::size_of, sync::Arc};
 
 use conduwuit::{
 	arrayvec::ArrayVec,
-	matrix::PduCount,
-	result::LogErr,
+	matrix::{PduCount, PduEvent},
 	utils::{
 		ReadyExt,
 		stream::{TryIgnore, WidebandExt},
@@ -12,14 +11,14 @@ use conduwuit::{
 };
 use database::Map;
 use futures::{Stream, StreamExt};
-use ruma::{EventId, UserId, api::Direction};
+use ruma::{EventId, api::Direction};
 
 use crate::{
 	Dep,
 	rooms::{
 		self,
-		short::{ShortEventId, ShortRoomId},
-		timeline::{PduId, PdusIterItem, RawPduId},
+		short::ShortRoomId,
+		timeline::{PduId, RawPduId},
 	},
 };
 
@@ -56,51 +55,34 @@ impl Data {
 
 	pub(super) fn get_relations<'a>(
 		&'a self,
-		user_id: &'a UserId,
 		shortroomid: ShortRoomId,
-		target: ShortEventId,
+		target: u64,
 		from: PduCount,
 		dir: Direction,
-	) -> impl Stream<Item = PdusIterItem> + Send + 'a {
-		// Query from exact position then filter excludes it (saturating_inc
-		// could skip events at min/max boundaries)
-		let from_unsigned = from.into_unsigned();
-		let mut current = ArrayVec::<u8, 16>::new();
-		current.extend(target.to_be_bytes());
-		current.extend(from_unsigned.to_be_bytes());
-		let current = current.as_slice();
+	) -> impl Stream<Item = (u64, PduEvent)> + Send + 'a {
+		let start_count = match from {
+			| PduCount::Normal(count) => count,
+			| PduCount::Backfilled(_) => 0,
+		};
+		let mut start = ArrayVec::<u8, 16>::new();
+		start.extend(target.to_be_bytes());
+		start.extend(start_count.to_be_bytes());
+		let start = start.as_slice();
 		match dir {
-			| Direction::Forward => self.tofrom_relation.raw_keys_from(current).boxed(),
-			| Direction::Backward => self.tofrom_relation.rev_raw_keys_from(current).boxed(),
+			| Direction::Forward => self.tofrom_relation.raw_keys_from(start).boxed(),
+			| Direction::Backward => self.tofrom_relation.rev_raw_keys_from(start).boxed(),
 		}
 		.ignore_err()
 		.ready_take_while(move |key| key.starts_with(&target.to_be_bytes()))
 		.map(|to_from| u64_from_u8(&to_from[8..16]))
-		.map(PduCount::from_unsigned)
-		.ready_filter(move |count| {
-			if from == PduCount::min() || from == PduCount::max() {
-				true
-			} else {
-				let count_unsigned = count.into_unsigned();
-				match dir {
-					| Direction::Forward => count_unsigned > from_unsigned,
-					| Direction::Backward => count_unsigned < from_unsigned,
-				}
-			}
-		})
-		.wide_filter_map(move |shorteventid| async move {
+		// The scan includes `from` itself, which should be skipped
+		.ready_filter(move |&count| PduCount::Normal(count).is_after(from, dir))
+		.wide_filter_map(move |count| async move {
+			let shorteventid = PduCount::Normal(count);
 			let pdu_id: RawPduId = PduId { shortroomid, shorteventid }.into();
+			let pdu = self.services.timeline.get_pdu_from_id(&pdu_id).await.ok()?;
 
-			let mut pdu = self.services.timeline.get_pdu_from_id(&pdu_id).await.ok()?;
-
-			let ctx = self
-				.services
-				.timeline
-				.get_unsigned_context(&pdu, Some(user_id))
-				.await;
-			pdu.set_unsigned(Some(ctx)).log_err().ok();
-
-			Some((shorteventid, pdu))
+			Some((count, pdu))
 		})
 	}
 

src/service/rooms/pdu_metadata/mod.rs

@@ -1,15 +1,20 @@
 mod bundled_aggregations;
 mod data;
-use std::sync::Arc;
+use std::{collections::HashSet, sync::Arc};
 
-use conduwuit::{Result, matrix::PduCount};
+use conduwuit::{
+	Result,
+	matrix::{PduCount, PduEvent},
+	result::LogErr,
+	utils::{IterStream, ReadyExt, stream::WidebandExt},
+};
 use futures::{StreamExt, future::try_join};
 use ruma::{EventId, RoomId, UserId, api::Direction};
 
 use self::data::Data;
 use crate::{
 	Dep,
-	rooms::{self, timeline::PdusIterItem},
+	rooms::{self, short::ShortRoomId, timeline::PdusIterItem},
 };
 
 pub struct Service {
@@ -51,7 +56,7 @@ impl Service {
 	}
 
 	#[allow(clippy::too_many_arguments)]
-	pub async fn get_relations<'a>(
+	pub async fn get_relations<'a, F>(
 		&'a self,
 		user_id: &'a UserId,
 		room_id: &'a RoomId,
@@ -60,59 +65,113 @@ impl Service {
 		limit: usize,
 		max_depth: u8,
 		dir: Direction,
-	) -> Vec<PdusIterItem> {
-		let room_id = self.services.short.get_shortroomid(room_id);
-
-		let target = self.services.timeline.get_pdu_count(target);
-
-		let Ok((room_id, target)) = try_join(room_id, target).await else {
+		matches: F,
+	) -> Vec<PdusIterItem>
+	where
+		F: Fn(&PduEvent) -> bool + Send + Sync,
+	{
+		let Some((room_id, target)) = self.relation_target(room_id, target).await else {
 			return Vec::new();
 		};
 
-		let target = match target {
-			| PduCount::Normal(c) => c,
-			// TODO: Support backfilled relations
-			| _ => 0, // This will result in an empty iterator
-		};
-
-		let mut pdus: Vec<_> = self
-			.db
-			.get_relations(user_id, room_id, target, from, dir)
-			.collect()
-			.await;
-
-		let mut stack: Vec<_> = pdus
-			.iter()
-			.filter(|_| max_depth > 0)
-			.map(|pdu| (pdu.clone(), 1))
-			.collect();
-
-		'limit: while let Some(stack_pdu) = stack.pop() {
-			let target = match stack_pdu.0.0 {
-				| PduCount::Normal(c) => c,
-				// TODO: Support backfilled relations
-				| PduCount::Backfilled(_) => 0, // This will result in an empty iterator
-			};
-
-			let relations: Vec<_> = self
+		// Direct relations are already in order
+		if max_depth == 0 {
+			return self
 				.db
-				.get_relations(user_id, room_id, target, from, dir)
+				.get_relations(room_id, target, from, dir)
+				.ready_filter(|(_, pdu)| matches(pdu))
+				.take(limit)
+				.wide_then(|relation| self.finish_relation(user_id, relation))
 				.collect()
 				.await;
-
-			for relation in relations {
-				if stack_pdu.1 < max_depth {
-					stack.push((relation.clone(), stack_pdu.1.saturating_add(1)));
-				}
-
-				pdus.push(relation);
-				if pdus.len() >= limit {
-					break 'limit;
-				}
-			}
 		}
 
-		pdus
+		let scan_from = match dir {
+			| Direction::Forward => PduCount::min(),
+			| Direction::Backward => from,
+		};
+
+		// An event can relate to two others in the tree, like a thread reply
+		// that also replies to another reply in the thread, so it's only kept
+		// and looked through the first time it's found
+		let mut seen = HashSet::new();
+		let mut relations = Vec::new();
+		let mut level: Vec<_> = self
+			.db
+			.get_relations(room_id, target, scan_from, dir)
+			.collect()
+			.await;
+		seen.extend(level.iter().map(|(count, _)| *count));
+		for _ in 0..max_depth {
+			let next: Vec<Vec<_>> = level
+				.iter()
+				.map(|(count, _)| *count)
+				.stream()
+				.wide_then(|target| {
+					self.db
+						.get_relations(room_id, target, scan_from, dir)
+						.collect::<Vec<_>>()
+				})
+				.collect()
+				.await;
+			relations.append(&mut level);
+			level = next
+				.into_iter()
+				.flatten()
+				.filter(|(count, _)| seen.insert(*count))
+				.collect();
+		}
+		relations.append(&mut level);
+
+		relations
+			.retain(|(count, pdu)| PduCount::Normal(*count).is_after(from, dir) && matches(pdu));
+		relations.sort_unstable_by(|(a, _), (b, _)| match dir {
+			| Direction::Forward => a.cmp(b),
+			| Direction::Backward => b.cmp(a),
+		});
+		relations.truncate(limit);
+
+		relations
+			.into_iter()
+			.stream()
+			.wide_then(|relation| self.finish_relation(user_id, relation))
+			.collect()
+			.await
+	}
+
+	/// The short room ID and count of `target`
+	async fn relation_target(
+		&self,
+		room_id: &RoomId,
+		target: &EventId,
+	) -> Option<(ShortRoomId, u64)> {
+		let room_id = self.services.short.get_shortroomid(room_id);
+		let target = self.services.timeline.get_pdu_count(target);
+
+		match try_join(room_id, target).await {
+			| Ok((room_id, PduCount::Normal(target))) => Some((room_id, target)),
+			// TODO: Support backfilled relations
+			| _ => None,
+		}
+	}
+
+	async fn finish_relation(
+		&self,
+		user_id: &UserId,
+		(count, pdu): (u64, PduEvent),
+	) -> PdusIterItem {
+		(PduCount::Normal(count), self.with_unsigned(user_id, pdu).await)
+	}
+
+	async fn with_unsigned(&self, user_id: &UserId, mut pdu: PduEvent) -> PduEvent {
+		let ctx = self
+			.services
+			.timeline
+			.get_unsigned_context(&pdu, Some(user_id))
+			.await;
+		pdu.set_unsigned(Some(ctx)).log_err().ok();
+
+		pdu
 	}
 
 	#[inline]