laburnum 1.17.1

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

use {
  super::Partition,
  crate::{
    ContentHash, Uri,
    database::{
      GenerationEpoch,
      handle::RecordHandle,
      hlist::{HCons, HNil, Here, There},
    },
    hash::{
      ContentHashMap, ContentHashMapExt,
    },
  },
  parking_lot::{RwLock, RwLockReadGuard},
  std::{collections::{BTreeMap, HashMap}, marker::PhantomData},
};

use {
  crate::record::CollectReferences,
  std::{
    borrow::Borrow,
    sync::atomic::{AtomicUsize, Ordering},
  },
};

pub struct PartitionStoreInner<P: Partition> {
  pub(crate) records: ContentHashMap<(P::Record, GenerationEpoch)>,
  /// Index entries keyed by sort key.
  ///
  /// Stores `P::IndexEntry` which may be a simple `HandleEntry<P>` (just a
  /// content hash reference) or a richer type with position-specific metadata.
  pub(crate) index: BTreeMap<String, P::IndexEntry>,
  /// Reference counts for each record.
  ///
  /// Tracks how many persistent references point to each record:
  /// - Index entries pointing to the record
  /// - Other records containing a RecordHandle to this record
  ///
  /// When refcount drops to zero and the reaper runs, the record can be removed.
  pub(crate) refcounts: ContentHashMap<AtomicUsize>,
  /// Span indexes keyed by URI.
  ///
  /// Each URI maps to a SpanIndex that enables O(log n) lookup of the
  /// record containing a given byte offset. Built during commit from
  /// staged intervals.
  pub(crate) span_indexes: HashMap<Uri, super::span_index::SpanIndex>,
}

impl<P: Partition> PartitionStoreInner<P> {
  pub fn new() -> Self {
    Self {
      records: ContentHashMap::new(),
      index: BTreeMap::new(),
      refcounts: ContentHashMap::new(),
      span_indexes: HashMap::new(),
    }
  }
}

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

/// A guard type that provides borrowed access to a content-addressed record
/// without cloning it.
///
/// `RecordRef` wraps an [`RwLockReadGuard`] alongside the record's
/// [`ContentHash`], so the read lock is held exactly as long as the caller
/// needs the data. This avoids the cost of cloning potentially large records
/// on every read while keeping the lock scope tied to the reference's
/// lifetime.
///
/// Obtain a `RecordRef` via [`PartitionStore::get`] or
/// [`PartitionStore::get_by_handle`]. Use [`record()`](Self::record) to
/// borrow the underlying record, or [`clone_record()`](Self::clone_record)
/// when you need ownership.
pub struct RecordRef<'a, P: Partition> {
  guard: RwLockReadGuard<'a, PartitionStoreInner<P>>,
  hash: ContentHash,
}

impl<'a, P: Partition> RecordRef<'a, P> {
  pub fn record(&self) -> Option<&P::Record> {
    self.guard.records.get(&self.hash).map(|(r, _)| r)
  }

  pub fn epoch(&self) -> Option<GenerationEpoch> {
    self.guard.records.get(&self.hash).map(|(_, e)| *e)
  }

  pub fn value(&self) -> Option<&(P::Record, GenerationEpoch)> {
    self.guard.records.get(&self.hash)
  }

  pub fn clone_record(&self) -> Option<P::Record>
  where
    P::Record: Clone,
  {
    self.record().cloned()
  }
}

pub struct PartitionStore<P: Partition> {
  inner: RwLock<PartitionStoreInner<P>>,
  _marker: PhantomData<P>,
}

impl<P: Partition> PartitionStore<P> {
  pub fn new() -> Self {
    Self {
      inner: RwLock::new(PartitionStoreInner::new()),
      _marker: PhantomData,
    }
  }

