From b01a93be60c6c6c880d9624c4afe5873d474596d Mon Sep 17 00:00:00 2001
From: Egor Suvorov <egor@neon.tech>
Date: Wed, 23 Nov 2022 23:09:23 +0200
Subject: [PATCH] persistent_range_query: second draft

Move same generic parameters to associated types.
Work around private fields of SumResult<T>
---
 libs/persistent_range_query/src/lib.rs        |  26 ++-
 libs/persistent_range_query/src/naive.rs      | 173 ++++++++---------
 libs/persistent_range_query/src/ops/mod.rs    |  28 +--
 libs/persistent_range_query/src/ops/rsq.rs    | 175 +++++++++---------
 libs/persistent_range_query/tests/rsq_test.rs | 100 +++++-----
 5 files changed, 239 insertions(+), 263 deletions(-)

diff --git a/libs/persistent_range_query/src/lib.rs b/libs/persistent_range_query/src/lib.rs
index 68b41dc819..e743608573 100644
--- a/libs/persistent_range_query/src/lib.rs
+++ b/libs/persistent_range_query/src/lib.rs
@@ -29,8 +29,8 @@ pub trait RangeQueryResult<Key>: Sized {
     }
 }
 
-pub trait LazyRangeInitializer<Key, R> {
-    fn get(&self, range: &Range<Key>) -> R;
+pub trait LazyRangeInitializer<Result: RangeQueryResult<Key>, Key> {
+    fn get(&self, range: &Range<Key>) -> Result;
 }
 
 /// Should be a monoid:
@@ -40,29 +40,25 @@ pub trait LazyRangeInitializer<Key, R> {
 /// 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<Result, Key> {
+pub trait RangeModification<Key> {
+    type Result: RangeQueryResult<Key>;
+
     fn no_op() -> Self;
-    fn apply<'a>(&self, result: &mut Result, range: &Range<Key>);
+    fn apply(&self, result: &mut Self::Result, range: &Range<Key>);
     fn compose(later: &Self, earlier: &mut Self);
 }
 
-pub trait VecVersion<
-    Key,
-    Result: RangeQueryResult<Key>,
-    Modification: RangeModification<Result, Key>,
->: Clone
-{
-    fn get(&self, keys: Range<Key>) -> Result;
+pub trait VecVersion<Modification: RangeModification<Key>, Key>: Clone {
+    fn get(&self, keys: Range<Key>) -> Modification::Result;
     fn modify(&mut self, keys: Range<Key>, modification: Modification);
 }
 
 pub trait PersistentVecStorage<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
     Key,
-    Result: RangeQueryResult<Key>,
-    Initializer: LazyRangeInitializer<Key, Result>,
-    Modification: RangeModification<Result, Key>,
 >
 {
-    type Version: VecVersion<Key, Result, Modification>;
+    type Version: VecVersion<Modification, Key>;
     fn new(all_keys: Range<Key>, initializer: Initializer) -> Self::Version;
 }
diff --git a/libs/persistent_range_query/src/naive.rs b/libs/persistent_range_query/src/naive.rs
index 193349975d..732f7576a1 100644
--- a/libs/persistent_range_query/src/naive.rs
+++ b/libs/persistent_range_query/src/naive.rs
@@ -1,93 +1,80 @@
-use crate::{
-    LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult, VecVersion,
-};
-use std::marker::PhantomData;
-use std::ops::Range;
-
-pub struct NaiveVecVersion<
-    Key: Clone,
-    Result: RangeQueryResult<Key> + Clone,
-    Modification: RangeModification<Result, Key>,
-> {
-    all_keys: Range<Key>,
-    values: Vec<Result>,
-    _modification: PhantomData<Modification>,
-}
-
-impl<
-        Key: Clone,
-        Result: RangeQueryResult<Key> + Clone,
-        Modification: RangeModification<Result, Key>,
-    > Clone for NaiveVecVersion<Key, Result, Modification>
-{
-    fn clone(&self) -> Self {
-        Self {
-            all_keys: self.all_keys.clone(),
-            values: self.values.clone(),
-            _modification: PhantomData,
-        }
-    }
-}
-
-pub trait IndexableKey: Clone {
-    fn index(all_keys: &Range<Self>, key: &Self) -> usize;
-    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self>;
-}
-
-impl<
-        Key: IndexableKey,
-        Result: Clone + RangeQueryResult<Key>,
-        Modification: RangeModification<Result, Key>,
-    > VecVersion<Key, Result, Modification> for NaiveVecVersion<Key, Result, Modification>
-{
-    fn get(&self, keys: Range<Key>) -> Result {
-        let mut result = Result::new_for_empty_range();
-        let mut result_range = keys.start.clone()..keys.start.clone();
-        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);
-            Result::add(
-                &mut result,
-                &result_range,
-                &self.values[index],
-                &element_range,
-            );
-            result_range.end = element_range.end;
-        }
-        result
-    }
-
-    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.values[index], &element_range);
-        }
-    }
-}
-
-pub struct NaiveVecStorage;
-
-impl<
-        Key: IndexableKey,
-        Result: Clone + RangeQueryResult<Key>,
-        Initializer: LazyRangeInitializer<Key, Result>,
-        Modification: RangeModification<Result, Key>,
-    > PersistentVecStorage<Key, Result, Initializer, Modification> for NaiveVecStorage
-{
-    type Version = NaiveVecVersion<Key, Result, Modification>;
-
-    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self::Version {
-        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)));
-        }
-        NaiveVecVersion {
-            all_keys,
-            values,
-            _modification: PhantomData,
-        }
-    }
-}
+use crate::{
+    LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult, VecVersion,
+};
+use std::ops::Range;
+
+pub struct NaiveVecVersion<Modification: RangeModification<Key>, Key> {
+    all_keys: Range<Key>,
+    values: Vec<Modification::Result>,
+}
+
+impl<Modification: RangeModification<Key>, Key: Clone> Clone for NaiveVecVersion<Modification, Key>
+where
+    Modification::Result: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            all_keys: self.all_keys.clone(),
+            values: self.values.clone(),
+        }
+    }
+}
+
+pub trait IndexableKey: Clone {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize;
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self>;
+}
+
+impl<Modification: RangeModification<Key>, Key: IndexableKey> VecVersion<Modification, Key>
+    for NaiveVecVersion<Modification, Key>
+where
+    Modification::Result: Clone,
+{
+    fn get(&self, 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(&self.all_keys, &keys.start)
+            ..IndexableKey::index(&self.all_keys, &keys.end)
+        {
+            let element_range = IndexableKey::element_range(&self.all_keys, index);
+            Modification::Result::add(
+                &mut result,
+                &result_range,
+                &self.values[index],
+                &element_range,
+            );
+            result_range.end = element_range.end;
+        }
+        result
+    }
+
+    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.values[index], &element_range);
+        }
+    }
+}
+
+pub struct NaiveVecStorage;
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: IndexableKey,
+    > PersistentVecStorage<Modification, Initializer, Key> for NaiveVecStorage
+where
+    Modification::Result: Clone,
+{
+    type Version = NaiveVecVersion<Modification, Key>;
+
+    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self::Version {
+        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)));
+        }
+        NaiveVecVersion { all_keys, values }
+    }
+}
diff --git a/libs/persistent_range_query/src/ops/mod.rs b/libs/persistent_range_query/src/ops/mod.rs
index 5fb1712db3..665824e6c2 100644
--- a/libs/persistent_range_query/src/ops/mod.rs
+++ b/libs/persistent_range_query/src/ops/mod.rs
@@ -1,14 +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,
-        }
-    }
-}
+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,
+        }
+    }
+}
diff --git a/libs/persistent_range_query/src/ops/rsq.rs b/libs/persistent_range_query/src/ops/rsq.rs
index e0b9d5f533..386bfa7979 100644
--- a/libs/persistent_range_query/src/ops/rsq.rs
+++ b/libs/persistent_range_query/src/ops/rsq.rs
@@ -1,85 +1,90 @@
-//! # Range Sum Query
-
-use crate::ops::SameElementsInitializer;
-use crate::{LazyRangeInitializer, RangeModification, RangeQueryResult};
-use std::borrow::Borrow;
-use std::ops::{AddAssign, Mul, Range, Sub};
-
-// 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: for<'a> AddAssign<&'a T> + From<u8>, Key> RangeQueryResult<Key> for SumResult<T> {
-    fn new_for_empty_range() -> Self {
-        SumResult { sum: 0.into() }
-    }
-
-    fn add(left: &mut Self, _left_range: &Range<Key>, right: &Self, _right_range: &Range<Key>) {
-        left.sum += &right.sum
-    }
-}
-
-impl<Key, T, TR: Borrow<T>, KeyDiff> LazyRangeInitializer<Key, SumResult<T>>
-    for SameElementsInitializer<TR>
-where
-    for<'a> &'a T: Mul<KeyDiff, Output = T>,
-    for<'b> &'b Key: Sub<Output = KeyDiff>,
-{
-    fn get(&self, range: &Range<Key>) -> SumResult<T> {
-        SumResult {
-            sum: self.initial_element_value.borrow() * (&range.end - &range.start),
-        }
-    }
-}
-
-#[derive(Copy, Clone, Debug)]
-pub enum AddAssignModification<T> {
-    None,
-    Add(T),
-    Assign(T),
-}
-
-impl<T: Clone + for<'a> AddAssign<&'a T>, Key> RangeModification<SumResult<T>, Key>
-    for AddAssignModification<T>
-where
-    for<'a> SameElementsInitializer<&'a T>: LazyRangeInitializer<Key, SumResult<T>>,
-{
-    fn no_op() -> Self {
-        AddAssignModification::None
-    }
-
-    fn apply<'a>(&self, result: &'a 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,
-        }
-    }
-}
+//! # Range Sum Query
+
+use crate::ops::SameElementsInitializer;
+use crate::{LazyRangeInitializer, RangeModification, RangeQueryResult};
+use std::borrow::Borrow;
+use std::ops::{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: for<'a> AddAssign<&'a T> + From<u8>, Key> RangeQueryResult<Key> for SumResult<T> {
+    fn new_for_empty_range() -> Self {
+        SumResult { sum: 0.into() }
+    }
+
+    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 apply<'a>(&self, result: &'a 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,
+        }
+    }
+}
diff --git a/libs/persistent_range_query/tests/rsq_test.rs b/libs/persistent_range_query/tests/rsq_test.rs
index 5c7db94a21..a9e95220e3 100644
--- a/libs/persistent_range_query/tests/rsq_test.rs
+++ b/libs/persistent_range_query/tests/rsq_test.rs
@@ -1,56 +1,44 @@
-use persistent_range_query::naive::*;
-use persistent_range_query::ops::rsq::*;
-use persistent_range_query::ops::SameElementsInitializer;
-use persistent_range_query::{PersistentVecStorage, VecVersion};
-use std::ops::{Mul, Range, Sub};
-
-#[derive(Copy, Clone, Debug)]
-struct K(u8);
-
-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 u8)..K(all_keys.start.0 + index as u8 + 1)
-    }
-}
-
-struct KDiff(i16);
-
-impl Sub<&K> for &K {
-    type Output = KDiff;
-
-    fn sub(self, rhs: &K) -> Self::Output {
-        KDiff((self.0 as i16) - (rhs.0 as i16))
-    }
-}
-
-impl<'a> Mul<KDiff> for &'a i32 {
-    type Output = i32;
-
-    fn mul(self, rhs: KDiff) -> Self::Output {
-        self * (rhs.0 as i32)
-    }
-}
-
-#[test]
-fn test_naive() {
-    let mut s: NaiveVecVersion<_, _, AddAssignModification<_>> =
-        NaiveVecStorage::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.clone();
-
-    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);
-}
+use persistent_range_query::naive::*;
+use persistent_range_query::ops::rsq::*;
+use persistent_range_query::ops::SameElementsInitializer;
+use persistent_range_query::{PersistentVecStorage, VecVersion};
+use std::ops::Range;
+
+#[derive(Copy, Clone, Debug)]
+struct K(u8);
+
+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 u8)..K(all_keys.start.0 + index as u8 + 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
+    }
+}
+
+#[test]
+fn test_naive() {
+    let mut s: NaiveVecVersion<AddAssignModification<i32>, _> =
+        NaiveVecStorage::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.clone();
+
+    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);
+}