persistent_range_query: add stress test

libs/persistent_range_query/tests/rsq_test.rs

@@ -1,12 +1,14 @@
 use persistent_range_query::naive::*;
+use persistent_range_query::ops::rsq::AddAssignModification::Add;
 use persistent_range_query::ops::rsq::*;
 use persistent_range_query::ops::SameElementsInitializer;
 use persistent_range_query::segment_tree::{MidpointableKey, PersistentSegmentTree};
 use persistent_range_query::{PersistentVecStorage, VecReadableVersion};
+use rand::{Rng, SeedableRng};
 use std::ops::Range;
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
-struct K(u8);
+struct K(u16);
 
 impl IndexableKey for K {
     fn index(all_keys: &Range<Self>, key: &Self) -> usize {
@@ -14,7 +16,7 @@ impl IndexableKey for K {
     }
 
     fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self> {
-        K(all_keys.start.0 + index as u8)..K(all_keys.start.0 + index as u8 + 1)
+        K(all_keys.start.0 + index as u16)..K(all_keys.start.0 + index as u16 + 1)
     }
 }
 
@@ -59,3 +61,56 @@ fn test_naive() {
 fn test_segment_tree() {
     test_storage::<PersistentSegmentTree<_, _, _>>();
 }
+
+#[test]
+fn test_stress() {
+    const LEN: u16 = 17_238;
+    const OPERATIONS: i32 = 20_000;
+
+    let mut rng = rand::rngs::StdRng::seed_from_u64(0);
+    let mut naive: NaiveVecStorage<AddAssignModification<i32>, _, _> =
+        NaiveVecStorage::new(K(0)..K(LEN), SameElementsInitializer::new(2i32));
+    let mut segm_tree: PersistentSegmentTree<AddAssignModification<i32>, _, _> =
+        PersistentSegmentTree::new(K(0)..K(LEN), SameElementsInitializer::new(2i32));
+
+    fn gen_range(rng: &mut impl Rng) -> Range<K> {
+        let l: u16 = rng.gen_range(0..LEN);
+        let r: u16 = rng.gen_range(0..LEN);
+        if l <= r {
+            K(l)..K(r)
+        } else {
+            K(r)..K(l)
+        }
+    }
+
+    for _ in 0..2 {
+        let checksum_range = gen_range(&mut rng);
+        let checksum_before: i32 = *naive.get(&checksum_range).sum();
+        assert_eq!(checksum_before, *segm_tree.get(&checksum_range).sum());
+
+        let naive_before = naive.freeze();
+        let segm_tree_before = segm_tree.freeze();
+        assert_eq!(checksum_before, *naive_before.get(&checksum_range).sum());
+        assert_eq!(checksum_before, *segm_tree.get(&checksum_range).sum());
+
+        for _ in 0..OPERATIONS {
+            {
+                let range = gen_range(&mut rng);
+                assert_eq!(naive.get(&range).sum(), segm_tree.get(&range).sum());
+            }
+            {
+                let range = gen_range(&mut rng);
+                let val = rng.gen_range(-10i32..=10i32);
+                let op = Add(val);
+                naive.modify(&range, &op);
+                segm_tree.modify(&range, &op);
+            }
+        }
+
+        assert_eq!(checksum_before, *naive_before.get(&checksum_range).sum());
+        assert_eq!(
+            checksum_before,
+            *segm_tree_before.get(&checksum_range).sum()
+        );
+    }
+}