diff --git a/src/service/rooms/search/mod.rs b/src/service/rooms/search/mod.rs index 241d952ff..71ed82610 100644 --- a/src/service/rooms/search/mod.rs +++ b/src/service/rooms/search/mod.rs @@ -3,7 +3,7 @@ use conduwuit::{ PduCount, PduEvent, Result, arrayvec::ArrayVec, - debug_warn, implement, + debug_warn, matrix::event::{Event, Matches}, utils::{ ArrayVecExt, IterStream, ReadyExt, set, @@ -70,158 +70,161 @@ fn build(args: crate::Args<'_>) -> Result> { fn name(&self) -> &str { crate::service::make_name(std::module_path!()) } } -#[implement(Service)] -pub fn index_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) { - let batch = tokenize(message_body) - .map(|word| { +impl Service { + /// Adds a single PDU to the search index. + pub fn index_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) { + let batch = tokenize(message_body) + .map(|word| { + let mut key = shortroomid.to_be_bytes().to_vec(); + key.extend_from_slice(word.as_bytes()); + key.push(0xFF); + key.extend_from_slice(pdu_id.as_ref()); // TODO: currently we save the room id a second time here + key + }) + .collect::>(); + + self.db + .tokenids + .insert_batch(batch.iter().map(|k| (k.as_slice(), &[]))); + } + + /// Removes a single PDU from the search index. + pub fn deindex_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) { + let batch = tokenize(message_body).map(|word| { let mut key = shortroomid.to_be_bytes().to_vec(); key.extend_from_slice(word.as_bytes()); key.push(0xFF); key.extend_from_slice(pdu_id.as_ref()); // TODO: currently we save the room id a second time here key - }) - .collect::>(); - - self.db - .tokenids - .insert_batch(batch.iter().map(|k| (k.as_slice(), &[]))); -} - -#[implement(Service)] -pub fn deindex_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) { - let batch = tokenize(message_body).map(|word| { - let mut key = shortroomid.to_be_bytes().to_vec(); - key.extend_from_slice(word.as_bytes()); - key.push(0xFF); - key.extend_from_slice(pdu_id.as_ref()); // TODO: currently we save the room id a second time here - key - }); - - for token in batch { - self.db.tokenids.remove(&token); - } -} - -#[implement(Service)] -pub async fn search_pdus<'a>( - &'a self, - query: &'a RoomQuery<'a>, - sender_user: &'a UserId, -) -> Result<(usize, impl Stream + Send + 'a)> { - let pdu_ids: Vec<_> = self.search_pdu_ids(query).await?.collect().await; - - let filter = &query.criteria.filter; - let count = pdu_ids.len(); - let pdus = pdu_ids - .into_iter() - .stream() - .wide_filter_map(move |result_pdu_id: RawPduId| async move { - self.services - .timeline - .get_pdu_from_id(&result_pdu_id) - .await - .ok() - }) - .ready_filter(|pdu| !pdu.is_redacted()) - .ready_filter(move |pdu| filter.matches(pdu)) - .wide_filter_map(move |pdu| async move { - self.services - .state_accessor - .user_can_see_event(query.user_id?, pdu.room_id().unwrap(), pdu.event_id()) - .await - .then_some(pdu) - }) - .skip(query.skip) - .take(query.limit) - .map(move |mut pdu| { - pdu.set_unsigned(query.user_id); - - pdu - }) - .then(async move |mut pdu| { - if let Err(e) = self - .services - .pdu_metadata - .add_bundled_aggregations_to_pdu(sender_user, &mut pdu) - .await - { - debug_warn!("Failed to add bundled aggregations: {e}"); - } - pdu }); - Ok((count, pdus)) -} - -// result is modeled as a stream such that callers don't have to be refactored -// though an additional async/wrap still exists for now -#[implement(Service)] -pub async fn search_pdu_ids<'a>( - &'a self, - query: &'a RoomQuery<'_>, -) -> Result + Send + 'a + use<'a>> { - let shortroomid = self.services.short.get_shortroomid(query.room_id).await?; - - let pdu_ids = self.search_pdu_ids_query_room(query, shortroomid).await; - - let iters = pdu_ids.into_iter().map(IntoIterator::into_iter); - - Ok(set::intersection(iters).stream()) -} - -#[implement(Service)] -async fn search_pdu_ids_query_room( - &self, - query: &RoomQuery<'_>, - shortroomid: ShortRoomId, -) -> Vec> { - tokenize(&query.criteria.search_term) - .stream() - .wide_then(|word| async move { - self.search_pdu_ids_query_words(shortroomid, &word) - .collect::>() - .await - }) - .collect::>() - .await -} - -/// Iterate over PduId's containing a word -#[implement(Service)] -fn search_pdu_ids_query_words<'a>( - &'a self, - shortroomid: ShortRoomId, - word: &'a str, -) -> impl Stream + Send + 'a { - self.search_pdu_ids_query_word(shortroomid, word) - .map(move |key| -> RawPduId { - let key = &key[prefix_len(word)..]; - key.into() - }) -} - -/// Iterate over raw database results for a word -#[implement(Service)] -fn search_pdu_ids_query_word<'a>( - &'a self, - shortroomid: ShortRoomId, - word: &'a str, -) -> impl Stream> + Send + 'a + use<'a> { - // rustc says const'ing this not yet stable - let end_id: RawPduId = PduId { - shortroomid, - shorteventid: PduCount::max(), + for token in batch { + self.db.tokenids.remove(&token); + } } - .into(); - // Newest pdus first - let end = make_tokenid(shortroomid, word, &end_id); - let prefix = make_prefix(shortroomid, word); - self.db - .tokenids - .rev_raw_keys_from(&end) - .ignore_err() - .ready_take_while(move |key| key.starts_with(&prefix)) + /// Searches through the search index to find PDUs matching the query. + /// Filters history visibility based on what the sender can see, and bundles + /// aggregations before returning. + pub async fn search_pdus<'a>( + &'a self, + query: &'a RoomQuery<'a>, + sender_user: &'a UserId, + ) -> Result<(usize, impl Stream + Send + 'a)> { + let pdu_ids: Vec<_> = self.search_pdu_ids(query).await?.collect().await; + + let filter = &query.criteria.filter; + let count = pdu_ids.len(); + let pdus = pdu_ids + .into_iter() + .stream() + .wide_filter_map(move |result_pdu_id: RawPduId| async move { + self.services + .timeline + .get_pdu_from_id(&result_pdu_id) + .await + .ok() + }) + .ready_filter(|pdu| !pdu.is_redacted()) + .ready_filter(move |pdu| filter.matches(pdu)) + .wide_filter_map(move |pdu| async move { + self.services + .state_accessor + .user_can_see_event(query.user_id?, pdu.room_id().unwrap(), pdu.event_id()) + .await + .then_some(pdu) + }) + .skip(query.skip) + .take(query.limit) + .map(move |mut pdu| { + pdu.set_unsigned(query.user_id); + + pdu + }) + .then(async move |mut pdu| { + if let Err(e) = self + .services + .pdu_metadata + .add_bundled_aggregations_to_pdu(sender_user, &mut pdu) + .await + { + debug_warn!("Failed to add bundled aggregations: {e}"); + } + pdu + }); + + Ok((count, pdus)) + } + + // TODO: The next 4 functions can be, theoretically, squashed into one. + /// Same as search_pdus (searches the search index for matching PDUs), but + /// only returns their short IDs with no filtering or bundled aggregations. + pub async fn search_pdu_ids<'a>( + &'a self, + query: &'a RoomQuery<'_>, + ) -> Result + Send + 'a + use<'a>> { + // result is modeled as a stream such that callers don't have to be refactored + // though an additional async/wrap still exists for now + let shortroomid = self.services.short.get_shortroomid(query.room_id).await?; + + let pdu_ids = self.search_pdu_ids_query_room(query, shortroomid).await; + + let iters = pdu_ids.into_iter().map(IntoIterator::into_iter); + + Ok(set::intersection(iters).stream()) + } + + async fn search_pdu_ids_query_room( + &self, + query: &RoomQuery<'_>, + shortroomid: ShortRoomId, + ) -> Vec> { + tokenize(&query.criteria.search_term) + .stream() + .wide_then(|word| async move { + self.search_pdu_ids_query_words(shortroomid, &word) + .collect::>() + .await + }) + .collect::>() + .await + } + + /// Iterate over PduId's containing a word + fn search_pdu_ids_query_words<'a>( + &'a self, + shortroomid: ShortRoomId, + word: &'a str, + ) -> impl Stream + Send + 'a { + self.search_pdu_ids_query_word(shortroomid, word) + .map(move |key| -> RawPduId { + let key = &key[prefix_len(word)..]; + key.into() + }) + } + + /// Iterate over raw database results for a word + fn search_pdu_ids_query_word<'a>( + &'a self, + shortroomid: ShortRoomId, + word: &'a str, + ) -> impl Stream> + Send + 'a + use<'a> { + // rustc says const'ing this not yet stable + let end_id: RawPduId = PduId { + shortroomid, + shorteventid: PduCount::max(), + } + .into(); + + // Newest pdus first + let end = make_tokenid(shortroomid, word, &end_id); + let prefix = make_prefix(shortroomid, word); + self.db + .tokenids + .rev_raw_keys_from(&end) + .ignore_err() + .ready_take_while(move |key| key.starts_with(&prefix)) + } } /// Splits a string into tokens used as keys in the search inverted index