drive/query/drive_document_ranked_query/

index_picker.rs

//! Covering-index picker for the ranked query, plus the shared
//! prefix-value encoding.
//!
//! Pure functions on the document type's index map plus the
//! `(group property, equality pins, axis, aggregate field)` tuple
//! [`super::mode_detection`] resolved. No Drive, no proof — the server
//! and the SDK verifier both call these so they land on the same index
//! (and therefore the same grove path) for the same request.

use super::{DocumentRankedMode, DriveDocumentRankedQuery, PrefixPin, RankedAxis};
use crate::error::query::QuerySyntaxError;
use crate::error::Error;
use crate::query::{index_admissible_for_resolved_time_range, ResolvedTimeRange};
use dpp::data_contract::document_type::methods::DocumentTypeV0Methods;
use dpp::data_contract::document_type::{DocumentTypeRef, Index};
use dpp::version::PlatformVersion;
use std::collections::BTreeMap;

/// Find the index that can serve `axis` ranking grouped by
/// `group_by_property` with the given equality pins, aggregating
/// `aggregate_field`.
///
/// An index qualifies when **all** of:
///
/// - one of its **ranked levels** for `axis` is `group_by_property`,
///   with exactly one pin per property before it. For the boolean axes
///   the ranked level is the index's **last** property; a prefix-level
///   `rankedCountable: { at }` hosts a Count secondary at the `at`
///   property as well, ranking its values by whole-subtree count — an
///   index may declare both, and the (group property, pin count) pair
///   singles out which level a request addresses. At a prefix level the
///   properties *after* `at` never appear in the request at all (they
///   are interior to the subtrees being counted);
/// - every property **before** the ranked level is pinned: each appears
///   (by name) among `equality_pin_fields`. Lengths matching plus the
///   pins being distinct (enforced upstream by
///   [`super::mode_detection::prefix_pins_from_where_clauses`]) makes
///   this set equality, so no pin is left over either;
/// - it declares the ranking keyword for `axis`
///   ([`RankedAxis::required_index_keyword`]);
/// - for [`RankedAxis::Sum`] / [`RankedAxis::Avg`], its `summable`
///   property is exactly `aggregate_field`. Both axes are derived from
///   the same running sum the index maintains (`Avg` is that sum over the
///   group's count), so summing a *different* field than the one the
///   index accumulates would silently answer about the wrong property.
///   A prefix-level Count ranking excludes both (the grammar rejects the
///   combination), so those arms never see an `at` index;
/// - it is admissible for the request's resolved time-range selections
///   ([`index_admissible_for_resolved_time_range`]): a request whose
///   leading pin was produced by `IN_TIME_RANGE` resolution may only be
///   served by the index bucketing that field with exactly that grid,
///   and a raw request never by a bucketed index — either mismatch
///   would be a validly-proven wrong answer. The resolved bucket-start
///   pin then descends the grid-qualified first level like any other
///   leading-property pin, into that window's own per-prefix secondary.
///
/// With no pins this degenerates to the original single-property rule —
/// or, for an index ranked at its FIRST property, to the global group
/// ranking ("top hashtags by total likes" with nothing pinned). A
/// partial pin (some but not all leading properties) matches nothing —
/// the per-prefix secondary lives under one value tree per leading
/// property, so there is no subtree an unpinned prefix could address —
/// and callers turn the `None` into a loud
/// [`crate::error::query::QuerySyntaxError`] naming what is missing.
///
/// Returns `None` when nothing qualifies; callers turn that into
/// [`crate::error::query::QuerySyntaxError::WhereClauseOnNonIndexedProperty`]
/// with a message naming the missing keyword.
///
/// At most one index can qualify for a given `(group property, pins,
/// axis, field)` tuple — rs-dpp rejects two indexes over the same
/// property set on one document type — so "first match wins" is not a
/// tie-break in practice. Should that ever change, the `BTreeMap`
/// iteration order (index name, ascending) keeps the choice
/// deterministic, which is what prover/verifier agreement actually
/// requires: both sides run this same function over the same contract
/// and must land on the same grove path.
///
/// Note that axis availability is decided from the index's `ranked_*`
/// flags, **not** from the element variant the write path laid down: a
/// `rankedCountable` index that also declares `rangeSummable` is stored
/// as a `ProvableCountProvableSumIndexedTree` carrying only the Count
/// axis, so the element variant alone would over-report what is rankable.
pub fn find_ranked_index_for_axis<'b>(
    indexes: &'b BTreeMap<String, Index>,
    group_by_property: &str,
    equality_pin_fields: &[String],
    axis: RankedAxis,
    aggregate_field: &str,
    resolved_time_ranges: &[ResolvedTimeRange],
) -> Option<&'b Index> {
    indexes.values().find(|index| {
        // Bucketed and raw indexes are never interchangeable, and one
        // grid's index never serves another grid's resolution — the same
        // provenance rule every other aggregate picker applies. This is
        // what keeps a raw request off bucketed indexes AND routes a
        // resolved request to exactly the grid it was resolved against.
        if !index_admissible_for_resolved_time_range(index, resolved_time_ranges) {
            return false;
        }
        // The positions whose levels host this axis's secondaries — for
        // the Count axis every `at` level plus the terminal when the
        // boolean is on (any subset of an index's levels may rank); empty
        // when the index does not declare the axis (or aggregates a
        // different field than requested).
        let candidate_positions: Vec<usize> = match axis {
            RankedAxis::Count => index
                .ranked_countable_at
                .iter()
                .filter_map(|at| index.properties.iter().position(|p| &p.name == at))
                .chain(
                    index
                        .ranked_countable
                        .then(|| index.properties.len().checked_sub(1))
                        .flatten(),
                )
                .collect(),
            RankedAxis::Sum => (index.ranked_summable
                && index.summable.as_deref() == Some(aggregate_field))
            .then(|| index.properties.len().checked_sub(1))
            .flatten()
            .into_iter()
            .collect(),
            RankedAxis::Avg => (index.ranked_averageable
                && index.summable.as_deref() == Some(aggregate_field))
            .then(|| index.properties.len().checked_sub(1))
            .flatten()
            .into_iter()
            .collect(),
        };
        // A candidate matches when its property is the grouping property
        // and every property before it is pinned exactly once (length
        // equality + distinct pins ⇒ set equality). At most one candidate
        // can match a given request: the two levels are distinct
        // positions, and the pin count singles one out.
        candidate_positions.into_iter().any(|ranked_position| {
            let Some(ranked_property) = index.properties.get(ranked_position) else {
                return false;
            };
            let leading = &index.properties[..ranked_position];
            ranked_property.name == group_by_property
                && leading.len() == equality_pin_fields.len()
                && leading
                    .iter()
                    .all(|property| equality_pin_fields.iter().any(|f| f == &property.name))
        })
    })
}

