persistent_range_query: add layer_map_test

Move same generic parameters to associated types.
Work around private fields of SumResult<T>

Cargo.lock

@@ -2255,6 +2255,14 @@ version = "2.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "478c572c3d73181ff3c2539045f6eb99e5491218eae919370993b890cdbdd98e"
 
+[[package]]
+name = "persistent_range_query"
+version = "0.1.0"
+dependencies = [
+ "rand",
+ "workspace_hack",
+]
+
 [[package]]
 name = "petgraph"
 version = "0.6.2"

libs/persistent_range_query/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "persistent_range_query"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+
+[dev-dependencies]
+rand = "0.8.3"

libs/persistent_range_query/src/lib.rs

@@ -0,0 +1,78 @@
+use std::ops::Range;
+
+pub mod naive;
+pub mod ops;
+pub mod segment_tree;
+
+/// Should be a monoid:
+/// * Identity element: for all a: combine(new_for_empty_range(), a) = combine(a, new_for_empty_range()) = a
+/// * Associativity: for all a, b, c: combine(combine(a, b), c) == combine(a, combine(b, c))
+pub trait RangeQueryResult<Key>: Sized + Clone {
+    // Clone is equivalent to combine with an empty range.
+
+    fn new_for_empty_range() -> Self;
+
+    // Contract: left_range.end == right_range.start
+    // left_range.start == left_range.end == right_range.start == right_range.end is still possible
+    fn combine(
+        left: &Self,
+        left_range: &Range<Key>,
+        right: &Self,
+        right_range: &Range<Key>,
+    ) -> Self;
+
+    fn add(left: &mut Self, left_range: &Range<Key>, right: &Self, right_range: &Range<Key>);
+}
+
+pub trait LazyRangeInitializer<Result: RangeQueryResult<Key>, Key> {
+    fn get(&self, range: &Range<Key>) -> Result;
+}
+
+/// Should be a monoid:
+/// * Identity element: for all op: compose(no_op(), op) == compose(op, no_op()) == op
+/// * Associativity: for all op_1, op_2, op_3: compose(compose(op_1, op_2), op_3) == compose(op_1, compose(op_2, op_3))
+///
+/// Should left act on Result:
+/// * Identity operation: for all r: no_op().apply(r) == r
+/// * Compatibility: for all op_1, op_2, r: op_1.apply(op_2.apply(r)) == compose(op_1, op_2).apply(r)
+pub trait RangeModification<Key> {
+    type Result: RangeQueryResult<Key>;
+
+    fn no_op() -> Self;
+    fn is_no_op(&self) -> bool;
+    fn is_reinitialization(&self) -> bool;
+    fn apply(&self, result: &mut Self::Result, range: &Range<Key>);
+    fn compose(later: &Self, earlier: &mut Self);
+}
+
+pub trait VecReadableVersion<Modification: RangeModification<Key>, Key> {
+    fn get(&self, keys: &Range<Key>) -> Modification::Result;
+}
+
+// TODO: use trait alias when stabilized
+pub trait VecFrozenVersion<Modification: RangeModification<Key>, Key>:
+    Clone + VecReadableVersion<Modification, Key>
+{
+}
+
+impl<
+        T: Clone + VecReadableVersion<Modification, Key>,
+        Modification: RangeModification<Key>,
+        Key,
+    > VecFrozenVersion<Modification, Key> for T
+{
+}
+
+pub trait PersistentVecStorage<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key,
+>: VecReadableVersion<Modification, Key>
+{
+    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self;
+
+    type FrozenVersion: VecFrozenVersion<Modification, Key>;
+
+    fn modify(&mut self, keys: &Range<Key>, modification: &Modification);
+    fn freeze(&mut self) -> Self::FrozenVersion;
+}

libs/persistent_range_query/src/naive.rs

