laburnum 1.17.1

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

use {
  crate::{
    ContentHash, Ident,
    database::{GenerationEpoch, MergeResult},
    hash::ContentHasher,
    record::LaburnumRecordRef,
  },
  serde::Serialize,
  std::{fmt::Debug, hash::Hash, sync::Arc},
};

/// Content-addressed storage where laburnum owns all partition stores.
///
/// This trait is the successor to [`Partitions`], designed for content-addressed
/// storage with structural sharing. Languages no longer implement storage logic -
/// they just declare:
/// 1. Partition types (via existing `Partition` trait)
/// 2. A `Partitions` type (HList of their partitions)
/// 3. A `RecordRef` enum for downcasting
///
/// Laburnum provides all DashMap storage, indexing, and GC.
///
/// # Example
///
/// ```ignore
/// use laburnum::define_partitions;
///
/// define_partitions! {
///     MyPartitions,
///     record_ref = MyRecordRef,
///     partitions = [CstPartition, AstPartition, SymbolsPartition],
/// }
///
/// type MyLaburnum = Laburnum<MyPartitions, MyLanguageServer>;
/// ```
pub trait Partitions: Send + Sync + 'static {
  /// The HList of PartitionStore<P> for each partition.
  ///
  /// Generated by the `partition_stores!` macro.
  type Stores: Send + Sync + crate::database::partitions::RefcountOps + crate::database::partitions::CollectCascadingRefs<Self>;

  /// Unified record reference for downcasting.
  ///
  /// This is an enum over all record types in all partitions.
  type RecordRef<'a>: LaburnumRecordRef
  where
    Self: 'a;

  /// Create empty stores.
  ///
  /// Generated by the `new_partition_stores!` macro.
  fn new_stores() -> Self::Stores;

  /// Dynamic dispatch for framework code.
  ///
  /// Look up a record by partition key and content hash. Returns None if
  /// not found or if the partition key is unknown.
  fn get_any(
    stores: &Self::Stores,
    partition: Ident,
    hash: ContentHash,
  ) -> Option<Self::RecordRef<'_>>;

  /// List all partition keys (for GC).
  ///
  /// Returns a static slice of all partition keys this Partitions type knows about.
  fn partition_keys() -> &'static [Ident];

  /// Merge records from source stores into the global store.
  ///
  /// This is used during chunk commit to merge all partition records
  /// into the global store with dynamic dispatch.
  ///
  /// Returns a [`MergeResult`] containing the set of content hashes that were
  /// newly inserted into the global store.
  fn merge_stores(
    global_stores: &Self::Stores,
    source: &Self::Stores,
    epoch: GenerationEpoch,
  ) -> MergeResult;

  /// Collect all content hashes from all partition indexes (GC roots).
  ///
  /// Each hash is tagged with its partition key so the GC can route
  /// items to the correct partition store during marking.
  fn collect_index_hashes(stores: &Self::Stores) -> Vec<crate::database::ContentHashRef>;

  /// Remove all index entries matching a prefix in the specified partition.
  ///
  /// Used for scoped deletion when re-processing a file.
  fn index_remove_prefix(
    stores: &Self::Stores,
    partition_key: Ident,
    prefix: &str,
  );

  /// Insert an entry into the index of the specified partition.
  fn index_insert(
    stores: &Self::Stores,
    partition_key: Ident,
    sort_key: String,
    content_hash: ContentHash,
  );

  /// Query index entries matching a prefix in the specified partition.
  ///
  /// Returns (partition_key, sort_key, content_hash) tuples.
  fn index_range(
    stores: &Self::Stores,
    partition_key: Ident,
    prefix: &str,
  ) -> Vec<(Ident, String, ContentHash)>;

  /// Get a content hash from the index by partition and sort key.
  fn index_get(
    stores: &Self::Stores,
    partition_key: Ident,
    sort_key: &str,
  ) -> Option<ContentHash>;

  /// Query index entries with sort key less than (or equal to) a value.
  fn index_less_than(
    stores: &Self::Stores,
    partition_key: Ident,
    value: &str,
    inclusive: bool,
  ) -> Vec<(Ident, String, ContentHash)>;

  /// Query index entries with sort key greater than (or equal to) a value.
  fn index_greater_than(
    stores: &Self::Stores,
    partition_key: Ident,
    value: &str,
    inclusive: bool,
  ) -> Vec<(Ident, String, ContentHash)>;

  /// Query index entries with sort key between two values (inclusive).
  fn index_between(
    stores: &Self::Stores,
    partition_key: Ident,
    from: &str,
    to: &str,
  ) -> Vec<(Ident, String, ContentHash)>;

  /// Replace the span index for a URI in the specified partition.
  fn span_index_replace(
    stores: &Self::Stores,
    partition_key: Ident,
    uri: crate::Uri,
    index: crate::database::partitions::span_index::SpanIndex,
  );

  /// Remove the span index for a URI in the specified partition.
  fn span_index_remove(
    stores: &Self::Stores,
    partition_key: Ident,
    uri: &crate::Uri,
  );

  /// Query the span index for a URI at a given byte offset.
  ///
  /// Returns the content hash of the innermost span containing the offset.
  fn span_index_query(
    stores: &Self::Stores,
    partition_key: Ident,
    uri: &crate::Uri,
    byte_offset: u64,
  ) -> Option<ContentHash>;

  /// Run one mark tick of the garbage collector across all partition types.
  ///
  /// Returns `true` when the gray queue is empty (marking complete).
  fn gc_mark_tick(
    gc: &crate::database::gc::GarbageCollector,
    stores: &Self::Stores,
    budget: usize,
  ) -> bool;

  /// Run the sweep phase of the garbage collector across all partition types.
  fn gc_sweep(
    gc: &crate::database::gc::GarbageCollector,
    stores: &Self::Stores,
  );
}

