From 7c37bb7d09181f3df1b851a2fb9ecea1e88a3c4b Mon Sep 17 00:00:00 2001
From: QuantumExplorer <quantum@dash.org>
Date: Tue, 20 Aug 2024 16:24:38 +0700
Subject: [PATCH] feat: PathQuery depth (#328)

---
 grovedb/build.rs             |  5 +++-
 grovedb/src/query/mod.rs     | 43 +++++++++++++++++++++++++++
 merk/src/proofs/query/mod.rs | 57 ++++++++++++++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 1 deletion(-)

diff --git a/grovedb/build.rs b/grovedb/build.rs
index 83a38900..f5792754 100644
--- a/grovedb/build.rs
+++ b/grovedb/build.rs
@@ -13,7 +13,10 @@ fn main() {
     let grovedbg_zip_path = out_dir.join("grovedbg.zip");
 
     if !grovedbg_zip_path.exists() {
-        let response = reqwest::blocking::get(format!("https://github.com/dashpay/grovedbg/releases/download/{GROVEDBG_VERSION}/grovedbg-{GROVEDBG_VERSION}.zip"))
+        let response = reqwest::blocking::get(format!(
+            "https://github.com/dashpay/grovedbg/releases/download/{GROVEDBG_VERSION}/\
+grovedbg-{GROVEDBG_VERSION}.zip"
+        ))
         .expect("can't download GroveDBG artifact");
 
         let mut grovedbg_zip = File::create(&grovedbg_zip_path).unwrap();
diff --git a/grovedb/src/query/mod.rs b/grovedb/src/query/mod.rs
index f140bb05..7f896cca 100644
--- a/grovedb/src/query/mod.rs
+++ b/grovedb/src/query/mod.rs
@@ -134,6 +134,14 @@ impl PathQuery {
         Self { path, query }
     }
 
+    /// The max depth of the query, this is the maximum layers we could get back
+    /// from grovedb
+    /// If the max depth can not be calculated we get None
+    /// This would occur if the recursion level was too high
+    pub fn max_depth(&self) -> Option<u16> {
+        self.query.query.max_depth()
+    }
+
     /// Gets the path of all terminal keys
     pub fn terminal_keys(
         &self,
@@ -1634,6 +1642,8 @@ mod tests {
             },
         };
 
+        assert_eq!(path_query.max_depth(), Some(4));
+
         {
             let path = vec![];
             let first = path_query
@@ -1710,4 +1720,37 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_max_depth_limit() {
+        /// Creates a `Query` with nested `SubqueryBranch` up to the specified
+        /// depth non-recursively.
+        fn create_non_recursive_query(subquery_depth: usize) -> Query {
+            let mut root_query = Query::new_range_full();
+            let mut current_query = &mut root_query;
+
+            for _ in 0..subquery_depth {
+                let new_query = Query::new_range_full();
+                current_query.default_subquery_branch = SubqueryBranch {
+                    subquery_path: None,
+                    subquery: Some(Box::new(new_query)),
+                };
+                current_query = current_query
+                    .default_subquery_branch
+                    .subquery
+                    .as_mut()
+                    .unwrap();
+            }
+
+            root_query
+        }
+
+        let query = create_non_recursive_query(100);
+
+        assert_eq!(query.max_depth(), Some(101));
+
+        let query = create_non_recursive_query(500);
+
+        assert_eq!(query.max_depth(), None);
+    }
 }
diff --git a/merk/src/proofs/query/mod.rs b/merk/src/proofs/query/mod.rs
index 669940cc..8bbb07e9 100644
--- a/merk/src/proofs/query/mod.rs
+++ b/merk/src/proofs/query/mod.rs
@@ -70,6 +70,36 @@ pub struct SubqueryBranch {
     pub subquery: Option<Box<Query>>,
 }
 
+impl SubqueryBranch {
+    /// Returns the depth of the subquery branch
+    /// This depth is how many GroveDB layers down we could query at maximum
+    #[inline]
+    pub fn max_depth(&self) -> Option<u16> {
+        self.max_depth_internal(u8::MAX)
+    }
+
+    /// Returns the depth of the subquery branch
+    /// This depth is how many GroveDB layers down we could query at maximum
+    #[inline]
+    fn max_depth_internal(&self, recursion_limit: u8) -> Option<u16> {
+        if recursion_limit == 0 {
+            return None;
+        }
+        let subquery_path_depth = self.subquery_path.as_ref().map_or(Some(0), |path| {
+            let path_len = path.len();
+            if path_len > u16::MAX as usize {
+                None
+            } else {
+                Some(path_len as u16)
+            }
+        })?;
+        let subquery_depth = self.subquery.as_ref().map_or(Some(0), |query| {
+            query.max_depth_internal(recursion_limit - 1)
+        })?;
+        subquery_path_depth.checked_add(subquery_depth)
+    }
+}
+
 #[cfg(any(feature = "full", feature = "verify"))]
 /// `Query` represents one or more keys or ranges of keys, which can be used to
 /// resolve a proof which will include all the requested values.
@@ -476,6 +506,33 @@ impl Query {
         // checks if all searched for items are keys
         self.items.iter().all(|a| a.is_key())
     }
+
+    /// Returns the depth of the subquery branch
+    /// This depth is how many GroveDB layers down we could query at maximum
+    pub fn max_depth(&self) -> Option<u16> {
+        self.max_depth_internal(u8::MAX)
+    }
+
+    /// Returns the depth of the subquery branch
+    /// This depth is how many GroveDB layers down we could query at maximum
+    pub(crate) fn max_depth_internal(&self, recursion_limit: u8) -> Option<u16> {
+        let default_subquery_branch_depth = self
+            .default_subquery_branch
+            .max_depth_internal(recursion_limit)?;
+        let conditional_subquery_branches_max_depth = self
+            .conditional_subquery_branches
+            .as_ref()
+            .map_or(Some(0), |condition_subqueries| {
+            condition_subqueries
+                .values()
+                .try_fold(0, |max_depth, conditional_subquery_branch| {
+                    conditional_subquery_branch
+                        .max_depth_internal(recursion_limit)
+                        .map(|depth| max_depth.max(depth))
+                })
+        })?;
+        1u16.checked_add(default_subquery_branch_depth.max(conditional_subquery_branches_max_depth))
+    }
 }
 
 #[cfg(feature = "full")]