  pub fn get(&self, hash: &ContentHash) -> Option<RecordRef<'_, P>> {
    let guard = self.inner.read();
    if guard.records.contains_key(hash) {
      Some(RecordRef { guard, hash: *hash })
    } else {
      None
    }
  }

  pub fn get_by_handle(
    &self,
    handle: &RecordHandle<P>,
  ) -> Option<RecordRef<'_, P>> {
    self.get(&handle.content_hash())
  }

  pub fn insert(
    &self,
    hash: ContentHash,
    record: P::Record,
    epoch: GenerationEpoch,
  ) -> (RecordHandle<P>, bool) {
    use std::collections::hash_map::Entry;
    let mut inner = self.inner.write();
    let is_new = match inner.records.entry(hash) {
      | Entry::Occupied(_) => false,
      | Entry::Vacant(entry) => {
        entry.insert((record, epoch));
        true
      },
    };
    (RecordHandle::new(hash), is_new)
  }

  pub fn len(&self) -> usize {
    self.inner.read().records.len()
  }

  pub fn is_empty(&self) -> bool {
    self.inner.read().records.is_empty()
  }

  pub fn retain<F>(&self, mut predicate: F)
  where
    F: FnMut(ContentHash, GenerationEpoch) -> bool,
  {
    let mut inner = self.inner.write();
    inner.records.retain(|k, v| predicate(*k, v.1));
    inner.records.shrink_to_fit();
  }

  pub fn iter_hashes(&self) -> Vec<(ContentHash, GenerationEpoch)> {
    let inner = self.inner.read();
    inner.records.iter().map(|(k, v)| (*k, v.1)).collect()
  }

  /// Get an index entry by sort key.
  pub fn index_get(&self, sort_key: &str) -> Option<P::IndexEntry> {
    let inner = self.inner.read();
    inner.index.get(sort_key).cloned()
  }

  /// Get all index entries matching a prefix.
  pub fn index_range(&self, prefix: &str) -> Vec<(String, P::IndexEntry)> {
    let inner = self.inner.read();
    inner
      .index
      .range(prefix.to_string()..)
      .take_while(|(sk, _)| sk.starts_with(prefix))
      .map(|(sk, entry)| (sk.clone(), entry.clone()))
      .collect()
  }

  /// Get all index entries with sort key less than (or equal to) a value.
  pub fn index_less_than(&self, value: &str, inclusive: bool) -> Vec<(String, P::IndexEntry)> {
    let inner = self.inner.read();
    use std::ops::Bound;
    let end = if inclusive {
      Bound::Included(value.to_string())
    } else {
      Bound::Excluded(value.to_string())
    };
    inner
      .index
      .range::<String, _>((Bound::Unbounded, end))
      .map(|(sk, entry)| (sk.clone(), entry.clone()))
      .collect()
  }

  /// Get all index entries with sort key greater than (or equal to) a value.
  pub fn index_greater_than(&self, value: &str, inclusive: bool) -> Vec<(String, P::IndexEntry)> {
    let inner = self.inner.read();
    use std::ops::Bound;
    let start = if inclusive {
      Bound::Included(value.to_string())
    } else {
      Bound::Excluded(value.to_string())
    };
    inner
      .index
      .range::<String, _>((start, Bound::Unbounded))
      .map(|(sk, entry)| (sk.clone(), entry.clone()))
      .collect()
  }

  /// Get all index entries with sort key between two values (inclusive).
  pub fn index_between(&self, from: &str, to: &str) -> Vec<(String, P::IndexEntry)> {
    let inner = self.inner.read();
    inner
      .index
      .range(from.to_string()..=to.to_string())
      .map(|(sk, entry)| (sk.clone(), entry.clone()))
      .collect()
  }

  /// Insert an index entry.
  ///
  /// Returns the old entry if one existed at this sort key.
  pub fn index_insert(&self, sort_key: String, entry: P::IndexEntry) -> Option<P::IndexEntry> {
    let mut inner = self.inner.write();
    inner.index.insert(sort_key, entry)
  }

  /// Remove all index entries matching a prefix.
  ///
  /// Returns the removed entries for GC refcount processing.
  pub fn index_remove_prefix(&self, prefix: &str) -> Vec<P::IndexEntry> {
    let mut inner = self.inner.write();
    let keys_to_remove: Vec<_> = inner
      .index
      .range(prefix.to_string()..)
      .take_while(|(sk, _)| sk.starts_with(prefix))
      .map(|(sk, _)| sk.clone())
      .collect();

    let mut removed = Vec::with_capacity(keys_to_remove.len());
    for key in keys_to_remove {
      if let Some(entry) = inner.index.remove(&key) {
        removed.push(entry);
      }
    }
    removed
  }

  /// Get all index entry values.
  pub fn index_values(&self) -> Vec<P::IndexEntry> {
    let inner = self.inner.read();
    inner.index.values().cloned().collect()
  }

  /// Get all content hashes from index entries.
  ///
  /// Uses `IndexEntry::primary_hash()` to extract the primary referenced
  /// hash from each index entry. Index entries that don't reference any
  /// CAS records (like `()`) return None and are filtered out.
  pub fn index_hashes(&self) -> Vec<ContentHash> {
    use super::IndexEntry;
    let inner = self.inner.read();
    inner
      .index
      .values()
      .filter_map(|entry| entry.primary_hash())
      .collect()
  }

  pub fn index_len(&self) -> usize {
    self.inner.read().index.len()
  }

  pub fn index_is_empty(&self) -> bool {
    self.inner.read().index.is_empty()
  }

  /// Get all index entries as (sort_key, entry) pairs.
  pub fn index_entries(&self) -> Vec<(String, P::IndexEntry)> {
    let inner = self.inner.read();
    inner
      .index
      .iter()
      .map(|(sk, entry)| (sk.clone(), entry.clone()))
      .collect()
  }

  /// Drain records and index entries from iterators into this store.
  pub fn drain_from<I>(
    &self,
    records: I,
    index_entries: Vec<(String, P::IndexEntry)>,
    epoch: GenerationEpoch,
  ) where
    I: IntoIterator<Item = (ContentHash, P::Record)>,
  {
    let mut inner = self.inner.write();
    for (hash, record) in records {
      inner.records.entry(hash).or_insert((record, epoch));
    }
    for (sort_key, entry) in index_entries {
      inner.index.insert(sort_key, entry);
    }
  }

  /// Increment the reference count for a record.
  ///
  /// Called when:
  /// - A new index entry points to the record
  /// - A new record is inserted that references this record
  pub fn increment_refcount(&self, hash: ContentHash) {
    let mut inner = self.inner.write();
    inner
      .refcounts
      .entry(hash)
      .or_insert_with(|| AtomicUsize::new(0))
      .fetch_add(1, Ordering::SeqCst);
  }

  /// Decrement the reference count for a record.
  ///
  /// Called when:
  /// - An index entry is overwritten or removed
  /// - A record that referenced this record is removed
  ///
  /// Returns the new reference count. If it returns 0, the record can be
  /// removed (but this should be done by the reaper for epoch safety).
  pub fn decrement_refcount(&self, hash: ContentHash) -> usize {
    let mut inner = self.inner.write();
    if let Some(count) = inner.refcounts.get(&hash) {
      let prev = count.fetch_sub(1, Ordering::SeqCst);
      if prev == 1 {
        // Refcount is now 0, remove the entry
        inner.refcounts.remove(&hash);
        return 0;
      }
      prev - 1
    } else {
      0
    }
  }

  /// Get the current reference count for a record.
  pub fn get_refcount(&self, hash: ContentHash) -> usize {
    let inner = self.inner.read();
    inner
      .refcounts
      .get(&hash)
      .map(|c| c.load(Ordering::SeqCst))
      .unwrap_or(0)
  }

  /// Remove a record from the store.
  ///
  /// Called by the reaper when a record's refcount drops to zero and
  /// all tasks from the decrement epoch have completed.
  ///
  /// Returns the removed record if it existed.
  pub fn remove_record(&self, hash: ContentHash) -> Option<P::Record>
  where
    P::Record: Clone,
  {
    let mut inner = self.inner.write();
    inner.records.remove(&hash).map(|(r, _)| r)
  }

  /// Query the span index for a URI at a given byte offset.
  ///
  /// Returns the content hash of the innermost span containing the offset.
  pub fn span_index_query(&self, uri: &Uri, byte_offset: u64) -> Option<ContentHash> {
    let inner = self.inner.read();
    inner.span_indexes.get(uri)?.query(byte_offset)
  }

  /// Replace the span index for a URI.
  pub fn span_index_replace(&self, uri: Uri, index: super::span_index::SpanIndex) {
    let mut inner = self.inner.write();
    inner.span_indexes.insert(uri, index);
  }

  /// Remove the span index for a URI.
  pub fn span_index_remove(&self, uri: &Uri) {
    let mut inner = self.inner.write();
    inner.span_indexes.remove(uri);
  }
}

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