/// Internal content-addressed storage for all partition records.
///
/// This type is not part of the public API. All read access goes through
/// [`QueryClient`](crate::database::query::QueryClient), which provides
/// typed index entry queries and CAS record lookups. All write access goes
/// through [`PartitionWriteContextRef`](crate::database::PartitionWriteContextRef).
///
/// Records are content-addressed by `ContentHash` and include an epoch
/// timestamp for snapshot isolation. Records with identical content share
/// the same storage location (structural sharing).
pub struct ContentAddressedStorage<P: Partitions> {
  stores: Arc<P::Stores>,
}

impl<P: Partitions> ContentAddressedStorage<P> {
  pub fn new() -> Self {
    Self {
      stores: Arc::new(P::new_stores()),
    }
  }

  /// Get a shared reference to the stores Arc.
  ///
  /// Used by the reaper and garbage collector which need shared
  /// ownership of the stores.
  pub(crate) fn stores_arc(&self) -> Arc<P::Stores> {
    Arc::clone(&self.stores)
  }

  /// Get a record by content hash from a specific partition.
  ///
  /// Returns a `RecordRef` guard that provides access to the record
  /// without cloning.
  pub(crate) fn get<Part>(
    &self,
    hash: ContentHash,
  ) -> Option<crate::database::RecordRef<'_, Part>>
  where
    Part: crate::database::partitions::Partition + 'static,
    P::Stores: crate::database::partitions::HasPartition<Part>,
  {
    use crate::database::partitions::HasPartition;
    <P::Stores as HasPartition<Part>>::store(&self.stores).get(&hash)
  }

