diff --git a/src/service/rooms/timeline/stitcher/algorithm.rs b/src/service/rooms/timeline/stitcher/algorithm.rs index 13f0a9be..679859ee 100644 --- a/src/service/rooms/timeline/stitcher/algorithm.rs +++ b/src/service/rooms/timeline/stitcher/algorithm.rs @@ -1,7 +1,4 @@ -use std::{ - cmp::Ordering, - collections::{BTreeSet, HashMap, HashSet}, -}; +use std::collections::HashSet; use indexmap::IndexSet; use itertools::Itertools; @@ -39,7 +36,7 @@ pub(super) struct Stitcher<'backend, B: StitcherBackend> { impl Stitcher<'_, B> { pub(super) fn new(backend: &B) -> Stitcher<'_, B> { Stitcher { backend } } - pub(super) fn stitch<'id>(&self, batch: Batch<'id>) -> OrderUpdates<'id, B::Key> { + pub(super) fn stitch<'id>(&self, batch: &Batch<'id>) -> OrderUpdates<'id, B::Key> { let mut gap_updates = Vec::new(); let mut all_new_events: HashSet<&'id str> = HashSet::new(); @@ -100,9 +97,12 @@ impl Stitcher<'_, B> { .sorted_by(batch.compare_by_dag_received()) .collect_vec(); - let mut items = Vec::with_capacity( - events_to_insert.capacity() + events_to_insert.capacity().div_euclid(2), - ); + // allocate 1.5x the size of the to-insert list + let items_capacity = events_to_insert + .capacity() + .saturating_add(events_to_insert.capacity().div_euclid(2)); + + let mut items = Vec::with_capacity(items_capacity); for event in events_to_insert { let missing_prev_events: HashSet = batch diff --git a/src/service/rooms/timeline/stitcher/mod.rs b/src/service/rooms/timeline/stitcher/mod.rs index 471d1357..e3d9e521 100644 --- a/src/service/rooms/timeline/stitcher/mod.rs +++ b/src/service/rooms/timeline/stitcher/mod.rs @@ -1,7 +1,4 @@ -use std::{ - cmp::Ordering, - collections::{BTreeMap, HashSet}, -}; +use std::{cmp::Ordering, collections::HashSet}; use indexmap::IndexMap; @@ -58,10 +55,10 @@ pub(super) struct Batch<'id> { } impl<'id> Batch<'id> { - pub(super) fn from_edges(edges: EventEdges<'_>) -> Batch<'_> { + pub(super) fn from_edges<'edges>(edges: &EventEdges<'edges>) -> Batch<'edges> { let mut events = IndexMap::new(); - for (event, prev_events) in &edges { + for (event, prev_events) in edges { let predecessor_set = Self::find_predecessor_set(event, &edges); events.insert(*event, EventPredecessors { diff --git a/src/service/rooms/timeline/stitcher/test/mod.rs b/src/service/rooms/timeline/stitcher/test/mod.rs index e40b37ed..4a89ab62 100644 --- a/src/service/rooms/timeline/stitcher/test/mod.rs +++ b/src/service/rooms/timeline/stitcher/test/mod.rs @@ -36,7 +36,9 @@ impl<'id> TestStitcherBackend<'id> { }, | None => unreachable!("we just checked that this index is valid"), } - gap_index + 1 + gap_index.checked_add(1).expect( + "should never allocate usize::MAX ids. what kind of test are you running", + ) }; let to_insert: Vec<_> = update @@ -97,8 +99,8 @@ fn run_testcase(testcase: parser::TestCase<'_>) { for (index, phase) in testcase.into_iter().enumerate() { let stitcher = Stitcher::new(&backend); - let batch = Batch::from_edges(phase.batch); - let updates = stitcher.stitch(batch); + let batch = Batch::from_edges(&phase.batch); + let updates = stitcher.stitch(&batch); println!(); println!("===== phase {index}"); diff --git a/src/service/rooms/timeline/stitcher/test/parser.rs b/src/service/rooms/timeline/stitcher/test/parser.rs index 961e0273..9a106430 100644 --- a/src/service/rooms/timeline/stitcher/test/parser.rs +++ b/src/service/rooms/timeline/stitcher/test/parser.rs @@ -1,8 +1,8 @@ -use std::collections::{BTreeMap, HashSet}; +use std::collections::HashSet; use indexmap::IndexMap; -use crate::rooms::timeline::stitcher::{StitchedItem, test}; +use crate::rooms::timeline::stitcher::StitchedItem; pub(super) type TestEventId<'id> = &'id str;