feat(grovedb,query): allow SizedQuery::limit on carrier AggregateCountOnRange

grovedb/src/operations/get/query.rs

@@ -20,7 +20,7 @@ use crate::{
     query_result_type::{QueryResultElement, QueryResultElements, QueryResultType},
     reference_path::ReferencePathType,
     util::TxRef,
-    Element, Error, GroveDb, PathQuery, TransactionArg,
+    Element, Error, GroveDb, PathQuery, SizedQuery, TransactionArg,
 };
 use grovedb_costs::cost_return_on_error_default;
 #[cfg(feature = "minimal")]
@@ -772,8 +772,15 @@ where {
     ///
     /// `path_query` must satisfy
     /// [`PathQuery::validate_aggregate_count_on_range`] in either
-    /// shape. Pagination is rejected. Each leaf subtree the walk
-    /// terminates in must be a `ProvableCountTree` or
+    /// shape. Pagination rules differ by shape: for **leaf** queries
+    /// both `SizedQuery::limit` and `SizedQuery::offset` are rejected
+    /// (a leaf returns a single `u64` and pagination would silently
+    /// change the answer); for **carrier** queries `SizedQuery::limit`
+    /// is accepted and caps the number of outer-key matches the walk
+    /// returns (each matched outer key still produces a complete
+    /// leaf-ACOR `u64`, the inner range is not capped), while
+    /// `SizedQuery::offset` is still rejected. Each leaf subtree the
+    /// walk terminates in must be a `ProvableCountTree` or
     /// `ProvableCountSumTree` — the merk-level walk rejects any other
     /// tree type.
     ///
@@ -833,9 +840,18 @@ where {
         // outer items at `path_query.path` without descending into the
         // subquery — we want just the matched outer keys, not the
         // (unproven) results of the leaf aggregate-count.
+        //
+        // Propagate `SizedQuery::limit` (validated as carrier-only
+        // above): it caps the number of outer-key matches the walk
+        // returns. Each matched outer key still produces a complete
+        // leaf-ACOR `u64` below. `offset` is rejected at validation, so
+        // we don't propagate it here.
         let mut shallow_query = grovedb_query::Query::new_with_direction(left_to_right);
         shallow_query.items = outer_items;
-        let shallow_pq = PathQuery::new_unsized(path_query.path.clone(), shallow_query);
+        let shallow_pq = PathQuery::new(
+            path_query.path.clone(),
+            SizedQuery::new(shallow_query, path_query.query.limit, None),
+        );
 
         let (matched, _skipped) = cost_return_on_error!(
             &mut cost,

grovedb/src/operations/proof/aggregate_count/classification.rs

@@ -38,8 +38,12 @@ pub(super) struct AggregateCountClassification {
 }
 
 /// Classify an `AggregateCountOnRange` path query and validate it at
-/// the PathQuery level — `SizedQuery::limit` / `offset` (which
-/// aggregate-count explicitly forbids) are enforced for both shapes.
+/// the PathQuery level. The shape-specific pagination rules are
+/// enforced through [`PathQuery::validate_aggregate_count_on_range`]:
+/// leaf queries reject both `SizedQuery::limit` and
+/// `SizedQuery::offset`; carrier queries accept `SizedQuery::limit`
+/// (caps the outer walk; threaded into the proof verifier via
+/// `path_query.query.limit`) but still reject `SizedQuery::offset`.
 pub(super) fn classify_aggregate_count_path_query(
     path_query: &PathQuery,
 ) -> Result<AggregateCountClassification, Error> {

grovedb/src/operations/proof/aggregate_count/helpers.rs

@@ -169,10 +169,20 @@ pub(super) struct OuterMatch {
 /// returning `(carrier_merk_root_hash, matched_outer_keys)`. Each
 /// `OuterMatch` carries the value bytes and the parent-recorded value_hash
 /// that the chain check will validate.
+///
+/// `outer_limit` is the `SizedQuery::limit` that bounds the outer walk
+/// (matching what the prover passed to `Merk::prove_unchecked_query_items`
+/// when it generated the carrier-layer merk proof). When the carrier
+/// query carries a non-`None` `SizedQuery::limit`, the prover truncates
+/// the outer walk after that many matched keys and emits structural
+/// Hash nodes for the rest; the verifier must therefore execute the
+/// proof with the same limit so that its merk walker stops at the same
+/// boundary instead of demanding KV data for the un-walked tail.
 pub(super) fn execute_carrier_layer_proof(
     merk_bytes: &[u8],
     outer_items: &[QueryItem],
     left_to_right: bool,
+    outer_limit: Option<u16>,
     path_query: &PathQuery,
 ) -> Result<(CryptoHash, Vec<OuterMatch>), Error> {
     // The grovedb_query::QueryItem and grovedb_merk::proofs::query::QueryItem
@@ -187,7 +197,7 @@ pub(super) fn execute_carrier_layer_proof(
     // walker stops at the first out-of-order boundary and only the
     // last key in the proof is returned.
     let (root_hash, merk_result) = level_query
-        .execute_proof(merk_bytes, None, left_to_right, 0)
+        .execute_proof(merk_bytes, outer_limit, left_to_right, 0)
         .unwrap()
         .map_err(|e| {
             Error::InvalidProof(

grovedb/src/operations/proof/aggregate_count/per_key.rs

@@ -104,6 +104,11 @@ fn verify_v1_per_key(
             merk_bytes,
             path_query,
             outer_items,
+            // `SizedQuery::limit` (validated as carrier-only at entry) caps
+            // the outer walk. The prover truncates after this many outer
+            // matches; the verifier must apply the same cap so its merk
+            // walker stops at the same boundary.
+            path_query.query.limit,
             classification,
             grove_version,
         ),
@@ -114,18 +119,23 @@ fn verify_v1_per_key(
 /// matched outer key descend the `subquery_path` (if any) and the
 /// leaf count proof, enforcing the chain at each step. Returns one
 /// `(outer_key, count)` entry per match in query-direction order.
+///
+/// `outer_limit` is the carrier's `SizedQuery::limit` (when set, the
+/// outer walk stops after that many matched outer keys).
 fn verify_v1_carrier_layer(
     layer: &LayerProof,
     merk_bytes: &[u8],
     path_query: &PathQuery,
     outer_items: &[QueryItem],
+    outer_limit: Option<u16>,
     classification: &AggregateCountClassification,
     grove_version: &GroveVersion,
 ) -> Result<(CryptoHash, Vec<(Vec<u8>, u64)>), Error> {
     let (carrier_root, matched) = execute_carrier_layer_proof(
         merk_bytes,
         outer_items,
         classification.carrier_left_to_right,
+        outer_limit,
         path_query,
     )?;
 

grovedb/src/query/mod.rs

@@ -123,11 +123,31 @@ impl SizedQuery {
     /// `default_subquery_branch.subquery` for carrier queries).
     ///
     /// This is the `SizedQuery`-level entry point: it forwards to
-    /// [`Query::validate_aggregate_count_on_range`] and additionally rejects
-    /// any non-`None` `limit` or `offset` (counting is an aggregate over the
-    /// full match set — pagination would silently change the answer).
+    /// [`Query::validate_aggregate_count_on_range`] and additionally
+    /// enforces the appropriate per-shape size-constraint rules:
+    ///
+    /// - **Leaf** shape (single `AggregateCountOnRange(_)` item, no
+    ///   subqueries): both `SizedQuery::limit` and `SizedQuery::offset`
+    ///   are rejected. A leaf returns a single `u64`; pagination would
+    ///   silently change the answer.
+    /// - **Carrier** shape (outer `Key`/`Range*` items routing to a leaf
+    ///   `AggregateCountOnRange` subquery): `SizedQuery::limit` is
+    ///   **allowed** and caps the number of outer-key matches the
+    ///   carrier walks (each matched outer key still produces a complete
+    ///   leaf-ACOR `u64`). `SizedQuery::offset` is still rejected —
+    ///   skipping outer matches changes which `(outer_key, u64)` pairs
+    ///   end up in the proof, and the use case for that hasn't been
+    ///   designed yet.
     pub fn validate_aggregate_count_on_range(&self) -> Result<&QueryItem, Error> {
-        self.check_aggregate_count_size_constraints()?;
+        // Inner classification first, then per-shape size-constraint
+        // check. Queries that aren't aggregate-count at all (neither leaf
+        // nor carrier) fall through to the Query-level validator below,
+        // which surfaces the canonical "no aggregate-count item" error.
+        if self.query.aggregate_count_on_range().is_some() {
+            self.check_leaf_aggregate_count_size_constraints()?;
+        } else if self.query.has_aggregate_count_on_range_anywhere() {
+            self.check_carrier_aggregate_count_size_constraints()?;
+        }
         self.query
             .validate_aggregate_count_on_range()
             .map_err(query_validation_error_to_static_str)
@@ -137,24 +157,49 @@ impl SizedQuery {
     /// Strict variant of [`Self::validate_aggregate_count_on_range`] that
     /// only accepts the **leaf** shape (single `AggregateCountOnRange(_)`
     /// item, no subqueries). Used by entry points that produce a single
-    /// `u64` and need to reject the carrier shape up front.
+    /// `u64` and need to reject the carrier shape up front. Pagination
+    /// (`SizedQuery::limit` / `SizedQuery::offset`) is rejected — see
+    /// [`Self::check_leaf_aggregate_count_size_constraints`].
     pub fn validate_leaf_aggregate_count_on_range(&self) -> Result<&QueryItem, Error> {
-        self.check_aggregate_count_size_constraints()?;
+        self.check_leaf_aggregate_count_size_constraints()?;
         self.query
             .validate_leaf_aggregate_count_on_range()
             .map_err(query_validation_error_to_static_str)
             .map_err(Error::InvalidQuery)
     }
 
-    fn check_aggregate_count_size_constraints(&self) -> Result<(), Error> {
+    /// Size-constraint check used for **leaf** `AggregateCountOnRange`
+    /// queries. A leaf returns a single `u64`; setting `limit` or
+    /// `offset` would silently change the answer, so both are rejected.
+    fn check_leaf_aggregate_count_size_constraints(&self) -> Result<(), Error> {
         if self.limit.is_some() {
             return Err(Error::InvalidQuery(
-                "AggregateCountOnRange queries may not set SizedQuery::limit",
+                "leaf AggregateCountOnRange queries may not set SizedQuery::limit — a leaf \
+                 returns a single u64 and pagination would silently change the answer",
+            ));
+        }
+        if self.offset.is_some() {
+            return Err(Error::InvalidQuery(
+                "leaf AggregateCountOnRange queries may not set SizedQuery::offset — same \
+                 reason as limit",
             ));
         }
+        Ok(())
+    }
+
+    /// Size-constraint check used for **carrier** `AggregateCountOnRange`
+    /// queries. `SizedQuery::limit` is allowed and caps the number of
+    /// outer-key matches the carrier walks (each matched outer key still
+    /// produces a complete leaf-ACOR `u64`; the inner range is *not*
+    /// capped). `SizedQuery::offset` is still rejected — paginating into
+    /// the outer dimension changes which `(outer_key, u64)` pairs end up
+    /// in the proof, and the use case for that hasn't been designed yet.
+    fn check_carrier_aggregate_count_size_constraints(&self) -> Result<(), Error> {
         if self.offset.is_some() {
             return Err(Error::InvalidQuery(
-                "AggregateCountOnRange queries may not set SizedQuery::offset",
+                "carrier AggregateCountOnRange queries may not set SizedQuery::offset — \
+                 skipping outer matches changes which (outer_key, u64) pairs end up in the \
+                 proof; the use case for this isn't designed yet",
             ));
         }
         Ok(())
@@ -2500,7 +2545,11 @@ mod tests {
     // ---------- SizedQuery / PathQuery AggregateCountOnRange validation ----------
 
     #[test]
-    fn sized_query_validate_acor_rejects_limit() {
+    fn sized_query_validate_leaf_acor_rejects_limit_and_offset() {
+        // Leaf shape (single AggregateCountOnRange item, no subqueries):
+        // both SizedQuery::limit and SizedQuery::offset are rejected
+        // because a leaf returns a single u64 and pagination would
+        // silently change the answer.
         let mut sq = SizedQuery::new(
             Query::new_aggregate_count_on_range(QueryItem::Range(b"a".to_vec()..b"z".to_vec())),
             Some(10),
@@ -2510,7 +2559,10 @@ mod tests {
             .validate_aggregate_count_on_range()
             .expect_err("limit must fail");
         match err {
-            Error::InvalidQuery(msg) => assert!(msg.contains("limit")),
+            Error::InvalidQuery(msg) => {
+                assert!(msg.contains("leaf"), "unexpected message: {msg}");
+                assert!(msg.contains("limit"), "unexpected message: {msg}");
+            }
             _ => panic!("expected InvalidQuery"),
         }
 
@@ -2521,7 +2573,47 @@ mod tests {
             .validate_aggregate_count_on_range()
             .expect_err("offset must fail");
         match err {
-            Error::InvalidQuery(msg) => assert!(msg.contains("offset")),
+            Error::InvalidQuery(msg) => {
+                assert!(msg.contains("leaf"), "unexpected message: {msg}");
+                assert!(msg.contains("offset"), "unexpected message: {msg}");
+            }
+            _ => panic!("expected InvalidQuery"),
+        }
+    }
+
+    #[test]
+    fn sized_query_validate_carrier_acor_accepts_limit_rejects_offset() {
+        // Carrier shape (outer Key/Range items + AggregateCountOnRange
+        // subquery): SizedQuery::limit is permitted (caps the outer
+        // walk), but SizedQuery::offset is still rejected pending a
+        // separate design pass.
+        let mut carrier = Query::new();
+        carrier.insert_key(b"k1".to_vec());
+        carrier.default_subquery_branch = SubqueryBranch {
+            subquery_path: None,
+            subquery: Some(Box::new(Query::new_aggregate_count_on_range(
+                QueryItem::Range(b"a".to_vec()..b"z".to_vec()),
+            ))),
+        };
+        let mut sq = SizedQuery::new(carrier, Some(20), None);
+
+        // limit=Some(20) is now accepted on the carrier shape.
+        let inner = sq
+            .validate_aggregate_count_on_range()
+            .expect("carrier with limit must validate");
+        assert!(matches!(inner, QueryItem::Range(_)));
+
+        // offset is still rejected, with a carrier-specific message.
+        sq.limit = None;
+        sq.offset = Some(3);
+        let err = sq
+            .validate_aggregate_count_on_range()
+            .expect_err("carrier offset must fail");
+        match err {
+            Error::InvalidQuery(msg) => {
+                assert!(msg.contains("carrier"), "unexpected message: {msg}");
+                assert!(msg.contains("offset"), "unexpected message: {msg}");
+            }
             _ => panic!("expected InvalidQuery"),
         }
     }
@@ -2591,4 +2683,29 @@ mod tests {
         let pq_regular = PathQuery::new_single_key(vec![b"p".to_vec()], b"k".to_vec());
         assert!(!pq_regular.has_aggregate_count_on_range());
     }
+
+    #[test]
+    fn query_validation_error_to_static_str_projects_invalid_operation_and_catches_other_variants()
+    {
+        use grovedb_query::error::Error as QueryError;
+
+        // The expected normal case: `InvalidOperation(&'static str)` is
+        // projected through unchanged.
+        let normal = QueryError::InvalidOperation("specific reason");
+        assert_eq!(
+            super::query_validation_error_to_static_str(normal),
+            "specific reason"
+        );
+
+        // The defensive catch-all: any other QueryError variant gets the
+        // generic fallback label. This branch shouldn't be reachable from
+        // real `Query::validate_aggregate_count_on_range` results — it's
+        // here to surface "an unrelated bug" rather than silently turning
+        // into a useless empty string.
+        let other = QueryError::NotSupported("anything not InvalidOperation".to_string());
+        assert_eq!(
+            super::query_validation_error_to_static_str(other),
+            "AggregateCountOnRange query validation failed"
+        );
+    }
 }