laburnum 1.17.3

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  crate::{
    ContentHash,
    database::{
      Database, DynPartition, GenerationEpoch, HasPartition, Partition,
      Partitions, RecordRef,
      query::{
        PartitionQueryBuilder, QueryBuilder, TypedPartitionQueryBuilder,
      },
      query_results::{QueryResults, RecordMetadata},
    },
    prelude::PartitionKey,
  },
  std::sync::Arc,
};

/// Snapshot-isolated read interface for the compilation database.
///
/// `QueryClient` is the single public API for reading data from the database.
/// It provides three categories of read access:
///
/// # Typed index entry queries
///
/// For partitions with rich index entries (e.g. `SymbolEntry`, `ResolutionEntry`),
/// use the typed index methods to get the full `Part::IndexEntry`:
///
/// ```ignore
/// // Exact lookup
/// let entry: Option<Part::IndexEntry> = client.index_get::<MyPartition>("sort_key");
///
/// // Prefix scan
/// let entries: Vec<(String, Part::IndexEntry)> = client.index_range::<MyPartition>("prefix");
///
/// // Range queries
/// let entries = client.index_less_than::<MyPartition>("value", true);
/// let entries = client.index_greater_than::<MyPartition>("value", false);
/// let entries = client.index_between::<MyPartition>("from", "to");
/// ```
///
/// # Typed CAS record lookup
///
/// To retrieve a content-addressed record by its hash (e.g. following a
/// `RecordHandle` from an index entry):
///
/// ```ignore
/// let record: Option<RecordRef<'_, MyPartition>> = client.get::<MyPartition>(content_hash);
/// ```
///
/// # Query builder API
///
/// For partitions with CAS records, the query builder returns `QueryResults`
/// which resolves index entries to their underlying records:
///
/// ```ignore
/// let results = client.query(MyPartition)
///     .sort_key_begins_with(prefix)
///     .execute()
///     .await;
///
/// for record_ref in results.iter() {
///     // work with resolved P::RecordRef
/// }
/// ```
///
/// Use `query_partition()` for standard partitions (`DynPartition` types)
/// and `query()` for custom partitions (`Partition` types).
#[derive(Clone)]
pub struct QueryClient<P: Partitions> {
  pub(crate) db: Database<P>,
  snapshot_epoch: GenerationEpoch,
}

impl<P: Partitions> QueryClient<P> {
  pub fn new(db: Database<P>) -> Self {
    let snapshot_epoch = db.get_current_epoch();
    Self { db, snapshot_epoch }
  }