@@ -0,0 +1,115 @@
+use crate::{
+    LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult,
+    VecReadableVersion,
+};
+use std::marker::PhantomData;
+use std::ops::Range;
+use std::rc::Rc;
+
+pub struct NaiveFrozenVersion<Modification: RangeModification<Key>, Key> {
+    all_keys: Range<Key>,
+    values: Rc<Box<Vec<Modification::Result>>>,
+}
+
+pub trait IndexableKey: Clone {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize;
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self>;
+}
+
+fn get<Modification: RangeModification<Key>, Key: IndexableKey>(
+    all_keys: &Range<Key>,
+    values: &Vec<Modification::Result>,
+    keys: &Range<Key>,
+) -> Modification::Result {
+    let mut result = Modification::Result::new_for_empty_range();
+    let mut result_range = keys.start.clone()..keys.start.clone();
+    for index in
+        IndexableKey::index(&all_keys, &keys.start)..IndexableKey::index(&all_keys, &keys.end)
+    {
+        let element_range = IndexableKey::element_range(&all_keys, index);
+        Modification::Result::add(&mut result, &result_range, &values[index], &element_range);
+        result_range.end = element_range.end;
+    }
+    result
+}
+
+impl<Modification: RangeModification<Key>, Key: IndexableKey> VecReadableVersion<Modification, Key>
+    for NaiveFrozenVersion<Modification, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        get::<Modification, Key>(&self.all_keys, &self.values, keys)
+    }
+}
+
+// Manual implementation of `Clone` becase `derive` requires `Modification: Clone`
+impl<Modification: RangeModification<Key>, Key: Clone> Clone
+    for NaiveFrozenVersion<Modification, Key>
+{
+    fn clone(&self) -> Self {
+        Self {
+            all_keys: self.all_keys.clone(),
+            values: self.values.clone(),
+        }
+    }
+}
+
+// TODO: is it at all possible to store previous versions in this struct,
+// without any Rc<>?
+pub struct NaiveVecStorage<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: IndexableKey,
+> {
+    all_keys: Range<Key>,
+    last_version: Vec<Modification::Result>,
+    _initializer: PhantomData<Initializer>,
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: IndexableKey,
+    > VecReadableVersion<Modification, Key> for NaiveVecStorage<Modification, Initializer, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        get::<Modification, Key>(&self.all_keys, &self.last_version, keys)
+    }
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: IndexableKey,
+    > PersistentVecStorage<Modification, Initializer, Key>
+    for NaiveVecStorage<Modification, Initializer, Key>
+{
+    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self {
+        let mut values = Vec::with_capacity(IndexableKey::index(&all_keys, &all_keys.end));
+        for index in 0..values.capacity() {
+            values.push(initializer.get(&IndexableKey::element_range(&all_keys, index)));
+        }
+        NaiveVecStorage {
+            all_keys,
+            last_version: values,
+            _initializer: PhantomData,
+        }
+    }
+
+    type FrozenVersion = NaiveFrozenVersion<Modification, Key>;
+
+    fn modify(&mut self, keys: &Range<Key>, modification: &Modification) {
+        for index in IndexableKey::index(&self.all_keys, &keys.start)
+            ..IndexableKey::index(&self.all_keys, &keys.end)
+        {
+            let element_range = IndexableKey::element_range(&self.all_keys, index);
+            modification.apply(&mut self.last_version[index], &element_range);
+        }
+    }
+
+    fn freeze(&mut self) -> Self::FrozenVersion {
+        NaiveFrozenVersion::<Modification, Key> {
+            all_keys: self.all_keys.clone(),
+            values: Rc::new(Box::new(self.last_version.clone())),
+        }
+    }
+}

libs/persistent_range_query/src/ops/mod.rs

@@ -0,0 +1,14 @@
+pub mod rsq;
+
+#[derive(Copy, Clone, Debug)]
+pub struct SameElementsInitializer<T> {
+    initial_element_value: T,
+}
+
+impl<T> SameElementsInitializer<T> {
+    pub fn new(initial_element_value: T) -> Self {
+        SameElementsInitializer {
+            initial_element_value,
+        }
+    }
+}