/// [`find_ranked_index_for_axis`] driven straight from a resolved
/// [`DocumentRankedMode`] — the shape every caller actually has.
pub fn find_ranked_index_for_mode<'b>(
    indexes: &'b BTreeMap<String, Index>,
    mode: &DocumentRankedMode,
    resolved_time_ranges: &[ResolvedTimeRange],
) -> Option<&'b Index> {
    let pin_fields: Vec<String> = mode
        .prefix_pins
        .iter()
        .map(|pin| pin.field.clone())
        .collect();
    find_ranked_index_for_axis(
        indexes,
        &mode.group_by_property,
        &pin_fields,
        mode.axis,
        &mode.aggregate_field,
        resolved_time_ranges,
    )
}

/// Resolve a validated [`DocumentRankedMode`] against a document type's
/// indexes into the executable [`DriveDocumentRankedQuery`]: pick the
/// covering index, encode the prefix pins into prefix **branches** (one
/// branch for all-`==` pins, one branch per element of the single
/// permitted `IN`), and assemble the query.
///
/// This is the **one** resolution path — the server's executors and the
/// SDK's proof helpers both call it, which is what guarantees a proof
/// and an unproven read (and the client's verification) are about the
/// same subtree.
///
/// `indexes` is threaded in separately rather than read off
/// `document_type` here because
/// [`DocumentTypeV0Getters::indexes`](dpp::data_contract::document_type::accessors::DocumentTypeV0Getters::indexes)
/// borrows its receiver — taking the map from the caller lets the
/// returned query's `&'a Index` outlive this frame. Callers pass
/// `document_type.indexes()`.
///
/// The main failure is "no index covers this", reported with the exact
/// contract keyword (and, for pinned requests, the exact index shape)
/// the request needs, so the caller can act on it without reading the
/// schema spec.
pub fn resolve_ranked_query_for_mode<'a>(
    contract_id: [u8; 32],
    document_type: DocumentTypeRef<'a>,
    document_type_name: String,
    indexes: &'a BTreeMap<String, Index>,
    mode: &DocumentRankedMode,
    resolved_time_ranges: &[ResolvedTimeRange],
    platform_version: &PlatformVersion,
) -> Result<DriveDocumentRankedQuery<'a>, Error> {
    let index =
        find_ranked_index_for_mode(indexes, mode, resolved_time_ranges).ok_or_else(|| {
            Error::Query(QuerySyntaxError::WhereClauseOnNonIndexedProperty(
                no_covering_index_message(
                    "ranked",
                    mode.axis,
                    &mode.group_by_property,
                    &mode.prefix_pins,
                    &mode.aggregate_field,
                ),
            ))
        })?;
    let prefix_branches =
        encode_prefix_branches(document_type, index, &mode.prefix_pins, platform_version)?;
    Ok(DriveDocumentRankedQuery {
        document_type,
        contract_id,
        document_type_name,
        index,
        prefix_branches,
        axis: mode.axis,
        descending: mode.descending,
        k: mode.k,
        offset: mode.offset,
    })
}