unsafe impl<P: Partition> Send for PartitionStore<P> where P::Record: Send {}
unsafe impl<P: Partition> Sync for PartitionStore<P> where P::Record: Sync {}

/// Trait for accessing a `PartitionStore<P>` at a specific type-level index.
///
/// This trait enables compile-time indexing into HLists. The `Idx` type parameter
/// encodes the position using [`Here`] and [`There`].
///
/// # Type Parameters
///
/// - `P`: The partition type to retrieve
/// - `Idx`: The type-level index (e.g., `Here`, `There<Here>`)
///
/// # Implementation
///
/// - `HCons<PartitionStore<P>, Tail>` implements `HasPartitionAt<P, Here>` (head access)
/// - `HCons<Head, Tail>` implements `HasPartitionAt<P, There<Idx>>` when `Tail: HasPartitionAt<P, Idx>` (recursive tail access)
pub trait HasPartitionAt<P: Partition, Idx> {
  /// Returns a reference to the `PartitionStore<P>` at the indexed position.
  fn store_at(&self) -> &PartitionStore<P>;
}

impl<P: Partition, Tail> HasPartitionAt<P, Here>
  for HCons<PartitionStore<P>, Tail>
{
  fn store_at(&self) -> &PartitionStore<P> {
    &self.head
  }
}

