feat: aggregate sum path query

grovedb-query/src/aggregate_sum_query/insert.rs

@@ -0,0 +1,176 @@
+use std::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};
+
+use super::AggregateSumQuery;
+use crate::QueryItem;
+
+impl AggregateSumQuery {
+    /// Adds an individual key to the query, so that its value (or its absence)
+    /// in the tree will be included in the resulting proof.
+    ///
+    /// If the key or a range including the key already exists in the query,
+    /// this will have no effect. If the query already includes a range that has
+    /// a non-inclusive bound equal to the key, the bound will be changed to be
+    /// inclusive.
+    pub fn insert_key(&mut self, key: Vec<u8>) {
+        let key = QueryItem::Key(key);
+        self.insert_item(key);
+    }
+
+    /// Adds multiple individual keys to the query, so that its value (or its
+    /// absence) in the tree will be included in the resulting proof.
+    ///
+    /// If the key or a range including the key already exists in the query,
+    /// this will have no effect. If the query already includes a range that has
+    /// a non-inclusive bound equal to the key, the bound will be changed to be
+    /// inclusive.
+    pub fn insert_keys(&mut self, keys: Vec<Vec<u8>>) {
+        for key in keys {
+            let key = QueryItem::Key(key);
+            self.insert_item(key);
+        }
+    }
+
+    /// Adds a range to the query, so that all the entries in the tree with keys
+    /// in the range will be included in the resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be joined together.
+    pub fn insert_range(&mut self, range: Range<Vec<u8>>) {
+        let range = QueryItem::Range(range);
+        self.insert_item(range);
+    }
+
+    /// Adds an inclusive range to the query, so that all the entries in the
+    /// tree with keys in the range will be included in the resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be merged together.
+    pub fn insert_range_inclusive(&mut self, range: RangeInclusive<Vec<u8>>) {
+        let range = QueryItem::RangeInclusive(range);
+        self.insert_item(range);
+    }
+
+    /// Adds a range until a certain included value to the query, so that all
+    /// the entries in the tree with keys in the range will be included in the
+    /// resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be joined together.
+    pub fn insert_range_to_inclusive(&mut self, range: RangeToInclusive<Vec<u8>>) {
+        let range = QueryItem::RangeToInclusive(range);
+        self.insert_item(range);
+    }
+
+    /// Adds a range from a certain included value to the query, so that all
+    /// the entries in the tree with keys in the range will be included in the
+    /// resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be joined together.
+    pub fn insert_range_from(&mut self, range: RangeFrom<Vec<u8>>) {
+        let range = QueryItem::RangeFrom(range);
+        self.insert_item(range);
+    }
+
+    /// Adds a range until a certain non included value to the query, so that
+    /// all the entries in the tree with keys in the range will be included
+    /// in the resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be joined together.
+    pub fn insert_range_to(&mut self, range: RangeTo<Vec<u8>>) {
+        let range = QueryItem::RangeTo(range);
+        self.insert_item(range);
+    }
+
+    /// Adds a range after the first value, so that all the entries in the tree
+    /// with keys in the range will be included in the resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be joined together.
+    pub fn insert_range_after(&mut self, range: RangeFrom<Vec<u8>>) {
+        let range = QueryItem::RangeAfter(range);
+        self.insert_item(range);
+    }
+
+    /// Adds a range after the first value, until a certain non included value
+    /// to the query, so that all the entries in the tree with keys in the
+    /// range will be included in the resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be joined together.
+    pub fn insert_range_after_to(&mut self, range: Range<Vec<u8>>) {
+        let range = QueryItem::RangeAfterTo(range);
+        self.insert_item(range);
+    }
+
+    /// Adds a range after the first value, until a certain included value to
+    /// the query, so that all the entries in the tree with keys in the
+    /// range will be included in the resulting proof.
+    ///
+    /// If a range including the range already exists in the query, this will
+    /// have no effect. If the query already includes a range that overlaps with
+    /// the range, the ranges will be joined together.
+    pub fn insert_range_after_to_inclusive(&mut self, range: RangeInclusive<Vec<u8>>) {
+        let range = QueryItem::RangeAfterToInclusive(range);
+        self.insert_item(range);
+    }
+
+    /// Adds a range of all potential values to the query, so that the query
+    /// will return all values
+    ///
+    /// All other items in the query will be discarded as you are now getting
+    /// back all elements.
+    pub fn insert_all(&mut self) {
+        let range = QueryItem::RangeFull(RangeFull);
+        self.insert_item(range);
+    }
+
+    /// Adds the `QueryItem` to the query, first checking to see if it collides
+    /// with any existing ranges or keys. All colliding items will be removed
+    /// then merged together so that the query includes the minimum number of
+    /// items (with no items covering any duplicate parts of keyspace) while
+    /// still including every key or range that has been added to the query.
+    pub fn insert_item(&mut self, mut item: QueryItem) {
+        // since `QueryItem::eq` considers items equal if they collide at all
+        // (including keys within ranges or ranges which partially overlap),
+        // `items.take` will remove the first item which collides
+
+        self.items = self
+            .items
+            .iter()
+            .filter_map(|our_item| {
+                if our_item.is_key() && item.is_key() && our_item == &item {
+                    None
+                } else if our_item.collides_with(&item) {
+                    item.merge_assign(our_item);
+                    None
+                } else {
+                    Some(our_item.clone()) // todo: manage this without a clone
+                }
+            })
+            .collect();
+
+        // since we need items to be sorted we do
+        match self.items.binary_search(&item) {
+            Ok(_) => {
+                unreachable!("this shouldn't be possible")
+            }
+            Err(pos) => self.items.insert(pos, item),
+        }
+    }
+
+    /// Performs an insert_item on each item in the vector.
+    pub fn insert_items(&mut self, items: Vec<QueryItem>) {
+        for item in items {
+            self.insert_item(item)
+        }
+    }
+}