/// The "no index covers this request" rejection text, shared by the
/// ranked and having-range resolutions (and the SDK's mirrors of them)
/// so a rejected request reads identically everywhere. Names the exact
/// index the request needs: property list (pins first, in request
/// order, then the grouping property), ranking keyword, and `summable`
/// field where applicable.
pub fn no_covering_index_message(
    surface: &str,
    axis: RankedAxis,
    group_by_property: &str,
    prefix_pins: &[PrefixPin],
    aggregate_field: &str,
) -> String {
    let pin_fields = || {
        prefix_pins
            .iter()
            .map(|pin| pin.field.as_str())
            .collect::<Vec<_>>()
            .join(", ")
    };
    let index_shape = if prefix_pins.is_empty() {
        format!("a single-property index on `{group_by_property}`")
    } else {
        format!(
            "a compound index on [{}, {group_by_property}] (every leading property pinned \
             by an equality or `IN` `where` clause, the trailing property grouped over)",
            pin_fields()
        )
    };
    format!(
        "no ranked index covers `group_by = [{group_by_property}]`{} on the {axis:?} axis \
         for this {surface} query: the document type needs {index_shape} declaring `{}`{}",
        if prefix_pins.is_empty() {
            String::new()
        } else {
            format!(" with pins on [{}]", pin_fields())
        },
        axis.required_index_keyword(),
        if aggregate_field.is_empty() {
            String::new()
        } else {
            format!(" with `summable: \"{aggregate_field}\"`")
        }
    )
}