impl<P: Partition, Head, Tail, Idx> HasPartitionAt<P, There<Idx>>
  for HCons<Head, Tail>
where
  Tail: HasPartitionAt<P, Idx>,
{
  fn store_at(&self) -> &PartitionStore<P> {
    self.tail.store_at()
  }
}

/// Ergonomic trait for accessing a `PartitionStore<P>` without specifying the index.
///
/// Unlike [`HasPartitionAt`], this trait doesn't require the caller to know the
/// position of `P` in the HList. The [`impl_has_partition!`](crate::impl_has_partition)
/// macro generates implementations that delegate to the correct `HasPartitionAt` impl.
///
/// # Type Parameters
///
/// - `P`: The partition type to retrieve
///
/// # Example
///
/// ```ignore
/// fn get_tokens<S: HasPartition<TokenPartition>>(stores: &S) -> &PartitionStore<TokenPartition> {
///     stores.store()
/// }
/// ```
pub trait HasPartition<P: Partition> {
  /// Returns a reference to the `PartitionStore<P>`.
  fn store(&self) -> &PartitionStore<P>;
}

/// Trait for accessing index storage for a partition in an HList.
///
/// This trait enables routing to the correct index storage position in an HList
/// of index stores. It mirrors `HasPartition` but for index entries rather than
/// CAS records.
///
/// Index stores are `BTreeMap<String, P::IndexEntry>` for indexed partitions,
/// or `PhantomData` for index-free partitions (where `IndexEntry = ()`).
pub trait HasIndexPartition<Part: super::Partition> {
  /// Access the index entries for this partition.
  fn index_map(&self) -> &std::collections::BTreeMap<String, Part::IndexEntry>;

  /// Access the index entries mutably for this partition.
  fn index_map_mut(&mut self) -> &mut std::collections::BTreeMap<String, Part::IndexEntry>;
}

/// Trait for constructing new HList instances with default `PartitionStore`s.
///
/// This trait provides a more idiomatic recursive trait implementation for constructing
/// HLists of `PartitionStore`s. Each HList type implements this to construct itself
/// with fresh `PartitionStore` instances.
pub trait NewStores {
  /// Creates a new HList with default-initialized `PartitionStore`s.
  fn new_stores() -> Self;
}

impl NewStores for HNil {
  fn new_stores() -> Self {
    HNil
  }
}

impl<P: Partition, T: NewStores> NewStores for HCons<PartitionStore<P>, T> {
  fn new_stores() -> Self {
    HCons::new(PartitionStore::new(), T::new_stores())
  }
}

/// Result of merging records from one store into another.
///
/// Contains the set of content hashes that were newly inserted (did not exist
/// in the target store before the merge). Each entry is a partition-tagged
/// [`ContentHashRef`](crate::database::ContentHashRef) so the GC write barrier
/// can route new records to the correct partition.
pub struct MergeResult {
  pub new_hashes: Vec<crate::database::ContentHashRef>,
}

impl MergeResult {
  pub fn new() -> Self {
    Self {
      new_hashes: Vec::new(),
    }
  }

  pub fn extend(&mut self, other: MergeResult) {
    self.new_hashes.extend(other.new_hashes);
  }
}

impl Default for MergeResult {
  fn default() -> Self {
    Self::new()
  }
}