  pub(crate) fn query_any<'a>(
    &'a mut self,
    partition_key: crate::SpannedIdent,
  ) -> QueryBuilder<'a, P> {
    QueryBuilder::new(self, partition_key)
  }

  pub fn query_partition<'a, Part: DynPartition + PartitionKey>(
    &'a mut self,
    _partition: Part,
  ) -> PartitionQueryBuilder<'a, P, Part> {
    PartitionQueryBuilder::new(self)
  }

  /// Exact lookup against a dynamic partition by its `DynSortKey`.
  ///
  /// Resolves the dyn key into `P::SortKey` via runtime downcast (ADR0011);
  /// returns no records if the partition is not in `P`.
  pub async fn get_record_dyn<D: DynPartition + PartitionKey>(
    &mut self,
    _partition: D,
    key: D::DynSortKey,
  ) -> QueryResults<P> {
    match P::wrap_dyn_sort_key(<D as PartitionKey>::KEY, Box::new(key)) {
      | Some(k) => {
        self
          .get_record_internal(<D as PartitionKey>::KEY, Some(k))
          .await
      },
      | None => QueryResults::new(Vec::new(), Arc::clone(&self.db.cas)),
    }
  }

  pub fn query<
    'a,
    Part: Partition + PartitionKey + crate::database::partitions::SortKeyOf<P>,
  >(
    &'a mut self,
    _partition: Part,
  ) -> TypedPartitionQueryBuilder<'a, P, Part> {
    TypedPartitionQueryBuilder::new(self)
  }

  pub async fn get_record(
    &mut self,
    partition_key: crate::SpannedIdent,
    sort_key: P::SortKey,
  ) -> QueryResults<P> {
    self
      .get_record_internal(partition_key, Some(sort_key))
      .await
  }

  /// Query the span index for the record at a given byte offset in a URI.
  ///
  /// Returns the innermost record whose indexed span contains the given
  /// byte offset, or `None` if no record covers that position.
  pub fn span_index_get<D: PartitionKey>(
    &self,
    uri: &crate::Uri,
    byte_offset: u64,
  ) -> Option<P::RecordRef<'_>> {
    let hash = self.db.span_index_query(D::KEY, uri, byte_offset)?;
    self.db.cas.get_any(D::KEY, hash)
  }

  /// Query the span index for the raw content hash at a byte offset.
  ///
  /// Unlike [`span_index_get`], this returns the staged `ContentHash` itself
  /// rather than resolving it against partition `D`'s CAS. Use it when the
  /// span index of `D` stores a handle into a *different* partition's CAS
  /// (e.g. the `FileSymbols` span index stores `Symbols`-CAS handles), then
  /// resolve it with [`get_by_hash`] against the owning partition.
  pub fn span_index_hash<D: PartitionKey>(
    &self,
    uri: &crate::Uri,
    byte_offset: u64,
  ) -> Option<ContentHash> {
    self.db.span_index_query(D::KEY, uri, byte_offset)
  }

  /// Get a CAS record by content hash from a specific partition.
  ///
  /// Use this to resolve a `ContentHash` obtained from one partition's record
  /// into a record in another partition. The type parameter `D` provides the
  /// target partition key at compile time.
  pub fn get_by_hash<D: PartitionKey>(
    &self,
    hash: ContentHash,
  ) -> Option<P::RecordRef<'_>> {
    self.db.cas.get_any(D::KEY, hash)
  }

  pub async fn batch_get_items(
    &mut self,
    items: Vec<(crate::SpannedIdent, P::SortKey)>,
  ) -> QueryResults<P> {
    let mut all_records = Vec::new();

    for (partition_key, sort_key) in items {
      let results = self
        .get_record_internal(partition_key, Some(sort_key))
        .await;
      all_records.extend(results.records);
    }

    all_records.sort_by(|a, b| a.sort_key.cmp(&b.sort_key));

    QueryResults::new(all_records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn get_record_internal(
    &mut self,
    partition_key: crate::SpannedIdent,
    sort_key: Option<P::SortKey>,
  ) -> QueryResults<P> {
    match sort_key {
      | Some(sk) => match self.db.index_get(partition_key, &sk) {
        | Some(content_hash) => {
          let records = vec![RecordMetadata {
            partition_key,
            sort_key: sk,
            content_hash,
          }];
          QueryResults::new(records, Arc::clone(&self.db.cas))
        },
        | None => QueryResults::new(Vec::new(), Arc::clone(&self.db.cas)),
      },
      | None => {
        let entries = self.db.index_all(partition_key);
        if entries.is_empty() {
          return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
        }

        let records: Vec<RecordMetadata<P>> = entries
          .into_iter()
          .map(|(_, sort_key, content_hash)| RecordMetadata {
            partition_key,
            sort_key,
            content_hash,
          })
          .collect();

        QueryResults::new(records, Arc::clone(&self.db.cas))
      },
    }
  }

  pub(crate) async fn less_than_internal(
    &mut self,
    partition_key: crate::SpannedIdent,
    value: P::SortKey,
    inclusive: bool,
  ) -> QueryResults<P> {
    let entries = self.db.index_less_than(partition_key, &value, inclusive);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata<P>> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn greater_than_internal(
    &mut self,
    partition_key: crate::SpannedIdent,
    value: P::SortKey,
    inclusive: bool,
  ) -> QueryResults<P> {
    let entries = self.db.index_greater_than(partition_key, &value, inclusive);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata<P>> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn between_internal(
    &mut self,
    partition_key: crate::SpannedIdent,
    from: P::SortKey,
    to: P::SortKey,
  ) -> QueryResults<P> {
    let entries = self.db.index_between(partition_key, &from, &to);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata<P>> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn prefix_internal(
    &mut self,
    partition_key: crate::SpannedIdent,
    prefix: P::SortKey,
  ) -> QueryResults<P> {
    let entries = self.db.index_range(partition_key, &prefix);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata<P>> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  // -- Typed index entry queries ----------------------------------------

  /// Get an index entry by exact sort key, returning the full `Part::IndexEntry`.
  pub fn index_get<Part>(
    &self,
    sort_key: &Part::SortKey,
  ) -> Option<Part::IndexEntry>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_get(sort_key)
  }

  /// Get all index entries whose sort key has `prefix` as a prefix.
  pub fn index_range<Part>(
    &self,
    prefix: &Part::SortKey,
  ) -> Vec<(Part::SortKey, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_range(prefix)
  }

  /// Get all index entries with sort key less than (or equal to) `value`.
  pub fn index_less_than<Part>(
    &self,
    value: &Part::SortKey,
    inclusive: bool,
  ) -> Vec<(Part::SortKey, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_less_than(value, inclusive)
  }

  /// Get all index entries with sort key greater than (or equal to) `value`.
  pub fn index_greater_than<Part>(
    &self,
    value: &Part::SortKey,
    inclusive: bool,
  ) -> Vec<(Part::SortKey, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_greater_than(value, inclusive)
  }

  /// Get all index entries with sort key between `from` and `to` (inclusive).
  pub fn index_between<Part>(
    &self,
    from: &Part::SortKey,
    to: &Part::SortKey,
  ) -> Vec<(Part::SortKey, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_between(from, to)
  }

  /// Get every index entry for a partition (typed "list all").
  pub fn index_entries<Part>(&self) -> Vec<(Part::SortKey, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_entries()
  }

  // -- Typed CAS record lookup ------------------------------------------

  /// Get a CAS record by content hash from a specific partition.
  ///
  /// Use this to follow a `RecordHandle` from an index entry to the
  /// underlying record.
  pub fn get<Part>(&self, hash: ContentHash) -> Option<RecordRef<'_, Part>>
  where
    Part: Partition + 'static,
    P::Stores: HasPartition<Part>,
  {
    self.db.cas.get::<Part>(hash)
  }

  pub(crate) fn refresh_snapshot(&mut self) {
    self.snapshot_epoch = self.db.get_current_epoch();
  }
}