/// Encode the resolved prefix pins into **branches** — one
/// `Vec<Vec<u8>>` of prefix path segments per branch, in index-property
/// order (the same order and encoding the write path used to key those
/// prefix value trees). A request with only `==` pins yields exactly
/// one branch; the (at most one) `IN` pin yields one branch per
/// element.
///
/// This is part of the prover/verifier agreement: server executors and
/// the SDK's proof helpers both come through here, so a pinned value
/// can only ever name one subtree — and a branch *set* only ever one
/// ordered subtree list — identically on both sides. Branch order is
/// canonical: ascending by encoded segment bytes, independent of the
/// caller's element order (which also makes `null`, the empty segment,
/// sort first deterministically).
///
/// `index` must have been picked by [`find_ranked_index_for_axis`]
/// against these same pins — every leading property is then guaranteed
/// a pin. A value the property's type cannot encode is a caller error
/// naming the property; two `IN` elements that encode to the same
/// segment (two spellings of one value) are one branch and are rejected
/// as a duplicate rather than walked twice.
pub fn encode_prefix_branches(
    document_type: DocumentTypeRef,
    index: &Index,
    prefix_pins: &[PrefixPin],
    platform_version: &PlatformVersion,
) -> Result<Vec<Vec<Vec<u8>>>, Error> {
    // The pinnable properties end at the ranked level the pin count
    // addresses (an index may host secondaries at both its `at` property
    // and its terminal — the pin count singles one out; properties past
    // it are interior to the counted subtrees and never pinned).
    let (leading, _) = super::path::ranked_level_split(index, prefix_pins.len())?;
    // Enforced BEFORE any encoding: the ceiling bounds every downstream
    // cost (encode, sort, clone, walk, proof size), so an oversized pin
    // must not buy that work first. The post-product branch count check
    // below stays as a backstop.
    if prefix_pins
        .iter()
        .any(|pin| pin.values.len() > super::MAX_PREFIX_IN_BRANCHES)
    {
        return Err(Error::Query(
            QuerySyntaxError::InvalidWhereClauseComponents(
                "an `IN` prefix pin fans out into more branches than the ranked surface serves \
             — narrow the element list or issue several requests",
            ),
        ));
    }
    let per_property: Vec<Vec<Vec<u8>>> = leading
        .iter()
        .map(|property| {
            let pin = prefix_pins
                .iter()
                .find(|pin| pin.field == property.name)
                .ok_or_else(|| {
                    Error::Query(QuerySyntaxError::InvalidWhereClauseComponents(
                        "internal resolution mismatch: the picked compound ranked index has \
                         a leading property with no pin — the index picker and the prefix \
                         encoder disagreed on the pins",
                    ))
                })?;
            let mut encoded = pin
                .values
                .iter()
                .map(|value| {
                    // A null pin addresses the subtree the write walkers
                    // create for an **absent** value: they encode it as
                    // `get_raw_for_document_type(..).unwrap_or_default()`
                    // — an empty path segment — for user and system
                    // properties alike. Null must short-circuit here
                    // because the system-property encoders (`$updatedAt`,
                    // `$creatorId`, …) reject null before any encoding
                    // happens, which would make the stored empty-segment
                    // prefix unaddressable.
                    if value.is_null() {
                        return Ok(Vec::new());
                    }
                    document_type
                        .serialize_value_for_key(&property.name, value, platform_version)
                        .map_err(|e| {
                            Error::Query(QuerySyntaxError::InvalidParameter(format!(
                                "the pin on `{}` does not encode as that property's \
                                 index key: {e}",
                                property.name
                            )))
                        })
                })
                .collect::<Result<Vec<_>, Error>>()?;
            if encoded.len() > 1 {
                encoded.sort();
                if encoded.windows(2).any(|pair| pair[0] == pair[1]) {
                    return Err(Error::Query(
                        QuerySyntaxError::InvalidWhereClauseComponents(
                            "an `IN` pin's elements encode to the same index key: two \
                             spellings of one value are one prefix branch — deduplicate \
                             the element list",
                        ),
                    ));
                }
            }
            Ok(encoded)
        })
        .collect::<Result<Vec<_>, Error>>()?;

    // Defense in depth at the shared choke point: the grammar enforces
    // both invariants upstream, but this function is `pub` and the
    // prover/verifier agreement hangs off it, so a mis-built pin set
    // must fail here rather than collapse to zero branches (a
    // downstream panic) or fan out into an unbounded cartesian product
    // (which would also break the one-varying-position assumption
    // `in_key` and the merge order rely on).
    if per_property.iter().any(|candidates| candidates.is_empty()) {
        return Err(Error::Query(
            QuerySyntaxError::InvalidWhereClauseComponents(
                "internal resolution mismatch: a prefix pin carries no values",
            ),
        ));
    }
    if per_property
        .iter()
        .filter(|candidates| candidates.len() > 1)
        .count()
        > 1
    {
        return Err(Error::Query(
            QuerySyntaxError::InvalidWhereClauseComponents(
                "internal resolution mismatch: more than one branching pin — the grammar \
             admits at most one `IN` across the prefix properties",
            ),
        ));
    }

    // A single `null` pin encodes as the empty path segment; the branched
    // proof grammar (`PathQuery::new_branched_axis`) cannot address an
    // empty segment in the shared prefix or suffix, so a null `==` pin
    // combined with an `IN` would serve the unproved read and fail the
    // prove — the exact proved/unproved divergence this surface forbids.
    // Rejected for any non-branching pin position, conservatively: issue
    // one request per `IN` element to combine null pins with multiple
    // prefixes. `null` as an ELEMENT of the `IN` itself stays legal — it
    // is a branch key, which the envelope addresses and authenticates
    // like any other.
    let has_branching_pin = per_property.iter().any(|candidates| candidates.len() > 1);
    if has_branching_pin
        && per_property
            .iter()
            .any(|candidates| candidates.len() == 1 && candidates[0].is_empty())
    {
        return Err(Error::Query(
            QuerySyntaxError::InvalidWhereClauseComponents(
                "an `IN` prefix pin cannot be combined with a `null` pin: null addresses the \
             absent-value prefix through an empty path segment, which the branched proof \
             cannot express — issue one request per `IN` element instead",
            ),
        ));
    }

    // The grammar admits at most one multi-value pin, so this product
    // is |IN| branches (or exactly one), already in canonical order
    // because the only varying position was sorted above.
    let mut branches: Vec<Vec<Vec<u8>>> = vec![Vec::with_capacity(leading.len())];
    for candidates in per_property {
        branches = branches
            .into_iter()
            .flat_map(|prefix| {
                candidates.iter().map(move |segment| {
                    let mut branch = prefix.clone();
                    branch.push(segment.clone());
                    branch
                })
            })
            .collect();
    }
    // The documented hard ceiling on branch fan-out, enforced at the
    // shared choke point too: this function is `pub`, and everything
    // downstream (encoding, sorting, per-branch walks, proof size) is
    // linear in the branch count.
    if branches.len() > super::MAX_PREFIX_IN_BRANCHES {
        return Err(Error::Query(
            QuerySyntaxError::InvalidWhereClauseComponents(
                "an `IN` prefix pin fans out into more branches than the ranked surface serves \
             — narrow the element list or issue several requests",
            ),
        ));
    }
    Ok(branches)
}