/// Trait for merging records from one HList of stores into another.
///
/// Used during chunk commit to copy records from a task-local store into
/// the global database. Each partition's records are merged independently.
///
/// Returns a [`MergeResult`] containing the set of content hashes that were
/// newly inserted into the target store.
pub trait MergeFrom {
  /// Merges all records from `source` into `self` with the given epoch.
  ///
  /// Returns a [`MergeResult`] containing the content hashes that were newly
  /// inserted (did not already exist in `self`).
  fn merge_from(&self, source: &Self, epoch: GenerationEpoch) -> MergeResult;
}

impl MergeFrom for HNil {
  fn merge_from(&self, _source: &Self, _epoch: GenerationEpoch) -> MergeResult {
    MergeResult::new()
  }
}

impl<P: Partition, T: MergeFrom> MergeFrom for HCons<PartitionStore<P>, T>
where
  P::Record: Clone,
{
  fn merge_from(&self, source: &Self, epoch: GenerationEpoch) -> MergeResult {
    let mut result = MergeResult::new();
    // Merge CAS records
    for (hash, _) in source.head.iter_hashes() {
      if let Some(record_ref) = source.head.get(&hash)
        && let Some(record) = record_ref.record()
      {
        let (_handle, is_new) = self.head.insert(hash, record.clone(), epoch);
        if is_new {
          result.new_hashes.push(crate::database::ContentHashRef::new(P::KEY, hash));
        }
      }
    }
    // Merge index entries
    for (sort_key, entry) in source.head.index_entries() {
      self.head.index_insert(sort_key, entry);
    }
    result.extend(self.tail.merge_from(&source.tail, epoch));
    result
  }
}

/// Trait for collecting all content hashes from partition indexes.
///
/// Used by garbage collection to find GC roots — all hashes that are
/// reachable via the index and should not be collected. Each hash is
/// tagged with its partition key so the GC can route items to the
/// correct partition store during marking.
pub trait CollectIndexHashes {
  /// Extends `result` with all partition-tagged content hashes from all partition indexes.
  fn collect_index_hashes(&self, result: &mut Vec<crate::database::ContentHashRef>);
}

impl CollectIndexHashes for HNil {
  fn collect_index_hashes(&self, _result: &mut Vec<crate::database::ContentHashRef>) {}
}

impl<P: Partition, T: CollectIndexHashes> CollectIndexHashes
  for HCons<PartitionStore<P>, T>
{
  fn collect_index_hashes(&self, result: &mut Vec<crate::database::ContentHashRef>) {
    for hash in self.head.index_hashes() {
      result.push(crate::database::ContentHashRef::new(P::KEY, hash));
    }
    self.tail.collect_index_hashes(result);
  }
}

/// Trait for routing refcount operations through an HList by partition key.
///
/// This enables dynamic dispatch of refcount operations when we only know
/// the partition key at runtime (e.g., during reaping from a queue).
pub trait RefcountOps {
  /// Increment the refcount for a record in the specified partition.
  fn increment_refcount(&self, partition: crate::Ident, hash: ContentHash);

  /// Decrement the refcount for a record in the specified partition.
  ///
  /// Returns the new refcount (0 means the record can be removed).
  fn decrement_refcount(
    &self,
    partition: crate::Ident,
    hash: ContentHash,
  ) -> usize;

  /// Remove a record from the specified partition.
  ///
  /// Returns true if the record was present and removed.
  fn remove_record(&self, partition: crate::Ident, hash: ContentHash) -> bool;

  /// Get the refcount for a record in the specified partition.
  fn get_refcount(&self, partition: crate::Ident, hash: ContentHash) -> usize;
}

impl RefcountOps for HNil {
  fn increment_refcount(&self, _partition: crate::Ident, _hash: ContentHash) {}

  fn decrement_refcount(
    &self,
    _partition: crate::Ident,
    _hash: ContentHash,
  ) -> usize {
    0
  }

  fn remove_record(
    &self,
    _partition: crate::Ident,
    _hash: ContentHash,
  ) -> bool {
    false
  }

  fn get_refcount(
    &self,
    _partition: crate::Ident,
    _hash: ContentHash,
  ) -> usize {
    0
  }
}

