1use std::{collections::HashSet, pin::Pin};
86
87use async_stream::try_stream;
88use futures_util::{Stream, StreamExt as _, TryStreamExt as _, stream};
89use matrix_sdk_base::{RoomStateFilter, deserialized_responses::TimelineEvent};
90use matrix_sdk_search::error::IndexError;
91#[cfg(doc)]
92use matrix_sdk_search::index::RoomIndex;
93pub use matrix_sdk_search::index::SearchResult;
94use ruma::{MilliSecondsSinceUnixEpoch, OwnedEventId, OwnedRoomId};
95
96use crate::{Client, Room};
97
98const SEARCH_RESULTS_PAGE_SIZE: usize = 100;
101
102const GLOBAL_SEARCH_PAGE_SIZE: usize = 20;
106
107const MAX_CONCURRENT_EVENT_LOADS: usize = 10;
110
111type RoomResultStream = Pin<
114 Box<dyn Stream<Item = Result<(MilliSecondsSinceUnixEpoch, OwnedEventId), IndexError>> + Send>,
115>;
116
117struct RoomStreamCursor {
120 room_id: OwnedRoomId,
122
123 stream: RoomResultStream,
125
126 next_result: Option<(MilliSecondsSinceUnixEpoch, OwnedEventId)>,
131}
132
133impl Room {
134 pub async fn search(
137 &self,
138 query: &str,
139 max_number_of_results: usize,
140 pagination_offset: Option<usize>,
141 ) -> Result<SearchResult, IndexError> {
142 let mut search_index_guard = self.client.search_index().lock().await;
143 search_index_guard.search(query, max_number_of_results, pagination_offset, self.room_id())
144 }
145}
146
147#[derive(thiserror::Error, Debug)]
150pub enum SearchError {
151 #[error(transparent)]
153 IndexError(#[from] IndexError),
154 #[error(transparent)]
156 EventLoadError(#[from] crate::Error),
157}
158
159impl Room {
160 pub fn search_messages(
163 &self,
164 query: String,
165 ) -> impl Stream<Item = Result<SearchResult, IndexError>> + use<> {
166 let room = self.clone();
167
168 try_stream! {
171 let mut offset = 0;
172 loop {
173 let page = room.search(&query, SEARCH_RESULTS_PAGE_SIZE, Some(offset)).await?;
174 if page.events.is_empty() {
175 break;
176 }
177 offset += page.events.len();
178 yield page;
179 }
180 }
181 }
182
183 pub fn search_messages_events(
187 &self,
188 query: String,
189 ) -> impl Stream<Item = Result<Vec<TimelineEvent>, SearchError>> + use<> {
190 let room = self.clone();
191
192 try_stream! {
193 let mut pages = Box::pin(room.search_messages(query));
194
195 while let Some(page) = pages.next().await {
196 let page = page?;
197 let mut events = Vec::with_capacity(page.events.len());
198
199 for (_timestamp, event_id) in page.events {
200 events.push(room.load_or_fetch_event(&event_id, None).await?);
201 }
202
203 yield events;
204 }
205 }
206 }
207}
208
209#[derive(Debug)]
212pub struct GlobalSearchBuilder {
213 client: Client,
214
215 query: String,
217
218 room_set: Vec<Room>,
220}
221
222impl GlobalSearchBuilder {
223 fn new(client: Client, query: String) -> Self {
225 let room_set = client.rooms_filtered(RoomStateFilter::JOINED);
226 Self { client, query, room_set }
227 }
228
229 pub async fn only_dm_rooms(mut self) -> Result<Self, crate::Error> {
231 let mut to_remove = HashSet::new();
232 for room in &self.room_set {
233 if !room.compute_is_dm().await? {
234 to_remove.insert(room.room_id().to_owned());
235 }
236 }
237 self.room_set.retain(|room| !to_remove.contains(room.room_id()));
238 Ok(self)
239 }
240
241 pub async fn no_dms(mut self) -> Result<Self, crate::Error> {
243 let mut to_remove = HashSet::new();
244 for room in &self.room_set {
245 if room.compute_is_dm().await? {
246 to_remove.insert(room.room_id().to_owned());
247 }
248 }
249 self.room_set.retain(|room| !to_remove.contains(room.room_id()));
250 Ok(self)
251 }
252
253 pub fn build(
257 self,
258 ) -> impl Stream<
259 Item = Result<Vec<(OwnedRoomId, MilliSecondsSinceUnixEpoch, OwnedEventId)>, IndexError>,
260 > {
261 let client = self.client;
262 let query = self.query;
263 let rooms = self.room_set;
264
265 try_stream! {
266 client
267 .search_index()
268 .open_indexes(rooms.iter().map(|room| room.room_id().to_owned()))
269 .await;
270
271 let mut cursors: Vec<RoomStreamCursor> = Vec::with_capacity(rooms.len());
274 for room in rooms {
275 let room_id = room.room_id().to_owned();
276
277 let stream = Box::pin(Self::flatten_pages(room.search_messages(query.clone())));
278 cursors.push(RoomStreamCursor { room_id, stream, next_result: None });
279 }
280
281 for cursor in &mut cursors {
282 cursor.next_result = match cursor.stream.next().await {
283 Some(result) => Some(result?),
284 None => None,
285 };
286 }
287
288 let mut page = Vec::with_capacity(GLOBAL_SEARCH_PAGE_SIZE);
289
290 loop {
291 let best = cursors
293 .iter()
294 .enumerate()
295 .filter_map(|(index, cursor)| {
296 cursor.next_result.as_ref().map(|(timestamp, _)| (index, *timestamp))
297 })
298 .max_by_key(|(_, timestamp)| *timestamp);
299
300 let Some((index, _)) = best else {
301 break;
303 };
304
305 let cursor = &mut cursors[index];
306 let (timestamp, event_id) =
307 cursor.next_result.take().expect("the chosen room must have a next result");
308 let room_id = cursor.room_id.clone();
309
310 cursor.next_result = match cursor.stream.next().await {
312 Some(result) => Some(result?),
313 None => None,
314 };
315
316 page.push((room_id, timestamp, event_id));
317 if page.len() == GLOBAL_SEARCH_PAGE_SIZE {
318 yield std::mem::take(&mut page);
319 }
320 }
321
322 if !page.is_empty() {
323 yield page;
324 }
325 }
326 }
327
328 pub fn build_events(
332 self,
333 ) -> impl Stream<Item = Result<Vec<(OwnedRoomId, TimelineEvent)>, SearchError>> {
334 let client = self.client.clone();
335 let pages = self.build();
336
337 try_stream! {
338 let mut pages = Box::pin(pages);
339
340 while let Some(page) = pages.next().await {
341 let loads = page?.into_iter().filter_map(|(room_id, _timestamp, event_id)| {
344 let room = client.get_room(&room_id)?;
345 Some(async move {
346 let event = room.load_or_fetch_event(&event_id, None).await?;
347 Ok::<_, crate::Error>((room_id, event))
348 })
349 });
350
351 let events: Vec<_> =
352 stream::iter(loads).buffered(MAX_CONCURRENT_EVENT_LOADS).try_collect().await?;
353
354 yield events;
355 }
356 }
357 }
358
359 fn flatten_pages(
362 pages: impl Stream<Item = Result<SearchResult, IndexError>>,
363 ) -> impl Stream<Item = Result<(MilliSecondsSinceUnixEpoch, OwnedEventId), IndexError>> {
364 try_stream! {
365 let mut pages = Box::pin(pages);
366
367 while let Some(page) = pages.next().await {
368 for result in page?.events {
369 yield result;
370 }
371 }
372 }
373 }
374}
375
376impl Client {
377 pub fn search_messages(&self, query: String) -> GlobalSearchBuilder {
380 GlobalSearchBuilder::new(self.clone(), query)
381 }
382}
383
384#[cfg(test)]
385mod tests {
386 use std::time::Duration;
387
388 use futures_util::TryStreamExt as _;
389 use matrix_sdk_search::index::SearchResult;
390 use matrix_sdk_test::{JoinedRoomBuilder, User, async_test, event_factory::EventFactory};
391 use ruma::{MilliSecondsSinceUnixEpoch, OwnedEventId, OwnedRoomId, event_id, room_id, user_id};
392
393 use super::GLOBAL_SEARCH_PAGE_SIZE;
394 use crate::{sleep::sleep, test_utils::mocks::MatrixMockServer};
395
396 #[async_test]
397 async fn test_room_message_search() {
398 let server = MatrixMockServer::new().await;
399 let client = server.client_builder().build().await;
400
401 let room_id = room_id!("!room_id:localhost");
402 let room = server.sync_joined_room(&client, room_id).await;
403
404 let f = EventFactory::new().room(room_id).sender(user_id!("@user_id:localhost"));
405
406 let event_id = event_id!("$event_id:localhost");
407
408 server
409 .sync_room(
410 &client,
411 JoinedRoomBuilder::new(room_id)
412 .add_timeline_event(f.text_msg("hello world").event_id(event_id)),
413 )
414 .await;
415
416 sleep(Duration::from_millis(200)).await;
418
419 {
421 let results: Vec<SearchResult> =
422 room.search_messages("search query".to_owned()).try_collect().await.unwrap();
423 assert!(results.is_empty());
424 }
425
426 {
428 let results: Vec<SearchResult> =
429 room.search_messages("world".to_owned()).try_collect().await.unwrap();
430 assert_eq!(results[0].events.len(), 1);
431 assert_eq!(results[0].events[0].1, event_id);
432 }
433
434 {
436 let results: Vec<_> =
437 room.search_messages_events("world".to_owned()).try_concat().await.unwrap();
438 assert_eq!(results.len(), 1);
439 assert_eq!(results[0].event_id().unwrap(), event_id);
440 }
441 }
442
443 #[async_test]
444 async fn test_global_message_search() {
445 let server = MatrixMockServer::new().await;
446 let client = server.client_builder().build().await;
447
448 let room_id1 = room_id!("!r1:localhost");
449 let room_id2 = room_id!("!r2:localhost");
450
451 let f = EventFactory::new().sender(user_id!("@user_id:localhost"));
452
453 let result_event_id1 = event_id!("$result1:localhost");
454 let result_event_id2 = event_id!("$result2:localhost");
455
456 server
457 .mock_sync()
458 .ok_and_run(&client, |sync_builder| {
459 sync_builder
460 .add_joined_room(
461 JoinedRoomBuilder::new(room_id1)
462 .add_timeline_event(
463 f.text_msg("hello world").room(room_id1).event_id(result_event_id1),
464 )
465 .add_timeline_event(f.text_msg("hello back").room(room_id1)),
466 )
467 .add_joined_room(JoinedRoomBuilder::new(room_id2).add_timeline_event(
468 f.text_msg("it's a mad world").room(room_id2).event_id(result_event_id2),
469 ));
470 })
471 .await;
472
473 sleep(Duration::from_millis(200)).await;
475
476 {
478 let results: Vec<(OwnedRoomId, MilliSecondsSinceUnixEpoch, OwnedEventId)> = client
479 .search_messages("search query".to_owned())
480 .build()
481 .try_concat()
482 .await
483 .unwrap();
484 assert!(results.is_empty());
485 }
486
487 {
489 let results: Vec<(OwnedRoomId, MilliSecondsSinceUnixEpoch, OwnedEventId)> =
490 client.search_messages("world".to_owned()).build().try_concat().await.unwrap();
491 assert_eq!(results.len(), 2);
492 assert!(results.iter().any(|(room_id, _, event_id)| {
495 room_id == room_id1 && event_id == result_event_id1
496 }));
497 assert!(results.iter().any(|(room_id, _, event_id)| {
498 room_id == room_id2 && event_id == result_event_id2
499 }));
500 }
501
502 {
504 let results: Vec<_> = client
505 .search_messages("world".to_owned())
506 .build_events()
507 .try_concat()
508 .await
509 .unwrap();
510 assert_eq!(results.len(), 2);
511 assert!(results.iter().any(|(room_id, event)| {
514 room_id == room_id1 && event.event_id() == Some(result_event_id1)
515 }));
516 assert!(results.iter().any(|(room_id, event)| {
517 room_id == room_id2 && event.event_id() == Some(result_event_id2)
518 }));
519 }
520 }
521
522 #[async_test]
523 async fn test_global_message_search_timestamp_ordering() {
524 let server = MatrixMockServer::new().await;
525 let client = server.client_builder().build().await;
526
527 let room_id1 = room_id!("!r1:localhost");
528 let room_id2 = room_id!("!r2:localhost");
529
530 let f = EventFactory::new().sender(user_id!("@user_id:localhost"));
531
532 let r1_rank1 = event_id!("$r1_rank1:localhost");
539 let r2_rank2 = event_id!("$r2_rank2:localhost");
540 let r1_rank3 = event_id!("$r1_rank3:localhost");
541 let r2_rank4 = event_id!("$r2_rank4:localhost");
542
543 server
544 .mock_sync()
545 .ok_and_run(&client, |sync_builder| {
546 sync_builder
547 .add_joined_room(
548 JoinedRoomBuilder::new(room_id1)
549 .add_timeline_event(
550 f.text_msg("world world")
551 .room(room_id1)
552 .event_id(r1_rank3)
553 .server_ts(2_000),
554 )
555 .add_timeline_event(
556 f.text_msg("world")
557 .room(room_id1)
558 .event_id(r1_rank1)
559 .server_ts(4_000),
560 ),
561 )
562 .add_joined_room(
563 JoinedRoomBuilder::new(room_id2)
564 .add_timeline_event(
565 f.text_msg("world world world world")
566 .room(room_id2)
567 .event_id(r2_rank4)
568 .server_ts(1_000),
569 )
570 .add_timeline_event(
571 f.text_msg("world")
572 .room(room_id2)
573 .event_id(r2_rank2)
574 .server_ts(3_000),
575 ),
576 );
577 })
578 .await;
579
580 sleep(Duration::from_millis(200)).await;
581
582 let results: Vec<(OwnedRoomId, MilliSecondsSinceUnixEpoch, OwnedEventId)> =
583 client.search_messages("world".to_owned()).build().try_concat().await.unwrap();
584 assert_eq!(results.len(), 4);
585
586 assert_eq!((&results[0].0, &results[0].2), (&room_id1.to_owned(), &r1_rank1.to_owned()));
589 assert_eq!((&results[1].0, &results[1].2), (&room_id2.to_owned(), &r2_rank2.to_owned()));
590 assert_eq!((&results[2].0, &results[2].2), (&room_id1.to_owned(), &r1_rank3.to_owned()));
591 assert_eq!((&results[3].0, &results[3].2), (&room_id2.to_owned(), &r2_rank4.to_owned()));
592 }
593
594 #[async_test]
595 async fn test_global_message_search_events_pages() {
596 let server = MatrixMockServer::new().await;
597 let client = server.client_builder().build().await;
598
599 let room_id = room_id!("!r1:localhost");
600 let f = EventFactory::new().room(room_id).sender(user_id!("@user_id:localhost"));
601
602 let total = GLOBAL_SEARCH_PAGE_SIZE + 5;
603 let event_ids: Vec<OwnedEventId> =
604 (0..total).map(|i| format!("$event{i}:localhost").try_into().unwrap()).collect();
605
606 server
607 .sync_room(
608 &client,
609 event_ids.iter().enumerate().fold(
610 JoinedRoomBuilder::new(room_id),
611 |builder, (i, event_id)| {
612 let ts = (i as u64 + 1) * 1_000;
614 builder.add_timeline_event(
615 f.text_msg("world").event_id(event_id).server_ts(ts),
616 )
617 },
618 ),
619 )
620 .await;
621
622 sleep(Duration::from_millis(200)).await;
623
624 let pages: Vec<Vec<(OwnedRoomId, _)>> =
625 client.search_messages("world".to_owned()).build_events().try_collect().await.unwrap();
626
627 assert_eq!(pages.iter().map(Vec::len).collect::<Vec<_>>(), [GLOBAL_SEARCH_PAGE_SIZE, 5]);
630 let loaded: Vec<_> =
631 pages.iter().flatten().map(|(_, event)| event.event_id().unwrap()).collect();
632 let expected: Vec<_> = event_ids.into_iter().rev().collect();
633 assert_eq!(loaded, expected);
634 }
635
636 #[async_test]
637 async fn test_global_message_search_indexes_large_syncs() {
638 let server = MatrixMockServer::new().await;
639 let client = server.client_builder().build().await;
640
641 let f = EventFactory::new().sender(user_id!("@user_id:localhost"));
642
643 let room_ids: Vec<OwnedRoomId> =
647 (0..300).map(|i| format!("!room{i}:localhost").try_into().unwrap()).collect();
648
649 let index_guard = client.search_index().lock().await;
650
651 server
652 .mock_sync()
653 .ok_and_run(&client, |sync_builder| {
654 for room_id in &room_ids {
655 sync_builder.add_joined_room(
656 JoinedRoomBuilder::new(room_id)
657 .add_timeline_event(f.text_msg("hello world").room(room_id)),
658 );
659 }
660 })
661 .await;
662
663 drop(index_guard);
664
665 let mut indexed = 0;
667 for _ in 0..100 {
668 let results: Vec<(OwnedRoomId, MilliSecondsSinceUnixEpoch, OwnedEventId)> =
669 client.search_messages("world".to_owned()).build().try_concat().await.unwrap();
670 indexed = results.len();
671 if indexed == room_ids.len() {
672 break;
673 }
674 sleep(Duration::from_millis(100)).await;
675 }
676 assert_eq!(indexed, room_ids.len());
677 }
678
679 #[async_test]
680 async fn test_global_message_search_dm_or_groups() {
681 let server = MatrixMockServer::new().await;
682 let client = server.client_builder().build().await;
683
684 let room_id1 = room_id!("!r1:localhost");
686 let room_id2 = room_id!("!r2:localhost");
688
689 let f = EventFactory::new().sender(user_id!("@user_id:localhost"));
690
691 let result_event_id1 = event_id!("$result1:localhost");
692 let result_event_id2 = event_id!("$result2:localhost");
693
694 server
695 .mock_sync()
696 .ok_and_run(&client, |sync_builder| {
697 sync_builder
698 .add_joined_room(
699 JoinedRoomBuilder::new(room_id1)
700 .add_timeline_event(
701 f.text_msg("hello world").room(room_id1).event_id(result_event_id1),
702 )
703 .add_timeline_event(f.text_msg("hello back").room(room_id1)),
704 )
705 .add_joined_room(JoinedRoomBuilder::new(room_id2).add_timeline_event(
706 f.text_msg("it's a mad world").room(room_id2).event_id(result_event_id2),
707 ))
708 .add_global_account_data(f.direct().add_user(User::Bob, room_id1));
710 })
711 .await;
712
713 sleep(Duration::from_millis(200)).await;
715
716 {
718 let results: Vec<(OwnedRoomId, MilliSecondsSinceUnixEpoch, OwnedEventId)> = client
719 .search_messages("world".to_owned())
720 .only_dm_rooms()
721 .await
722 .unwrap()
723 .build()
724 .try_concat()
725 .await
726 .unwrap();
727
728 assert_eq!(results.len(), 1);
729 assert_eq!(
730 (&results[0].0, &results[0].2),
731 (&room_id1.to_owned(), &result_event_id1.to_owned())
732 );
733 }
734
735 {
737 let results: Vec<_> = client
738 .search_messages("world".to_owned())
739 .no_dms()
740 .await
741 .unwrap()
742 .build_events()
743 .try_concat()
744 .await
745 .unwrap();
746
747 assert_eq!(results.len(), 1);
748 assert_eq!(results[0].0, room_id2);
749 assert_eq!(results[0].1.event_id().unwrap(), result_event_id2);
750 }
751 }
752}