1use std::{collections::HashSet, pin::Pin};
84
85use async_stream::try_stream;
86use futures_util::{Stream, StreamExt as _};
87use matrix_sdk_base::{RoomStateFilter, deserialized_responses::TimelineEvent};
88use matrix_sdk_search::error::IndexError;
89#[cfg(doc)]
90use matrix_sdk_search::index::RoomIndex;
91use ruma::{OwnedEventId, OwnedRoomId};
92
93use crate::{Client, Room};
94
95const SEARCH_RESULTS_PAGE_SIZE: usize = 100;
98
99type RoomResultStream = Pin<Box<dyn Stream<Item = Result<(f32, OwnedEventId), IndexError>> + Send>>;
102
103struct RoomStreamCursor {
106 room_id: OwnedRoomId,
108
109 stream: RoomResultStream,
111
112 next_result: Option<(f32, OwnedEventId)>,
117}
118
119impl Room {
120 pub async fn search(
123 &self,
124 query: &str,
125 max_number_of_results: usize,
126 pagination_offset: Option<usize>,
127 ) -> Result<Vec<(f32, OwnedEventId)>, IndexError> {
128 let mut search_index_guard = self.client.search_index().lock().await;
129 search_index_guard.search(query, max_number_of_results, pagination_offset, self.room_id())
130 }
131}
132
133#[derive(thiserror::Error, Debug)]
136pub enum SearchError {
137 #[error(transparent)]
139 IndexError(#[from] IndexError),
140 #[error(transparent)]
142 EventLoadError(#[from] crate::Error),
143}
144
145impl Room {
146 pub fn search_messages(
150 &self,
151 query: String,
152 ) -> impl Stream<Item = Result<Vec<(f32, OwnedEventId)>, IndexError>> + use<> {
153 let room = self.clone();
154
155 try_stream! {
158 let mut offset = 0;
159 loop {
160 let page = room.search(&query, SEARCH_RESULTS_PAGE_SIZE, Some(offset)).await?;
161 if page.is_empty() {
162 break;
163 }
164 offset += page.len();
165 yield page;
166 }
167 }
168 }
169
170 pub fn search_messages_events(
174 &self,
175 query: String,
176 ) -> impl Stream<Item = Result<Vec<TimelineEvent>, SearchError>> + use<> {
177 let room = self.clone();
178
179 try_stream! {
180 let mut pages = Box::pin(room.search_messages(query));
181 while let Some(page) = pages.next().await {
182 let page = page?;
183 let mut events = Vec::with_capacity(page.len());
184 for (_score, event_id) in page {
185 events.push(room.load_or_fetch_event(&event_id, None).await?);
186 }
187 yield events;
188 }
189 }
190 }
191}
192
193#[derive(Debug)]
196pub struct GlobalSearchBuilder {
197 client: Client,
198
199 query: String,
201
202 room_set: Vec<Room>,
204}
205
206impl GlobalSearchBuilder {
207 fn new(client: Client, query: String) -> Self {
209 let room_set = client.rooms_filtered(RoomStateFilter::JOINED);
210 Self { client, query, room_set }
211 }
212
213 pub async fn only_dm_rooms(mut self) -> Result<Self, crate::Error> {
215 let mut to_remove = HashSet::new();
216 for room in &self.room_set {
217 if !room.compute_is_dm().await? {
218 to_remove.insert(room.room_id().to_owned());
219 }
220 }
221 self.room_set.retain(|room| !to_remove.contains(room.room_id()));
222 Ok(self)
223 }
224
225 pub async fn no_dms(mut self) -> Result<Self, crate::Error> {
227 let mut to_remove = HashSet::new();
228 for room in &self.room_set {
229 if room.compute_is_dm().await? {
230 to_remove.insert(room.room_id().to_owned());
231 }
232 }
233 self.room_set.retain(|room| !to_remove.contains(room.room_id()));
234 Ok(self)
235 }
236
237 pub fn build(
241 self,
242 ) -> impl Stream<Item = Result<Vec<(OwnedRoomId, f32, OwnedEventId)>, IndexError>> {
243 let query = self.query;
244 let rooms = self.room_set;
245
246 try_stream! {
247 let mut cursors: Vec<RoomStreamCursor> = Vec::with_capacity(rooms.len());
250 for room in rooms {
251 let room_id = room.room_id().to_owned();
252 let stream = Box::pin(Self::flatten_pages(room.search_messages(query.clone())));
253 cursors.push(RoomStreamCursor { room_id, stream, next_result: None });
254 }
255 for cursor in &mut cursors {
256 cursor.next_result = match cursor.stream.next().await {
257 Some(result) => Some(result?),
258 None => None,
259 };
260 }
261
262 let mut page = Vec::with_capacity(SEARCH_RESULTS_PAGE_SIZE);
263 loop {
264 let best = cursors
266 .iter()
267 .enumerate()
268 .filter_map(|(index, cursor)| {
269 cursor.next_result.as_ref().map(|(score, _)| (index, *score))
270 })
271 .max_by(|(_, a), (_, b)| a.total_cmp(b));
272
273 let Some((index, _)) = best else {
274 break;
276 };
277
278 let cursor = &mut cursors[index];
279 let (score, event_id) =
280 cursor.next_result.take().expect("the chosen room must have a next result");
281 let room_id = cursor.room_id.clone();
282
283 cursor.next_result = match cursor.stream.next().await {
285 Some(result) => Some(result?),
286 None => None,
287 };
288
289 page.push((room_id, score, event_id));
290 if page.len() == SEARCH_RESULTS_PAGE_SIZE {
291 yield std::mem::take(&mut page);
292 }
293 }
294
295 if !page.is_empty() {
296 yield page;
297 }
298 }
299 }
300
301 pub fn build_events(
305 self,
306 ) -> impl Stream<Item = Result<Vec<(OwnedRoomId, TimelineEvent)>, SearchError>> {
307 let client = self.client.clone();
308 let pages = self.build();
309
310 try_stream! {
311 let mut pages = Box::pin(pages);
312 while let Some(page) = pages.next().await {
313 let page = page?;
314 let mut events = Vec::with_capacity(page.len());
315 for (room_id, _score, event_id) in page {
316 let Some(room) = client.get_room(&room_id) else {
317 continue;
318 };
319 events.push((room_id, room.load_or_fetch_event(&event_id, None).await?));
320 }
321 yield events;
322 }
323 }
324 }
325
326 fn flatten_pages(
329 pages: impl Stream<Item = Result<Vec<(f32, OwnedEventId)>, IndexError>>,
330 ) -> impl Stream<Item = Result<(f32, OwnedEventId), IndexError>> {
331 try_stream! {
332 let mut pages = Box::pin(pages);
333 while let Some(page) = pages.next().await {
334 for result in page? {
335 yield result;
336 }
337 }
338 }
339 }
340}
341
342impl Client {
343 pub fn search_messages(&self, query: String) -> GlobalSearchBuilder {
346 GlobalSearchBuilder::new(self.clone(), query)
347 }
348}
349
350#[cfg(test)]
351mod tests {
352 use std::time::Duration;
353
354 use futures_util::TryStreamExt as _;
355 use matrix_sdk_test::{BOB, JoinedRoomBuilder, async_test, event_factory::EventFactory};
356 use ruma::{OwnedEventId, OwnedRoomId, event_id, room_id, user_id};
357
358 use crate::{sleep::sleep, test_utils::mocks::MatrixMockServer};
359
360 #[async_test]
361 async fn test_room_message_search() {
362 let server = MatrixMockServer::new().await;
363 let client = server.client_builder().build().await;
364
365 let event_cache = client.event_cache();
366 event_cache.subscribe().unwrap();
367
368 let room_id = room_id!("!room_id:localhost");
369 let room = server.sync_joined_room(&client, room_id).await;
370
371 let f = EventFactory::new().room(room_id).sender(user_id!("@user_id:localhost"));
372
373 let event_id = event_id!("$event_id:localhost");
374
375 server
376 .sync_room(
377 &client,
378 JoinedRoomBuilder::new(room_id)
379 .add_timeline_event(f.text_msg("hello world").event_id(event_id)),
380 )
381 .await;
382
383 sleep(Duration::from_millis(200)).await;
385
386 {
388 let results: Vec<(f32, OwnedEventId)> =
389 room.search_messages("search query".to_owned()).try_concat().await.unwrap();
390 assert!(results.is_empty());
391 }
392
393 {
395 let results: Vec<(f32, OwnedEventId)> =
396 room.search_messages("world".to_owned()).try_concat().await.unwrap();
397 assert_eq!(results.len(), 1);
398 assert_eq!(results[0].1, event_id);
399 }
400
401 {
403 let events: Vec<_> =
404 room.search_messages_events("world".to_owned()).try_concat().await.unwrap();
405 assert_eq!(events.len(), 1);
406 assert_eq!(events[0].event_id().unwrap(), event_id);
407 }
408 }
409
410 #[async_test]
411 async fn test_global_message_search() {
412 let server = MatrixMockServer::new().await;
413 let client = server.client_builder().build().await;
414
415 let event_cache = client.event_cache();
416 event_cache.subscribe().unwrap();
417
418 let room_id1 = room_id!("!r1:localhost");
419 let room_id2 = room_id!("!r2:localhost");
420
421 let f = EventFactory::new().sender(user_id!("@user_id:localhost"));
422
423 let result_event_id1 = event_id!("$result1:localhost");
424 let result_event_id2 = event_id!("$result2:localhost");
425
426 server
427 .mock_sync()
428 .ok_and_run(&client, |sync_builder| {
429 sync_builder
430 .add_joined_room(
431 JoinedRoomBuilder::new(room_id1)
432 .add_timeline_event(
433 f.text_msg("hello world").room(room_id1).event_id(result_event_id1),
434 )
435 .add_timeline_event(f.text_msg("hello back").room(room_id1)),
436 )
437 .add_joined_room(JoinedRoomBuilder::new(room_id2).add_timeline_event(
438 f.text_msg("it's a mad world").room(room_id2).event_id(result_event_id2),
439 ));
440 })
441 .await;
442
443 sleep(Duration::from_millis(200)).await;
445
446 {
448 let results: Vec<(OwnedRoomId, f32, OwnedEventId)> = client
449 .search_messages("search query".to_owned())
450 .build()
451 .try_concat()
452 .await
453 .unwrap();
454 assert!(results.is_empty());
455 }
456
457 {
459 let results: Vec<(OwnedRoomId, f32, OwnedEventId)> =
460 client.search_messages("world".to_owned()).build().try_concat().await.unwrap();
461 assert_eq!(results.len(), 2);
462 assert!(results.iter().any(|(room_id, _, event_id)| {
465 room_id == room_id1 && event_id == result_event_id1
466 }));
467 assert!(results.iter().any(|(room_id, _, event_id)| {
468 room_id == room_id2 && event_id == result_event_id2
469 }));
470 }
471
472 {
474 let results: Vec<_> = client
475 .search_messages("world".to_owned())
476 .build_events()
477 .try_concat()
478 .await
479 .unwrap();
480 assert_eq!(results.len(), 2);
481 assert!(results.iter().any(|(room_id, event)| {
484 room_id == room_id1 && event.event_id() == Some(result_event_id1)
485 }));
486 assert!(results.iter().any(|(room_id, event)| {
487 room_id == room_id2 && event.event_id() == Some(result_event_id2)
488 }));
489 }
490 }
491
492 #[async_test]
493 async fn test_global_message_search_score_ordering() {
494 let server = MatrixMockServer::new().await;
495 let client = server.client_builder().build().await;
496
497 let event_cache = client.event_cache();
498 event_cache.subscribe().unwrap();
499
500 let room_id1 = room_id!("!r1:localhost");
501 let room_id2 = room_id!("!r2:localhost");
502
503 let f = EventFactory::new().sender(user_id!("@user_id:localhost"));
504
505 let r1_rank1 = event_id!("$r1_rank1:localhost"); let r2_rank2 = event_id!("$r2_rank2:localhost"); let r1_rank3 = event_id!("$r1_rank3:localhost"); let r2_rank4 = event_id!("$r2_rank4:localhost"); server
518 .mock_sync()
519 .ok_and_run(&client, |sync_builder| {
520 sync_builder
521 .add_joined_room(
522 JoinedRoomBuilder::new(room_id1)
523 .add_timeline_event(
524 f.text_msg("world world world world filler filler filler filler filler filler")
525 .room(room_id1)
526 .event_id(r1_rank1),
527 )
528 .add_timeline_event(
529 f.text_msg("world world filler filler filler filler filler filler filler filler")
530 .room(room_id1)
531 .event_id(r1_rank3),
532 ),
533 )
534 .add_joined_room(
535 JoinedRoomBuilder::new(room_id2)
536 .add_timeline_event(
537 f.text_msg("world world world filler filler filler filler filler filler filler")
538 .room(room_id2)
539 .event_id(r2_rank2),
540 )
541 .add_timeline_event(
542 f.text_msg("world filler filler filler filler filler filler filler filler filler")
543 .room(room_id2)
544 .event_id(r2_rank4),
545 ),
546 );
547 })
548 .await;
549
550 sleep(Duration::from_millis(200)).await;
551
552 let results: Vec<(OwnedRoomId, f32, OwnedEventId)> =
553 client.search_messages("world".to_owned()).build().try_concat().await.unwrap();
554 assert_eq!(results.len(), 4);
555
556 assert_eq!((&results[0].0, &results[0].2), (&room_id1.to_owned(), &r1_rank1.to_owned()));
558 assert_eq!((&results[1].0, &results[1].2), (&room_id2.to_owned(), &r2_rank2.to_owned()));
559 assert_eq!((&results[2].0, &results[2].2), (&room_id1.to_owned(), &r1_rank3.to_owned()));
560 assert_eq!((&results[3].0, &results[3].2), (&room_id2.to_owned(), &r2_rank4.to_owned()));
561 }
562
563 #[async_test]
564 async fn test_global_message_search_dm_or_groups() {
565 let server = MatrixMockServer::new().await;
566 let client = server.client_builder().build().await;
567
568 let event_cache = client.event_cache();
569 event_cache.subscribe().unwrap();
570
571 let room_id1 = room_id!("!r1:localhost");
573 let room_id2 = room_id!("!r2:localhost");
575
576 let f = EventFactory::new().sender(user_id!("@user_id:localhost"));
577
578 let result_event_id1 = event_id!("$result1:localhost");
579 let result_event_id2 = event_id!("$result2:localhost");
580
581 server
582 .mock_sync()
583 .ok_and_run(&client, |sync_builder| {
584 sync_builder
585 .add_joined_room(
586 JoinedRoomBuilder::new(room_id1)
587 .add_timeline_event(
588 f.text_msg("hello world").room(room_id1).event_id(result_event_id1),
589 )
590 .add_timeline_event(f.text_msg("hello back").room(room_id1)),
591 )
592 .add_joined_room(JoinedRoomBuilder::new(room_id2).add_timeline_event(
593 f.text_msg("it's a mad world").room(room_id2).event_id(result_event_id2),
594 ))
595 .add_global_account_data(
597 f.direct().add_user((*BOB).to_owned().into(), room_id1),
598 );
599 })
600 .await;
601
602 sleep(Duration::from_millis(200)).await;
604
605 {
607 let results: Vec<(OwnedRoomId, f32, OwnedEventId)> = client
608 .search_messages("world".to_owned())
609 .only_dm_rooms()
610 .await
611 .unwrap()
612 .build()
613 .try_concat()
614 .await
615 .unwrap();
616
617 assert_eq!(results.len(), 1);
618 assert_eq!(
619 (&results[0].0, &results[0].2),
620 (&room_id1.to_owned(), &result_event_id1.to_owned())
621 );
622 }
623
624 {
626 let results: Vec<_> = client
627 .search_messages("world".to_owned())
628 .no_dms()
629 .await
630 .unwrap()
631 .build_events()
632 .try_concat()
633 .await
634 .unwrap();
635
636 assert_eq!(results.len(), 1);
637 assert_eq!(results[0].0, room_id2);
638 assert_eq!(results[0].1.event_id().unwrap(), result_event_id2);
639 }
640 }
641}