impl<P: Partition, T: RefcountOps> RefcountOps for HCons<PartitionStore<P>, T>
where
  P::Record: Clone,
{
  fn increment_refcount(&self, partition: crate::Ident, hash: ContentHash) {
    if partition == P::KEY {
      self.head.increment_refcount(hash);
    } else {
      self.tail.increment_refcount(partition, hash);
    }
  }

  fn decrement_refcount(
    &self,
    partition: crate::Ident,
    hash: ContentHash,
  ) -> usize {
    if partition == P::KEY {
      self.head.decrement_refcount(hash)
    } else {
      self.tail.decrement_refcount(partition, hash)
    }
  }

  fn remove_record(&self, partition: crate::Ident, hash: ContentHash) -> bool {
    if partition == P::KEY {
      self.head.remove_record(hash).is_some()
    } else {
      self.tail.remove_record(partition, hash)
    }
  }

  fn get_refcount(&self, partition: crate::Ident, hash: ContentHash) -> usize {
    if partition == P::KEY {
      self.head.get_refcount(hash)
    } else {
      self.tail.get_refcount(partition, hash)
    }
  }
}

/// Trait for collecting references from a record for cascading deletes.
///
/// This trait enables the reaper to collect all records that a given record
/// references, so their refcounts can be decremented when the parent record
/// is removed.
pub trait CollectCascadingRefs<P: crate::database::storage::Partitions + ?Sized> {
  /// Collect all (partition, hash) pairs that the record at the given
  /// location references.
  ///
  /// Returns an empty Vec if the partition key is unknown or the record
  /// doesn't exist.
  fn collect_cascading_refs(
    &self,
    partition: crate::Ident,
    hash: ContentHash,
  ) -> Vec<(crate::Ident, ContentHash)>;
}

impl<P: crate::database::storage::Partitions> CollectCascadingRefs<P> for HNil {
  fn collect_cascading_refs(
    &self,
    _partition: crate::Ident,
    _hash: ContentHash,
  ) -> Vec<(crate::Ident, ContentHash)> {
    Vec::new()
  }
}

impl<Part, T, P> CollectCascadingRefs<P> for HCons<PartitionStore<Part>, T>
where
  Part: Partition,
  Part::Record: Clone + crate::record::CollectReferences<P>,
  T: CollectCascadingRefs<P>,
  P: crate::database::storage::Partitions,
  P::Stores: HasPartition<Part>,
{
  fn collect_cascading_refs(
    &self,
    partition: crate::Ident,
    hash: ContentHash,
  ) -> Vec<(crate::Ident, ContentHash)> {
    if partition == Part::KEY {
      if let Some(record_ref) = self.head.get(&hash)
        && let Some(record) = record_ref.record()
      {
        let mut collector = CascadingRefCollector::<P>::new();
        record.collect_references(&mut collector);
        return collector.into_refs();
      }
      Vec::new()
    } else {
      self.tail.collect_cascading_refs(partition, hash)
    }
  }
}

/// Collector for cascading references.
///
/// Implements `References<P>` to collect all (partition, hash) pairs from
/// a record's references.
pub struct CascadingRefCollector<P: crate::database::storage::Partitions> {
  refs: Vec<(crate::Ident, ContentHash)>,
  _marker: std::marker::PhantomData<P>,
}

impl<P: crate::database::storage::Partitions> CascadingRefCollector<P> {
  pub fn new() -> Self {
    Self {
      refs: Vec::new(),
      _marker: std::marker::PhantomData,
    }
  }

  pub fn into_refs(self) -> Vec<(crate::Ident, ContentHash)> {
    self.refs
  }
}

impl<P: crate::database::storage::Partitions> Default
  for CascadingRefCollector<P>
{
  fn default() -> Self {
    Self::new()
  }
}

impl<P: crate::database::storage::Partitions> crate::record::References<P>
  for CascadingRefCollector<P>
{
  fn add<Part: Partition>(
    &mut self,
    handle: crate::database::handle::RecordHandle<Part>,
  ) where
    P::Stores: HasPartition<Part>,
  {
    self.refs.push((Part::KEY, handle.content_hash()));
  }
}

/// Trait for merging records with reference counting.
///
/// This extends `MergeFrom` to also increment refcounts for records that
/// newly inserted records reference.
pub trait MergeFromWithRefcount<P: crate::database::storage::Partitions>:
  MergeFrom
{
  /// Merges all records from `source` into `self` with the given epoch,
  /// incrementing refcounts for all records that newly inserted records reference.
  fn merge_from_with_refcount(
    &self,
    source: &Self,
    epoch: GenerationEpoch,
  ) -> MergeResult
  where
    P::Stores: RefcountOps;
}