libs/persistent_range_query/src/ops/rsq.rs

@@ -0,0 +1,118 @@
+//! # Range Sum Query
+
+use crate::ops::SameElementsInitializer;
+use crate::{LazyRangeInitializer, RangeModification, RangeQueryResult};
+use std::borrow::Borrow;
+use std::ops::{Add, AddAssign, Range};
+
+// TODO: commutative Add
+
+#[derive(Clone, Copy, Debug)]
+pub struct SumResult<T> {
+    sum: T,
+}
+
+impl<T> SumResult<T> {
+    pub fn sum(&self) -> &T {
+        &self.sum
+    }
+}
+
+impl<T: Clone + for<'a> AddAssign<&'a T> + From<u8>, Key> RangeQueryResult<Key> for SumResult<T>
+where
+    for<'a> &'a T: Add<&'a T, Output = T>,
+{
+    fn new_for_empty_range() -> Self {
+        SumResult { sum: 0.into() }
+    }
+
+    fn combine(
+        left: &Self,
+        _left_range: &Range<Key>,
+        right: &Self,
+        _right_range: &Range<Key>,
+    ) -> Self {
+        SumResult {
+            sum: &left.sum + &right.sum,
+        }
+    }
+
+    fn add(left: &mut Self, _left_range: &Range<Key>, right: &Self, _right_range: &Range<Key>) {
+        left.sum += &right.sum
+    }
+}
+
+pub trait SumOfSameElements<Key> {
+    fn sum(initial_element_value: &Self, keys: &Range<Key>) -> Self;
+}
+
+impl<T: SumOfSameElements<Key>, TB: Borrow<T>, Key> LazyRangeInitializer<SumResult<T>, Key>
+    for SameElementsInitializer<TB>
+where
+    SumResult<T>: RangeQueryResult<Key>,
+{
+    fn get(&self, range: &Range<Key>) -> SumResult<T> {
+        SumResult {
+            sum: SumOfSameElements::sum(self.initial_element_value.borrow(), range),
+        }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub enum AddAssignModification<T> {
+    None,
+    Add(T),
+    Assign(T),
+}
+
+impl<T: Clone + for<'a> AddAssign<&'a T>, Key> RangeModification<Key> for AddAssignModification<T>
+where
+    SumResult<T>: RangeQueryResult<Key>,
+    for<'a> SameElementsInitializer<&'a T>: LazyRangeInitializer<SumResult<T>, Key>,
+{
+    type Result = SumResult<T>;
+
+    fn no_op() -> Self {
+        AddAssignModification::None
+    }
+
+    fn is_no_op(&self) -> bool {
+        match self {
+            AddAssignModification::None => true,
+            _ => false,
+        }
+    }
+
+    fn is_reinitialization(&self) -> bool {
+        match self {
+            AddAssignModification::Assign(_) => true,
+            _ => false,
+        }
+    }
+
+    fn apply(&self, result: &mut SumResult<T>, range: &Range<Key>) {
+        use AddAssignModification::*;
+        match self {
+            None => {}
+            Add(x) | Assign(x) => {
+                let to_add = SameElementsInitializer::new(x).get(range).sum;
+                if let Assign(_) = self {
+                    result.sum = to_add;
+                } else {
+                    result.sum += &to_add;
+                }
+            }
+        }
+    }
+
+    fn compose(later: &Self, earlier: &mut Self) {
+        use AddAssignModification::*;
+        match (later, earlier) {
+            (_, e @ None) => *e = later.clone(),
+            (None, _) => {}
+            (Assign(_), e) => *e = later.clone(),
+            (Add(x), Add(y)) => *y += x,
+            (Add(x), Assign(value)) => *value += x,
+        }
+    }
+}

libs/persistent_range_query/src/segment_tree.rs

@@ -0,0 +1,255 @@
+//! # Segment Tree
+//! It is a competitive programming folklore data structure. Do not confuse with the interval tree.
+
+use crate::{LazyRangeInitializer, PersistentVecStorage, RangeQueryResult, VecReadableVersion};
+use std::ops::Range;
+use std::rc::Rc;
+
+pub trait MidpointableKey: Clone + Ord + Sized {
+    fn midpoint(range: &Range<Self>) -> Self;
+}
+
+pub trait RangeModification<Key>: Clone + crate::RangeModification<Key> {}
+
+// TODO: use trait alias when stabilized
+impl<T: Clone + crate::RangeModification<Key>, Key> RangeModification<Key> for T {}
+
+#[derive(Debug)]
+struct Node<Modification: RangeModification<Key>, Key> {
+    result: Modification::Result,
+    modify_children: Modification,
+    left: Option<Rc<Self>>,
+    right: Option<Rc<Self>>,
+}
+
+// Manual implementation because we don't need `Key: Clone` for this, unlike with `derive`.
+impl<Modification: RangeModification<Key>, Key> Clone for Node<Modification, Key> {
+    fn clone(&self) -> Self {
+        Node {
+            result: self.result.clone(),
+            modify_children: self.modify_children.clone(),
+            left: self.left.clone(),
+            right: self.right.clone(),
+        }
+    }
+}
+
+impl<Modification: RangeModification<Key>, Key> Node<Modification, Key> {
+    fn new<Initializer: LazyRangeInitializer<Modification::Result, Key>>(
+        range: &Range<Key>,
+        initializer: &Initializer,
+    ) -> Self {
+        Node {
+            result: initializer.get(range),
+            modify_children: Modification::no_op(),
+            left: None,
+            right: None,
+        }
+    }
+
+    pub fn apply(&mut self, modification: &Modification, range: &Range<Key>) {
+        modification.apply(&mut self.result, range);
+        Modification::compose(modification, &mut self.modify_children);
+        if self.modify_children.is_reinitialization() {
+            self.left = None;
+            self.right = None;
+        }
+    }
+
+    pub fn force_children<Initializer: LazyRangeInitializer<Modification::Result, Key>>(
+        &mut self,
+        initializer: &Initializer,
+        range_left: &Range<Key>,
+        range_right: &Range<Key>,
+    ) {
+        let left = Rc::make_mut(
+            self.left
+                .get_or_insert_with(|| Rc::new(Node::new(&range_left, initializer))),
+        );
+        let right = Rc::make_mut(
+            self.right
+                .get_or_insert_with(|| Rc::new(Node::new(&range_right, initializer))),
+        );
+        left.apply(&self.modify_children, &range_left);
+        right.apply(&self.modify_children, &range_right);
+        self.modify_children = Modification::no_op();
+    }
+
+    pub fn recalculate_from_children(&mut self, range_left: &Range<Key>, range_right: &Range<Key>) {
+        assert!(self.modify_children.is_no_op());
+        assert!(self.left.is_some());
+        assert!(self.right.is_some());
+        self.result = Modification::Result::combine(
+            &self.left.as_ref().unwrap().result,
+            &range_left,
+            &self.right.as_ref().unwrap().result,
+            &range_right,
+        );
+    }
+}
+
+fn split_range<Key: MidpointableKey>(range: &Range<Key>) -> (Range<Key>, Range<Key>) {
+    let range_left = range.start.clone()..MidpointableKey::midpoint(range);
+    let range_right = range_left.end.clone()..range.end.clone();
+    (range_left, range_right)
+}
+
+pub struct PersistentSegmentTreeVersion<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: Clone,
+> {
+    root: Rc<Node<Modification, Key>>,
+    all_keys: Range<Key>,
+    initializer: Rc<Initializer>,
+}
+
+// Manual implementation because we don't need `Key: Clone` for this, unlike with `derive`.
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: Clone,
+    > Clone for PersistentSegmentTreeVersion<Modification, Initializer, Key>
+{
+    fn clone(&self) -> Self {
+        Self {
+            root: self.root.clone(),
+            all_keys: self.all_keys.clone(),
+            initializer: self.initializer.clone(),
+        }
+    }
+}
+
+fn get<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: MidpointableKey,
+>(
+    node: &mut Rc<Node<Modification, Key>>,
+    node_keys: &Range<Key>,
+    initializer: &Initializer,
+    keys: &Range<Key>,
+) -> Modification::Result {
+    if node_keys.end <= keys.start || keys.end <= node_keys.start {
+        return Modification::Result::new_for_empty_range();
+    }
+    if keys.start <= node_keys.start && node_keys.end <= keys.end {
+        return node.result.clone();
+    }
+    let node = Rc::make_mut(node);
+    let (left_keys, right_keys) = split_range(node_keys);
+    node.force_children(initializer, &left_keys, &right_keys);
+    let mut result = get(node.left.as_mut().unwrap(), &left_keys, initializer, keys);
+    Modification::Result::add(
+        &mut result,
+        &left_keys,
+        &get(node.right.as_mut().unwrap(), &right_keys, initializer, keys),
+        &right_keys,
+    );
+    result
+}
+
+fn modify<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: MidpointableKey,
+>(
+    node: &mut Rc<Node<Modification, Key>>,
+    node_keys: &Range<Key>,
+    initializer: &Initializer,
+    keys: &Range<Key>,
+    modification: &Modification,
+) {
+    if modification.is_no_op() || node_keys.end <= keys.start || keys.end <= node_keys.start {
+        return;
+    }
+    let node = Rc::make_mut(node);
+    if keys.start <= node_keys.start && node_keys.end <= keys.end {
+        node.apply(modification, node_keys);
+        return;
+    }
+    let (left_keys, right_keys) = split_range(node_keys);
+    node.force_children(initializer, &left_keys, &right_keys);
+    modify(
+        node.left.as_mut().unwrap(),
+        &left_keys,
+        initializer,
+        keys,
+        &modification,
+    );
+    modify(
+        node.right.as_mut().unwrap(),
+        &right_keys,
+        initializer,
+        keys,
+        &modification,
+    );
+    node.recalculate_from_children(&left_keys, &right_keys);
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: MidpointableKey,
+    > VecReadableVersion<Modification, Key>
+    for PersistentSegmentTreeVersion<Modification, Initializer, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        get(
+            &mut self.root.clone(), // TODO: do not always force a branch
+            &self.all_keys,
+            self.initializer.as_ref(),
+            keys,
+        )
+    }
+}
+
+pub struct PersistentSegmentTree<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: MidpointableKey,
+>(PersistentSegmentTreeVersion<Modification, Initializer, Key>);
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: MidpointableKey,
+    > VecReadableVersion<Modification, Key>
+    for PersistentSegmentTree<Modification, Initializer, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        self.0.get(keys)
+    }
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: MidpointableKey,
+    > PersistentVecStorage<Modification, Initializer, Key>
+    for PersistentSegmentTree<Modification, Initializer, Key>
+{
+    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self {
+        PersistentSegmentTree(PersistentSegmentTreeVersion {
+            root: Rc::new(Node::new(&all_keys, &initializer)),
+            all_keys: all_keys,
+            initializer: Rc::new(initializer),
+        })
+    }
+
+    type FrozenVersion = PersistentSegmentTreeVersion<Modification, Initializer, Key>;
+
+    fn modify(&mut self, keys: &Range<Key>, modification: &Modification) {
+        modify(
+            &mut self.0.root, // TODO: do not always force a branch
+            &self.0.all_keys,
+            self.0.initializer.as_ref(),
+            keys,
+            modification,
+        )
+    }
+
+    fn freeze(&mut self) -> Self::FrozenVersion {
+        self.0.clone()
+    }
+}