  /// Get a record by partition key and hash, returning a dynamic RecordRef.
  ///
  /// This provides dynamic dispatch for framework code that doesn't know
  /// the partition type at compile time.
  pub(crate) fn get_any(
    &self,
    partition: Ident,
    hash: ContentHash,
  ) -> Option<P::RecordRef<'_>> {
    P::get_any(&self.stores, partition, hash)
  }

  pub(crate) fn stores(&self) -> &P::Stores {
    &self.stores
  }

  /// Collect all content hashes from all partition indexes (GC roots).
  pub(crate) fn collect_index_hashes(&self) -> Vec<crate::database::ContentHashRef> {
    P::collect_index_hashes(&self.stores)
  }

  /// Remove all index entries matching a prefix in the specified partition.
  pub(in crate::database) fn index_remove_prefix(&self, partition_key: Ident, prefix: &str) {
    P::index_remove_prefix(&self.stores, partition_key, prefix);
  }

  pub(crate) fn index_range(
    &self,
    partition_key: Ident,
    prefix: &str,
  ) -> Vec<(Ident, String, ContentHash)> {
    P::index_range(&self.stores, partition_key, prefix)
  }

  pub(crate) fn index_get(
    &self,
    partition_key: Ident,
    sort_key: &str,
  ) -> Option<ContentHash> {
    P::index_get(&self.stores, partition_key, sort_key)
  }

  pub(crate) fn index_less_than(
    &self,
    partition_key: Ident,
    value: &str,
    inclusive: bool,
  ) -> Vec<(Ident, String, ContentHash)> {
    P::index_less_than(&self.stores, partition_key, value, inclusive)
  }

  pub(crate) fn index_greater_than(
    &self,
    partition_key: Ident,
    value: &str,
    inclusive: bool,
  ) -> Vec<(Ident, String, ContentHash)> {
    P::index_greater_than(&self.stores, partition_key, value, inclusive)
  }

  pub(crate) fn index_between(
    &self,
    partition_key: Ident,
    from: &str,
    to: &str,
  ) -> Vec<(Ident, String, ContentHash)> {
    P::index_between(&self.stores, partition_key, from, to)
  }

  /// Replace the span index for a URI in the specified partition.
  pub(in crate::database) fn span_index_replace(
    &self,
    partition_key: Ident,
    uri: crate::Uri,
    index: crate::database::partitions::span_index::SpanIndex,
  ) {
    P::span_index_replace(&self.stores, partition_key, uri, index);
  }

  /// Query the span index for a URI at a given byte offset.
  pub(crate) fn span_index_query(
    &self,
    partition_key: Ident,
    uri: &crate::Uri,
    byte_offset: u64,
  ) -> Option<crate::ContentHash> {
    P::span_index_query(&self.stores, partition_key, uri, byte_offset)
  }

  /// Increment the refcount for a record in the specified partition.
  pub(in crate::database) fn increment_refcount(&self, partition_key: Ident, hash: ContentHash)
  where
    P::Stores: crate::database::partitions::RefcountOps,
  {
    use crate::database::partitions::RefcountOps;
    self.stores.increment_refcount(partition_key, hash);
  }

  /// Decrement the refcount for a record in the specified partition.
  ///
  /// Returns the new refcount (0 means the record can be removed).
  #[cfg_attr(not(test), allow(dead_code))]
  pub(in crate::database) fn decrement_refcount(&self, partition_key: Ident, hash: ContentHash) -> usize
  where
    P::Stores: crate::database::partitions::RefcountOps,
  {
    use crate::database::partitions::RefcountOps;
    self.stores.decrement_refcount(partition_key, hash)
  }

  /// Get the refcount for a record in the specified partition.
  #[cfg_attr(not(test), allow(dead_code))]
  pub(crate) fn get_refcount(&self, partition_key: Ident, hash: ContentHash) -> usize
  where
    P::Stores: crate::database::partitions::RefcountOps,
  {
    use crate::database::partitions::RefcountOps;
    self.stores.get_refcount(partition_key, hash)
  }
}

impl<P: Partitions> Default for ContentAddressedStorage<P> {
  fn default() -> Self {
    Self::new()
  }
}

impl<P: Partitions> std::fmt::Debug for ContentAddressedStorage<P> {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.debug_struct("ContentAddressedStorage")
      .finish_non_exhaustive()
  }
}

/// Immutable storage for compilation records organized by index.
///
/// Each chunk contains a `RecordStorage` that holds the actual data. Records
/// are stored internally and retrieved by an opaque `Index` type. The storage
/// must be content-hashable for chunk content addressing.
///
/// # Why This Trait?
///
/// The `RecordStorage` trait abstracts over different storage implementations,
/// allowing language servers to choose:
/// - **In-memory**: `Vec<Record>` (simple, fast)
/// - **Persistent**: Disk-backed storage for large projects
/// - **Compressed**: Packed representations to reduce memory
/// - **Distributed**: Remote storage for multi-machine compilation
///
/// # Implementation Requirements
///
/// - `Index`: Unique identifier for each stored record (e.g., `usize`, enum)
/// - `RecordRef<'a>`: Borrowed view of a record, must implement
///   [`LaburnumRecordRef`]
/// - `Builder`: Associated builder type for constructing the storage
///
/// # Example
///
/// ```ignore
/// pub struct MyStorage {
///   records: Vec<MyRecord>,
/// }
///
/// impl RecordStorage for MyStorage {
///   type Index = usize;
///   type RecordRef<'a> = &'a MyRecord;
///   type Builder = MyStorageBuilder;
///
///   fn get(&self, idx: &Self::Index) -> Option<Self::RecordRef<'_>> {
///     self.records.get(*idx)
///   }
///
///   fn hash_contents<H: Hasher>(&self, state: &mut H) {
///     self.records.hash(state);
///   }
/// }
/// ```
pub trait RecordStorage: Debug + Send + Sync {
  /// Unique identifier for each stored record.
  ///
  /// This is an opaque handle to a record. Simple implementations use `usize`
  /// (index into Vec), but more sophisticated implementations could use
  /// database row IDs, memory-mapped file offsets, etc.
  type Index: Serialize + Clone + Hash + Eq + Debug + Send + Sync;