impl<P: crate::database::storage::Partitions> MergeFromWithRefcount<P>
  for HNil
{
  fn merge_from_with_refcount(
    &self,
    _source: &Self,
    _epoch: GenerationEpoch,
  ) -> MergeResult
  where
    P::Stores: RefcountOps,
  {
    MergeResult::new()
  }
}

impl<Part, T, P> MergeFromWithRefcount<P> for HCons<PartitionStore<Part>, T>
where
  Part: Partition,
  Part::Record: Clone + crate::record::CollectReferences<P>,
  T: MergeFromWithRefcount<P>,
  P: crate::database::storage::Partitions,
  P::Stores: HasPartition<Part> + RefcountOps,
  // We need Self to be P::Stores for the incrementer
  Self: std::borrow::Borrow<P::Stores>,
{
  fn merge_from_with_refcount(
    &self,
    source: &Self,
    epoch: GenerationEpoch,
  ) -> MergeResult
  where
    P::Stores: RefcountOps,
  {
    let mut result = MergeResult::new();

    for (hash, _) in source.head.iter_hashes() {
      if let Some(record_ref) = source.head.get(&hash)
        && let Some(record) = record_ref.record()
      {
        let (_handle, is_new) = self.head.insert(hash, record.clone(), epoch);
        if is_new {
          result.new_hashes.push(crate::database::ContentHashRef::new(Part::KEY, hash));

          // Increment refcounts for all records this record references
          let stores: &P::Stores = self.borrow();
          let mut incrementer =
            crate::database::reaper::RefCountIncrementer::<P>::new(stores);
          record.collect_references(&mut incrementer);
        }
      }
    }

    result.extend(self.tail.merge_from_with_refcount(&source.tail, epoch));
    result
  }
}

/// Trait for rendering partition stores as bluegum tree nodes.
///
/// Walks the HList of `PartitionStore<P>` entries and renders each
/// partition's index entries and CAS records as bluegum nodes.
pub trait BluegumStores {
  fn bluegum_partition_nodes(
    &self,
    state: &(dyn crate::SpanResolver + 'static),
  ) -> Vec<bluegum::Builder>;
}

impl BluegumStores for HNil {
  fn bluegum_partition_nodes(
    &self,
    _state: &(dyn crate::SpanResolver + 'static),
  ) -> Vec<bluegum::Builder> {
    Vec::new()
  }
}

impl<P: Partition, T: BluegumStores> BluegumStores for HCons<PartitionStore<P>, T>
where
  P::Record: Clone,
{
  fn bluegum_partition_nodes(
    &self,
    state: &(dyn crate::SpanResolver + 'static),
  ) -> Vec<bluegum::Builder> {
    let mut nodes = Vec::new();

    let mut partition_builder = bluegum::Builder::new();
    match P::KEY_NAME {
      Some(name) => partition_builder.name(&format!("Partition({})", name)),
      None => partition_builder.name(&format!("Partition({})", P::KEY)),
    };

    let entries = self.head.index_entries();
    if !entries.is_empty() {
      let mut entry_builders = Vec::new();
      for (sort_key, entry) in &entries {
        let mut eb = bluegum::Builder::new();
        eb.name("IndexEntry")
          .field("sort_key", sort_key);
        <P::IndexEntry as bluegum::BluegumWithState<dyn crate::SpanResolver>>::node_with_state(entry, &mut eb, state);
        entry_builders.push(eb);
      }
      partition_builder.add_nodes_of_builders("index_entries", entry_builders);
    }

    let inner = self.head.inner.read();
    if !inner.records.is_empty() {
      let mut record_builders = Vec::new();
      for (hash, (record, _epoch)) in inner.records.iter() {
        let mut rb = bluegum::Builder::new();
        rb.field("hash", hash);
        <P::Record as bluegum::BluegumWithState<dyn crate::SpanResolver>>::node_with_state(record, &mut rb, state);
        record_builders.push(rb);
      }
      partition_builder.add_nodes_of_builders("records", record_builders);
    }

    nodes.push(partition_builder);
    nodes.extend(self.tail.bluegum_partition_nodes(state));
    nodes
  }
}