events: Add methods to compute poll results

crates/ruma-common/src/events/poll.rs

@@ -4,6 +4,71 @@
 //!
 //! [MSC3381]: https://github.com/matrix-org/matrix-spec-proposals/pull/3381
 
+use std::collections::{BTreeMap, BTreeSet};
+
+use indexmap::IndexMap;
+use js_int::uint;
+
+use crate::{MilliSecondsSinceUnixEpoch, UserId};
+
+use self::{response::OriginalSyncPollResponseEvent, start::PollContentBlock};
+
 pub mod end;
 pub mod response;
 pub mod start;
+
+/// Generate the current results with the given poll and responses.
+///
+/// If the `end_timestamp` is provided, any response with an `origin_server_ts` after that timestamp
+/// is ignored. If it is not provided, `MilliSecondsSinceUnixEpoch::now()` will be used instead.
+///
+/// This method will handle invalid responses, or several response from the same user so all
+/// responses to the poll should be provided.
+///
+/// Returns a map of answer ID to a set of user IDs that voted for them. When using `.iter()` or
+/// `.into_iter()` on the map, the results are sorted from the highest number of votes to the
+/// lowest.
+pub fn compile_poll_results<'a>(
+    poll: &'a PollContentBlock,
+    responses: impl IntoIterator<Item = &'a OriginalSyncPollResponseEvent>,
+    end_timestamp: Option<MilliSecondsSinceUnixEpoch>,
+) -> IndexMap<&'a str, BTreeSet<&'a UserId>> {
+    let end_ts = end_timestamp.unwrap_or_else(MilliSecondsSinceUnixEpoch::now);
+
+    let users_selections = responses
+        .into_iter()
+        .filter(|ev| {
+            // Filter out responses after the end_timestamp.
+            ev.origin_server_ts <= end_ts
+        })
+        .fold(BTreeMap::new(), |mut acc, ev| {
+            let response =
+                acc.entry(&*ev.sender).or_insert((MilliSecondsSinceUnixEpoch(uint!(0)), None));
+
+            // Only keep the latest selections for each user.
+            if response.0 < ev.origin_server_ts {
+                *response = (ev.origin_server_ts, ev.content.selections.validate(poll));
+            }
+
+            acc
+        });
+
+    // Aggregate the selections by answer.
+    let mut results =
+        IndexMap::from_iter(poll.answers.iter().map(|a| (a.id.as_str(), BTreeSet::new())));
+
+    for (user, (_, selections)) in users_selections {
+        if let Some(selections) = selections {
+            for selection in selections {
+                results
+                    .get_mut(selection)
+                    .expect("validated selections should only match possible answers")
+                    .insert(user);
+            }
+        }
+    }
+
+    results.sort_by(|_, a, _, b| b.len().cmp(&a.len()));
+
+    results
+}

crates/ruma-common/src/events/poll/end.rs

@@ -15,6 +15,11 @@ use crate::{
 };
 
 /// The payload for a poll end event.
+///
+/// This type can be generated from the poll start and poll response events with
+/// [`OriginalSyncPollStartEvent::compile_results()`].
+///
+/// [`OriginalSyncPollStartEvent::compile_results()`]: super::start::OriginalSyncPollStartEvent::compile_results
 #[derive(Clone, Debug, Serialize, Deserialize, EventContent)]
 #[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
 #[ruma_event(type = "org.matrix.msc3381.v2.poll.end", alias = "m.poll.end", kind = MessageLike)]