libs/persistent_range_query/tests/layer_map_test.rs

@@ -0,0 +1,295 @@
+use persistent_range_query::naive::{IndexableKey, NaiveVecStorage};
+use persistent_range_query::ops::SameElementsInitializer;
+use persistent_range_query::segment_tree::{MidpointableKey, PersistentSegmentTree};
+use persistent_range_query::{
+    LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult,
+    VecReadableVersion,
+};
+use std::cmp::Ordering;
+use std::ops::Range;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
+struct PageIndex(u32);
+type LayerId = String;
+
+impl IndexableKey for PageIndex {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize {
+        (key.0 as usize) - (all_keys.start.0 as usize)
+    }
+
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self> {
+        PageIndex(all_keys.start.0 + index as u32)..PageIndex(all_keys.start.0 + index as u32 + 1)
+    }
+}
+
+impl MidpointableKey for PageIndex {
+    fn midpoint(range: &Range<Self>) -> Self {
+        PageIndex(range.start.0 + (range.end.0 - range.start.0) / 2)
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+struct LayerMapInformation {
+    // Only make sense for a range of length 1.
+    last_layer: Option<LayerId>,
+    last_image_layer: Option<LayerId>,
+    // Work for all ranges
+    max_delta_layers: (usize, Range<PageIndex>),
+}
+
+impl LayerMapInformation {
+    fn last_layers(&self) -> (&Option<LayerId>, &Option<LayerId>) {
+        (&self.last_layer, &self.last_image_layer)
+    }
+
+    fn max_delta_layers(&self) -> &(usize, Range<PageIndex>) {
+        &self.max_delta_layers
+    }
+}
+
+fn merge_ranges(left: &Range<PageIndex>, right: &Range<PageIndex>) -> Range<PageIndex> {
+    if left.is_empty() {
+        right.clone()
+    } else if right.is_empty() {
+        left.clone()
+    } else if left.end == right.start {
+        left.start..right.end
+    } else {
+        left.clone()
+    }
+}
+
+impl RangeQueryResult<PageIndex> for LayerMapInformation {
+    fn new_for_empty_range() -> Self {
+        LayerMapInformation {
+            last_layer: None,
+            last_image_layer: None,
+            max_delta_layers: (0, PageIndex(0)..PageIndex(0)),
+        }
+    }
+
+    fn combine(
+        left: &Self,
+        _left_range: &Range<PageIndex>,
+        right: &Self,
+        _right_range: &Range<PageIndex>,
+    ) -> Self {
+        // Note that either range may be empty.
+        LayerMapInformation {
+            last_layer: left
+                .last_layer
+                .as_ref()
+                .or_else(|| right.last_layer.as_ref())
+                .cloned(),
+            last_image_layer: left
+                .last_image_layer
+                .as_ref()
+                .or_else(|| right.last_image_layer.as_ref())
+                .cloned(),
+            max_delta_layers: match left.max_delta_layers.0.cmp(&right.max_delta_layers.0) {
+                Ordering::Less => right.max_delta_layers.clone(),
+                Ordering::Greater => left.max_delta_layers.clone(),
+                Ordering::Equal => (
+                    left.max_delta_layers.0,
+                    merge_ranges(&left.max_delta_layers.1, &right.max_delta_layers.1),
+                ),
+            },
+        }
+    }
+
+    fn add(
+        left: &mut Self,
+        left_range: &Range<PageIndex>,
+        right: &Self,
+        right_range: &Range<PageIndex>,
+    ) {
+        *left = Self::combine(&left, left_range, right, right_range);
+    }
+}
+
+#[derive(Clone, Debug)]
+struct AddDeltaLayers {
+    last_layer: LayerId,
+    count: usize,
+}
+
+#[derive(Clone, Debug)]
+struct LayerMapModification {
+    add_image_layer: Option<LayerId>,
+    add_delta_layers: Option<AddDeltaLayers>,
+}
+
+impl LayerMapModification {
+    fn add_image_layer(layer: impl Into<LayerId>) -> Self {
+        LayerMapModification {
+            add_image_layer: Some(layer.into()),
+            add_delta_layers: None,
+        }
+    }
+
+    fn add_delta_layer(layer: impl Into<LayerId>) -> Self {
+        LayerMapModification {
+            add_image_layer: None,
+            add_delta_layers: Some(AddDeltaLayers {
+                last_layer: layer.into(),
+                count: 1,
+            }),
+        }
+    }
+}
+
+impl RangeModification<PageIndex> for LayerMapModification {
+    type Result = LayerMapInformation;
+
+    fn no_op() -> Self {
+        LayerMapModification {
+            add_image_layer: None,
+            add_delta_layers: None,
+        }
+    }
+
+    fn is_no_op(&self) -> bool {
+        self.add_image_layer.is_none() && self.add_delta_layers.is_none()
+    }
+
+    fn is_reinitialization(&self) -> bool {
+        self.add_image_layer.is_some()
+    }
+
+    fn apply(&self, result: &mut Self::Result, range: &Range<PageIndex>) {
+        if let Some(layer) = &self.add_image_layer {
+            result.last_layer = Some(layer.clone());
+            result.last_image_layer = Some(layer.clone());
+            result.max_delta_layers = (0, range.clone());
+        }
+        if let Some(AddDeltaLayers { last_layer, count }) = &self.add_delta_layers {
+            result.last_layer = Some(last_layer.clone());
+            result.max_delta_layers.0 += count;
+        }
+    }
+
+    fn compose(later: &Self, earlier: &mut Self) {
+        if later.add_image_layer.is_some() {
+            *earlier = later.clone();
+            return;
+        }
+        if let Some(AddDeltaLayers { last_layer, count }) = &later.add_delta_layers {
+            let res = earlier.add_delta_layers.get_or_insert(AddDeltaLayers {
+                last_layer: LayerId::default(),
+                count: 0,
+            });
+            res.last_layer = last_layer.clone();
+            res.count += count;
+        }
+    }
+}
+
+impl LazyRangeInitializer<LayerMapInformation, PageIndex> for SameElementsInitializer<()> {
+    fn get(&self, range: &Range<PageIndex>) -> LayerMapInformation {
+        LayerMapInformation {
+            last_layer: None,
+            last_image_layer: None,
+            max_delta_layers: (0, range.clone()),
+        }
+    }
+}
+
+fn test_layer_map<
+    S: PersistentVecStorage<LayerMapModification, SameElementsInitializer<()>, PageIndex>,
+>() {
+    let mut s = S::new(
+        PageIndex(0)..PageIndex(100),
+        SameElementsInitializer::new(()),
+    );
+    s.modify(
+        &(PageIndex(0)..PageIndex(70)),
+        &LayerMapModification::add_image_layer("Img0..70"),
+    );
+    s.modify(
+        &(PageIndex(50)..PageIndex(100)),
+        &LayerMapModification::add_image_layer("Img50..100"),
+    );
+    s.modify(
+        &(PageIndex(10)..PageIndex(60)),
+        &LayerMapModification::add_delta_layer("Delta10..60"),
+    );
+    let s_before_last_delta = s.freeze();
+    s.modify(
+        &(PageIndex(20)..PageIndex(80)),
+        &LayerMapModification::add_delta_layer("Delta20..80"),
+    );
+
+    assert_eq!(
+        s.get(&(PageIndex(5)..PageIndex(6))).last_layers(),
+        (&Some("Img0..70".to_owned()), &Some("Img0..70".to_owned()))
+    );
+    assert_eq!(
+        s.get(&(PageIndex(15)..PageIndex(16))).last_layers(),
+        (
+            &Some("Delta10..60".to_owned()),
+            &Some("Img0..70".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(25)..PageIndex(26))).last_layers(),
+        (
+            &Some("Delta20..80".to_owned()),
+            &Some("Img0..70".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(65)..PageIndex(66))).last_layers(),
+        (
+            &Some("Delta20..80".to_owned()),
+            &Some("Img50..100".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(95)..PageIndex(96))).last_layers(),
+        (
+            &Some("Img50..100".to_owned()),
+            &Some("Img50..100".to_owned())
+        )
+    );
+
+    assert_eq!(
+        s.get(&(PageIndex(0)..PageIndex(100))).max_delta_layers(),
+        &(2, PageIndex(20)..PageIndex(60)),
+    );
+    assert_eq!(
+        *s_before_last_delta
+            .get(&(PageIndex(0)..PageIndex(100)))
+            .max_delta_layers(),
+        (1, PageIndex(10)..PageIndex(60)),
+    );
+
+    assert_eq!(
+        *s.get(&(PageIndex(10)..PageIndex(30))).max_delta_layers(),
+        (2, PageIndex(20)..PageIndex(30))
+    );
+    assert_eq!(
+        *s.get(&(PageIndex(10)..PageIndex(20))).max_delta_layers(),
+        (1, PageIndex(10)..PageIndex(20))
+    );
+
+    assert_eq!(
+        *s.get(&(PageIndex(70)..PageIndex(80))).max_delta_layers(),
+        (1, PageIndex(70)..PageIndex(80))
+    );
+    assert_eq!(
+        *s_before_last_delta
+            .get(&(PageIndex(70)..PageIndex(80)))
+            .max_delta_layers(),
+        (0, PageIndex(70)..PageIndex(80))
+    );
+}
+
+#[test]
+fn test_naive() {
+    test_layer_map::<NaiveVecStorage<_, _, _>>();
+}
+
+#[test]
+fn test_segment_tree() {
+    test_layer_map::<PersistentSegmentTree<_, _, _>>();
+}

libs/persistent_range_query/tests/rsq_test.rs

@@ -0,0 +1,116 @@
+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(u16);
+
+impl IndexableKey for K {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize {
+        (key.0 as usize) - (all_keys.start.0 as usize)
+    }
+
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self> {
+        K(all_keys.start.0 + index as u16)..K(all_keys.start.0 + index as u16 + 1)
+    }
+}
+
+impl SumOfSameElements<K> for i32 {
+    fn sum(initial_element_value: &Self, keys: &Range<K>) -> Self {
+        initial_element_value * (keys.end.0 - keys.start.0) as Self
+    }
+}
+
+impl MidpointableKey for K {
+    fn midpoint(range: &Range<Self>) -> Self {
+        K(range.start.0 + (range.end.0 - range.start.0) / 2)
+    }
+}
+
+fn test_storage<
+    S: PersistentVecStorage<AddAssignModification<i32>, SameElementsInitializer<i32>, K>,
+>() {
+    let mut s = S::new(K(0)..K(12), SameElementsInitializer::new(0i32));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 0);
+
+    s.modify(&(K(2)..K(5)), &AddAssignModification::Add(3));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 3 + 3 + 3);
+    let s_old = s.freeze();
+
+    s.modify(&(K(3)..K(6)), &AddAssignModification::Assign(10));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 3 + 10 + 10 + 10);
+
+    s.modify(&(K(4)..K(7)), &AddAssignModification::Add(2));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 3 + 10 + 12 + 12 + 2);
+
+    assert_eq!(*s.get(&(K(4)..K(6))).sum(), 12 + 12);
+    assert_eq!(*s_old.get(&(K(4)..K(6))).sum(), 3);
+}
+
+#[test]
+fn test_naive() {
+    test_storage::<NaiveVecStorage<_, _, _>>();
+}
+
+#[test]
+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()
+        );
+    }
+}