  /// Borrowed reference to a record.
  ///
  /// Must implement [`LaburnumRecordRef`] to enable the framework to downcast
  /// to specific record types (diagnostics, symbols, etc.) without knowing your
  /// concrete enum type.
  type RecordRef<'a>: LaburnumRecordRef
  where
    Self: 'a;

  /// Builder type for constructing this storage.
  type Builder: PartitionsBuilder<Storage = Self>;

  /// Retrieves a record by its index.
  ///
  /// Returns `None` if the index is invalid (though this should be rare since
  /// indexes are returned by `Builder::push()`).
  fn get(&self, idx: &Self::Index) -> Option<Self::RecordRef<'_>>;

  /// Hashes all storage contents for chunk content addressing.
  ///
  /// This hash becomes part of the chunk ID. Hash semantic content (AST
  /// structure, types, symbols). Typically exclude source positions -
  /// prevents cache invalidation when unrelated code moves.
  ///
  /// If semantic content is identical, hash should be identical for cache hits.
  ///
  ///  ```rust-norun
  ///   fn hash_contents(&self, hasher: &mut ContentHasher) {
  ///     for record in &self.records {
  ///       hasher.update(&record);
  ///     }
  ///   }
  /// ```
  fn hash_contents(&self, hasher: &mut ContentHasher);
}

/// Builder for constructing [`RecordStorage`] instances.
///
/// Records are added via `push()` and the final storage is created with
/// `build()`. The builder is consumed when building the storage.
///
/// # Implementation Requirements
///
/// - Must implement `Default` for initialization
/// - Must support [`LaburnumRecord`](crate::record::LaburnumRecord) via `From`
///   trait on `Record` type
///
/// # Example
///
/// ```ignore
/// #[derive(Default)]
/// pub struct MyStorageBuilder {
///   records: Vec<MyRecord>,
/// }
///
/// impl RecordStorageBuilder for MyStorageBuilder {
///   type Storage = MyStorage;
///   type Record = MyRecord;
///
///   fn push(&mut self, record: Self::Record) -> usize {
///     let idx = self.records.len();
///     self.records.push(record);
///     idx
///   }
///
///   fn build(self) -> Self::Storage {
///     MyStorage { records: self.records }
///   }
/// }
/// ```
pub trait PartitionsBuilder: Default + Send {
  /// The storage type this builder constructs.
  type Storage: RecordStorage;

  /// The record type that can be added to this storage.
  ///
  /// # Framework Integration
  ///
  /// Must implement `From<LaburnumRecord>` to allow the framework to insert
  /// its own record types (for watchers, hooks, etc.). Typically this is a
  /// variant in your record enum:
  ///
  /// ```ignore
  /// pub enum MyRecord {
  ///   Laburnum(LaburnumRecord),  // Framework records
  ///   Function(Function),         // Your records
  ///   // ...
  /// }
  ///
  /// impl From<LaburnumRecord> for MyRecord {
  ///   fn from(lr: LaburnumRecord) -> Self {
  ///     MyRecord::Laburnum(lr)
  ///   }
  /// }
  /// ```
  type Record;

  /// Adds a record to the storage and returns its index.
  ///
  /// The returned index can be used later to retrieve this record via
  /// `RecordStorage::get()`. Indexes must be stable - the same index
  /// should always refer to the same record within a chunk.
  fn push(
    &mut self,
    record: Self::Record,
  ) -> <Self::Storage as RecordStorage>::Index;

  /// Consumes the builder and produces the final immutable storage.
  ///
  /// After calling `build()`, no more records can be added. The resulting
  /// storage is immutable and will be stored in a chunk.
  fn build(self) -> Self::Storage;
}