@@ -77,6 +82,17 @@ impl PollEndEventContent {
 #[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
 pub struct PollResultsContentBlock(BTreeMap<String, UInt>);
 
+impl PollResultsContentBlock {
+    /// Get these results sorted from the highest number of votes to the lowest.
+    ///
+    /// Returns a list of `(answer ID, number of votes)`.
+    pub fn sorted(&self) -> Vec<(&str, UInt)> {
+        let mut sorted = self.0.iter().map(|(id, count)| (id.as_str(), *count)).collect::<Vec<_>>();
+        sorted.sort_by(|(_, a), (_, b)| b.cmp(a));
+        sorted
+    }
+}
+
 impl From<BTreeMap<String, UInt>> for PollResultsContentBlock {
     fn from(value: BTreeMap<String, UInt>) -> Self {
         Self(value)

crates/ruma-common/src/events/poll/response.rs

@@ -7,6 +7,8 @@ use serde::{Deserialize, Serialize};
 
 use crate::{events::relation::Reference, OwnedEventId};
 
+use super::start::PollContentBlock;
+
 /// The payload for a poll response event.
 #[derive(Clone, Debug, Serialize, Deserialize, EventContent)]
 #[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
@@ -53,6 +55,23 @@ impl SelectionsContentBlock {
     pub fn is_empty(&self) -> bool {
         self.0.is_empty()
     }
+
+    /// Validate these selections against the given `PollContentBlock`.
+    ///
+    /// Returns the list of valid selections in this `SelectionsContentBlock`, or `None` if there is
+    /// no valid selection.
+    pub fn validate(&self, poll: &PollContentBlock) -> Option<impl Iterator<Item = &str>> {
+        // Vote is spoiled if any answer is unknown.
+        if self.0.iter().any(|s| !poll.answers.iter().any(|a| a.id == *s)) {
+            return None;
+        }
+
+        // Fallback to the maximum value for usize because we can't have more selections than that
+        // in memory.
+        let max_selections: usize = poll.max_selections.try_into().unwrap_or(usize::MAX);
+
+        Some(self.0.iter().take(max_selections).map(Deref::deref))
+    }
 }
 
 impl From<Vec<String>> for SelectionsContentBlock {

crates/ruma-common/src/events/poll/start.rs

@@ -12,6 +12,12 @@ use poll_answers_serde::PollAnswersDeHelper;
 
 use crate::{events::message::TextContentBlock, serde::StringEnum, PrivOwnedStr};
 
+use super::{
+    compile_poll_results,
+    end::{PollEndEventContent, PollResultsContentBlock},
+    response::OriginalSyncPollResponseEvent,
+};
+
 /// The payload for a poll start event.
 #[derive(Clone, Debug, Serialize, Deserialize, EventContent)]
 #[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
@@ -59,6 +65,74 @@ impl PollStartEventContent {
     }
 }
 
+impl OriginalSyncPollStartEvent {
+    /// Compile the results for this poll with the given response into a `PollEndEventContent`.
+    ///
+    /// It generates a default text representation of the results in English.
+    ///
+    /// This uses [`compile_poll_results()`] internally.
+    pub fn compile_results<'a>(
+        &'a self,
+        responses: impl IntoIterator<Item = &'a OriginalSyncPollResponseEvent>,
+    ) -> PollEndEventContent {
+        let computed = compile_poll_results(&self.content.poll, responses, None);
+
+        // Construct the results and get the top answer(s).
+        let mut top_answers = Vec::new();
+        let mut top_count = uint!(0);
+
+        let results =
+            PollResultsContentBlock::from_iter(computed.into_iter().map(|(id, users)| {
+                let count = users.len().try_into().unwrap_or(UInt::MAX);
+
+                if count >= top_count {
+                    top_answers.push(id);
+                    top_count = count;
+                }
+
+                (id.to_owned(), count)
+            }));
+
+        // Get the text representation of the best answers.
+        let top_answers_text = top_answers
+            .into_iter()
+            .map(|id| {
+                let text = &self
+                    .content
+                    .poll
+                    .answers
+                    .iter()
+                    .find(|a| a.id == id)
+                    .expect("top answer ID should be a valid answer ID")
+                    .text;
+                // Use the plain text representation, fallback to the first text representation or
+                // to the answer ID.
+                text.find_plain()
+                    .or_else(|| text.get(0).map(|t| t.body.as_ref()))
+                    .unwrap_or(id)
+                    .to_owned()
+            })
+            .collect::<Vec<_>>();
+
+        // Construct the plain text representation.
+        let plain_text = match top_answers_text.len() {
+            l if l > 1 => {
+                let answers = top_answers_text.join(", ");
+                format!("The poll has closed. Top answers: {answers}")
+            }
+            l if l == 1 => {
+                format!("The poll has closed. Top answer: {}", top_answers_text[0])
+            }
+            _ => "The poll has closed with no top answer".to_owned(),
+        };
+
+        let mut end = PollEndEventContent::with_plain_text(plain_text, self.event_id.clone());
+        end.poll_results = Some(results);
+
+        end
+    }
+}
+
 /// A block for poll content.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 #[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]

crates/ruma-common/tests/events/poll.rs

@@ -1,24 +1,25 @@
 #![cfg(feature = "unstable-msc3381")]
 
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, ops::Range};
 
 use assert_matches2::assert_matches;
-use js_int::uint;
+use js_int::{uint, UInt};
 use ruma_common::{
     events::{
         message::TextContentBlock,
         poll::{
+            compile_poll_results,
             end::PollEndEventContent,
-            response::PollResponseEventContent,
+            response::{OriginalSyncPollResponseEvent, PollResponseEventContent},
             start::{
-                PollAnswer, PollAnswers, PollAnswersError, PollContentBlock, PollKind,
+                OriginalSyncPollStartEvent, PollAnswer, PollAnswers, PollAnswersError,
-                PollStartEventContent,
+                PollContentBlock, PollKind, PollStartEventContent,
             },
         },
         relation::Reference,
         AnyMessageLikeEvent, MessageLikeEvent,
     },
-    owned_event_id,
+    owned_event_id, MilliSecondsSinceUnixEpoch,
 };
 use serde_json::{from_value as from_json_value, json, to_value as to_json_value};
 
@@ -340,3 +341,303 @@ fn end_event_deserialization() {
     assert_matches!(message_event.content.relates_to, Reference { event_id, .. });
     assert_eq!(event_id, "$related_event:notareal.hs");
 }
+
+fn new_poll_response(
+    event_id: &str,
+    user_id: &str,
+    ts: UInt,
+    selections: &[&str],
+) -> OriginalSyncPollResponseEvent {
+    from_json_value(json!({
+      "type": "org.matrix.msc3381.v2.poll.response",
+      "sender": user_id,
+      "origin_server_ts": ts,
+      "event_id": event_id,
+      "content": {
+        "m.relates_to": {
+          "rel_type": "m.reference",
+          "event_id": "$poll_start_event_id"
+        },
+        "org.matrix.msc3381.v2.selections": selections,
+      }
+    }))
+    .unwrap()
+}
+
+fn generate_poll_responses(
+    range: Range<usize>,
+    selections: &[&str],
+) -> Vec<OriginalSyncPollResponseEvent> {
+    let mut responses = Vec::with_capacity(range.len());
+
+    for i in range {
+        let event_id = format!("$valid_event_{i}");
+        let user_id = format!("@valid_user_{i}:localhost");
+        let ts = 1000 + i as u16;
+
+        responses.push(new_poll_response(&event_id, &user_id, ts.into(), selections));
+    }
+
+    responses
+}
+
+#[test]
+fn compute_results() {
+    let poll: OriginalSyncPollStartEvent = from_json_value(json!({
+        "type": "org.matrix.msc3381.v2.poll.start",
+        "sender": "@alice:localhost",
+        "event_id": "$poll_start_event_id",
+        "origin_server_ts": 1,
+        "content": {
+          "org.matrix.msc1767.text": [
+            {
+              "mimetype": "text/plain",
+              "body": "What should we order for the party?\n1. Pizza 🍕\n2. Poutine 🍟\n3. Italian 🍝\n4. Wings 🔥"
+            }
+          ],
+          "org.matrix.msc3381.v2.poll": {
+            "kind": "m.disclosed",
+            "max_selections": 2,
+            "question": {
+              "org.matrix.msc1767.text": [{"body": "What should we order for the party?"}]
+            },
+            "answers": [
+              {"org.matrix.msc3381.v2.id": "pizza", "org.matrix.msc1767.text": [{"body": "Pizza 🍕"}]},
+              {"org.matrix.msc3381.v2.id": "poutine", "org.matrix.msc1767.text": [{"body": "Poutine 🍟"}]},
+              {"org.matrix.msc3381.v2.id": "italian", "org.matrix.msc1767.text": [{"body": "Italian 🍝"}]},
+              {"org.matrix.msc3381.v2.id": "wings", "org.matrix.msc1767.text": [{"body": "Wings 🔥"}]},
+            ]
+          }
+        }
+      })).unwrap();
+
+    // Populate responses.
+    let mut responses = generate_poll_responses(0..5, &["pizza"]);
+    responses.extend(generate_poll_responses(5..10, &["poutine"]));
+    responses.extend(generate_poll_responses(10..15, &["italian"]));
+    responses.extend(generate_poll_responses(15..20, &["wings"]));
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 5);
+    assert_eq!(counted.get("poutine").unwrap().len(), 5);
+    assert_eq!(counted.get("italian").unwrap().len(), 5);
+    assert_eq!(counted.get("wings").unwrap().len(), 5);
+    let mut iter = counted.keys();
+    assert_eq!(iter.next(), Some(&"pizza"));
+    assert_eq!(iter.next(), Some(&"poutine"));
+    assert_eq!(iter.next(), Some(&"italian"));
+    assert_eq!(iter.next(), Some(&"wings"));
+    assert_eq!(iter.next(), None);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(5));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(5));
+    assert_eq!(*results.get("italian").unwrap(), uint!(5));
+    assert_eq!(*results.get("wings").unwrap(), uint!(5));
+    assert_eq!(
+        results.sorted().as_slice(),
+        &[("italian", uint!(5)), ("pizza", uint!(5)), ("poutine", uint!(5)), ("wings", uint!(5))]
+    );
+    assert_eq!(
+        poll_end.text.find_plain(),
+        Some("The poll has closed. Top answers: Pizza 🍕, Poutine 🍟, Italian 🍝, Wings 🔥")
+    );
+
+    responses.extend(vec![
+        new_poll_response(
+            "$multi_event_1",
+            "@multi_user_1:localhost",
+            uint!(2000),
+            &["poutine", "wings"],
+        ),
+        new_poll_response(
+            "$multi_event_2",
+            "@multi_user_2:localhost",
+            uint!(2200),
+            &["poutine", "italian"],
+        ),
+    ]);
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 5);
+    assert_eq!(counted.get("poutine").unwrap().len(), 7);
+    assert_eq!(counted.get("italian").unwrap().len(), 6);
+    assert_eq!(counted.get("wings").unwrap().len(), 6);
+    let mut iter = counted.keys();
+    assert_eq!(iter.next(), Some(&"poutine"));
+    assert_eq!(iter.next(), Some(&"italian"));
+    assert_eq!(iter.next(), Some(&"wings"));
+    assert_eq!(iter.next(), Some(&"pizza"));
+    assert_eq!(iter.next(), None);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(5));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(7));
+    assert_eq!(*results.get("italian").unwrap(), uint!(6));
+    assert_eq!(*results.get("wings").unwrap(), uint!(6));
+    assert_eq!(
+        results.sorted().as_slice(),
+        &[("poutine", uint!(7)), ("italian", uint!(6)), ("wings", uint!(6)), ("pizza", uint!(5))]
+    );
+    assert_eq!(poll_end.text.find_plain(), Some("The poll has closed. Top answer: Poutine 🍟"));
+
+    responses.extend(vec![
+        new_poll_response(
+            "$multi_same_event_1",
+            "@multi_same_user_1:localhost",
+            uint!(3000),
+            &["poutine", "poutine"],
+        ),
+        new_poll_response(
+            "$multi_same_event_2",
+            "@multi_same_user_2:localhost",
+            uint!(3300),
+            &["pizza", "pizza"],
+        ),
+    ]);
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 6);
+    assert_eq!(counted.get("poutine").unwrap().len(), 8);
+    assert_eq!(counted.get("italian").unwrap().len(), 6);
+    assert_eq!(counted.get("wings").unwrap().len(), 6);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(6));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(8));
+    assert_eq!(*results.get("italian").unwrap(), uint!(6));
+    assert_eq!(*results.get("wings").unwrap(), uint!(6));
+
+    let changing_user_1 = "@changing_user_1:localhost";
+    let changing_user_2 = "@changing_user_2:localhost";
+    let changing_user_3 = "@changing_user_3:localhost";
+
+    responses.extend(vec![
+        new_poll_response("$valid_for_now_event_1", changing_user_1, uint!(4000), &["wings"]),
+        new_poll_response("$valid_for_now_event_2", changing_user_2, uint!(4100), &["wings"]),
+        new_poll_response("$valid_for_now_event_3", changing_user_3, uint!(4200), &["wings"]),
+    ]);
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 6);
+    assert_eq!(counted.get("poutine").unwrap().len(), 8);
+    assert_eq!(counted.get("italian").unwrap().len(), 6);
+    assert_eq!(counted.get("wings").unwrap().len(), 9);
+    let mut iter = counted.keys();
+    assert_eq!(iter.next(), Some(&"wings"));
+    assert_eq!(iter.next(), Some(&"poutine"));
+    assert_eq!(iter.next(), Some(&"pizza"));
+    assert_eq!(iter.next(), Some(&"italian"));
+    assert_eq!(iter.next(), None);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(6));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(8));
+    assert_eq!(*results.get("italian").unwrap(), uint!(6));
+    assert_eq!(*results.get("wings").unwrap(), uint!(9));
+    assert_eq!(
+        results.sorted().as_slice(),
+        &[("wings", uint!(9)), ("poutine", uint!(8)), ("italian", uint!(6)), ("pizza", uint!(6))]
+    );
+    assert_eq!(poll_end.text.find_plain(), Some("The poll has closed. Top answer: Wings 🔥"));
+
+    // Change with new selection.
+    responses.push(new_poll_response(
+        "$change_vote_event",
+        changing_user_1,
+        uint!(4400),
+        &["italian"],
+    ));
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 6);
+    assert_eq!(counted.get("poutine").unwrap().len(), 8);
+    assert_eq!(counted.get("italian").unwrap().len(), 7);
+    assert_eq!(counted.get("wings").unwrap().len(), 8);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(6));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(8));
+    assert_eq!(*results.get("italian").unwrap(), uint!(7));
+    assert_eq!(*results.get("wings").unwrap(), uint!(8));
+
+    // Change with no selection.
+    responses.push(new_poll_response(
+        "$no_selection_vote_event",
+        changing_user_1,
+        uint!(4500),
+        &[],
+    ));
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 6);
+    assert_eq!(counted.get("poutine").unwrap().len(), 8);
+    assert_eq!(counted.get("italian").unwrap().len(), 6);
+    assert_eq!(counted.get("wings").unwrap().len(), 8);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(6));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(8));
+    assert_eq!(*results.get("italian").unwrap(), uint!(6));
+    assert_eq!(*results.get("wings").unwrap(), uint!(8));
+
+    // Change with invalid selection.
+    responses.push(new_poll_response(
+        "$invalid_vote_event",
+        changing_user_2,
+        uint!(4500),
+        &["indian"],
+    ));
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 6);
+    assert_eq!(counted.get("poutine").unwrap().len(), 8);
+    assert_eq!(counted.get("italian").unwrap().len(), 6);
+    assert_eq!(counted.get("wings").unwrap().len(), 7);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(6));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(8));
+    assert_eq!(*results.get("italian").unwrap(), uint!(6));
+    assert_eq!(*results.get("wings").unwrap(), uint!(7));
+
+    // Response older than most recent one is ignored.
+    responses.push(new_poll_response("$past_event", changing_user_3, uint!(1), &["pizza"]));
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 6);
+    assert_eq!(counted.get("poutine").unwrap().len(), 8);
+    assert_eq!(counted.get("italian").unwrap().len(), 6);
+    assert_eq!(counted.get("wings").unwrap().len(), 7);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(6));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(8));
+    assert_eq!(*results.get("italian").unwrap(), uint!(6));
+    assert_eq!(*results.get("wings").unwrap(), uint!(7));
+
+    // Response in the future is ignored.
+    let future_ts = MilliSecondsSinceUnixEpoch::now().0 + uint!(100_000);
+    responses.push(new_poll_response("$future_event", changing_user_3, future_ts, &["pizza"]));
+
+    let counted = compile_poll_results(&poll.content.poll, &responses, None);
+    assert_eq!(counted.get("pizza").unwrap().len(), 6);
+    assert_eq!(counted.get("poutine").unwrap().len(), 8);
+    assert_eq!(counted.get("italian").unwrap().len(), 6);
+    assert_eq!(counted.get("wings").unwrap().len(), 7);
+
+    let poll_end = poll.compile_results(&responses);
+    let results = poll_end.poll_results.unwrap();
+    assert_eq!(*results.get("pizza").unwrap(), uint!(6));
+    assert_eq!(*results.get("poutine").unwrap(), uint!(8));
+    assert_eq!(*results.get("italian").unwrap(), uint!(6));
+    assert_eq!(*results.get("wings").unwrap(), uint!(7));
+}