grovedb-query/src/aggregate_sum_query/merge.rs

@@ -0,0 +1,72 @@
+use super::AggregateSumQuery;
+use crate::error::Error;
+
+impl AggregateSumQuery {
+    /// Merge multiple aggregate sum queries into one.
+    pub fn merge_multiple(mut queries: Vec<AggregateSumQuery>) -> Result<Self, Error> {
+        if queries.is_empty() {
+            // We put sum 0 and limit 0 to represent a no-op query
+            return Ok(AggregateSumQuery::new(0, Some(0)));
+        }
+
+        // Slight performance improvement via swap_remove
+        let mut merged_query = queries.swap_remove(0);
+        let mut aggregate_sum_limit = merged_query.sum_limit;
+        let expected_left_to_right = merged_query.left_to_right;
+        let mut merged_limit: Option<u16> = merged_query.limit_of_items_to_check;
+
+        for query in queries {
+            if query.left_to_right != expected_left_to_right {
+                return Err(Error::NotSupported(
+                    "Cannot merge queries with differing left_to_right values".to_string(),
+                ));
+            }
+
+            aggregate_sum_limit = aggregate_sum_limit
+                .checked_add(query.sum_limit)
+                .ok_or(Error::Overflow("Overflow when merging sum limits"))?;
+
+            merged_limit = match (merged_limit, query.limit_of_items_to_check) {
+                (Some(a), Some(b)) => Some(
+                    a.checked_add(b)
+                        .ok_or(Error::Overflow("Overflow when merging item check limits"))?,
+                ),
+                _ => None, // if either is None, result is None
+            };
+
+            merged_query.insert_items(query.items);
+        }
+
+        merged_query.sum_limit = aggregate_sum_limit;
+        merged_query.limit_of_items_to_check = merged_limit;
+
+        Ok(merged_query)
+    }
+
+    /// Merge another aggregate sum query into this one.
+    pub fn merge_with(&mut self, other: AggregateSumQuery) -> Result<(), Error> {
+        if self.left_to_right != other.left_to_right {
+            return Err(Error::NotSupported(
+                "Cannot merge queries with differing left_to_right values".to_string(),
+            ));
+        }
+
+        self.sum_limit = self
+            .sum_limit
+            .checked_add(other.sum_limit)
+            .ok_or(Error::Overflow("Overflow when merging sum limits"))?;
+
+        self.limit_of_items_to_check =
+            match (self.limit_of_items_to_check, other.limit_of_items_to_check) {
+                (Some(a), Some(b)) => Some(
+                    a.checked_add(b)
+                        .ok_or(Error::Overflow("Overflow when merging item check limits"))?,
+                ),
+                _ => None,
+            };
+
+        self.insert_items(other.items);
+
+        Ok(())
+    }
+}