laburnum 1.17.3

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

use crate::{
  ContentHash,
  connect::lsp::ClientKind,
  database::{
    HasPartition,
    Partition,
    PartitionKey,
    Partitions,
    partitions::HandleEntry,
  },
  record::{
    Record,
    hash_record,
  },
};

pub struct Clients;

impl PartitionKey for Clients {
  const KEY: crate::SpannedIdent = crate::partitions::known_keys::CLIENTS;
}
impl Partition for Clients {
  type Record = ClientRecord;
  type IndexEntry = HandleEntry<Self>;
  type SortKey = ClientSortKey;

  fn index_entry_from_handle(
    handle: crate::database::RecordHandle<Self>,
  ) -> Self::IndexEntry {
    HandleEntry::new(handle)
  }
}

#[derive(Debug, Clone, Hash)]
pub enum ClientRecord {
  Connection {
    client_id: u64,
    kind:      ClientKind,
  },
  Subscription {
    client_id: u64,
    topic:     String,
  },
}

impl Record for ClientRecord {
  fn content_hash(&self) -> ContentHash {
    hash_record(self)
  }
}

impl bluegum::Bluegum for ClientRecord {
  fn node(&self, b: &mut bluegum::Builder) {
    match self {
      ClientRecord::Connection { client_id, kind } => {
        b.name("ClientRecord::Connection")
          .field("client_id", client_id)
          .field("kind", format!("{kind:?}"));
      }
      ClientRecord::Subscription { client_id, topic } => {
        b.name("ClientRecord::Subscription")
          .field("client_id", client_id)
          .field("topic", topic);
      }
    }
  }
}

impl bluegum::BluegumWithState<dyn crate::SourceResolver> for ClientRecord {}

#[derive(Debug, Clone)]
pub enum ClientSortKey {
  Connection { client_id: u64 },
  Subscription { client_id: u64, topic: String },
  All,
  ConnectionPrefix,
  SubscriptionPrefix,
  ClientPrefix { client_id: u64 },
}

impl std::fmt::Display for ClientSortKey {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    match self {
      | Self::Connection { client_id } => write!(f, "conn|{:020}", client_id),
      | Self::Subscription { client_id, topic } => {
        write!(f, "sub|{:020}|{}", client_id, topic)
      },
      | Self::All => Ok(()),
      | Self::ConnectionPrefix => write!(f, "conn|"),
      | Self::SubscriptionPrefix => write!(f, "sub|"),
      | Self::ClientPrefix { client_id } => write!(f, "sub|{:020}|", client_id),
    }
  }
}

// -- Sort key ordering (ADR0011) ----------------------------------------------

/// `Ns(0)` = connections, `Ns(1)` = subscriptions, mirroring the `conn` < `sub`
/// string ordering. `ClientPrefix` lives in the subscription namespace, so it
/// is a prefix of a client's `Subscription` keys.
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash)]
enum ClientToken {
  Ns(u8),
  ClientId(u64),
  Topic(String),
}

impl ClientSortKey {
  fn tokens(&self) -> Vec<ClientToken> {
    match self {
      | Self::All => Vec::new(),
      | Self::ConnectionPrefix => vec![ClientToken::Ns(0)],
      | Self::Connection { client_id } => {
        vec![ClientToken::Ns(0), ClientToken::ClientId(*client_id)]
      },
      | Self::SubscriptionPrefix => vec![ClientToken::Ns(1)],
      | Self::ClientPrefix { client_id } => {
        vec![ClientToken::Ns(1), ClientToken::ClientId(*client_id)]
      },
      | Self::Subscription { client_id, topic } => vec![
        ClientToken::Ns(1),
        ClientToken::ClientId(*client_id),
        ClientToken::Topic(topic.clone()),
      ],
    }
  }
}

impl PartialEq for ClientSortKey {
  fn eq(&self, other: &Self) -> bool {
    self.tokens() == other.tokens()
  }
}
impl Eq for ClientSortKey {}
impl std::hash::Hash for ClientSortKey {
  fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
    self.tokens().hash(state);
  }
}
impl PartialOrd for ClientSortKey {
  fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
    Some(self.cmp(other))
  }
}
impl Ord for ClientSortKey {
  fn cmp(&self, other: &Self) -> std::cmp::Ordering {
    self.tokens().cmp(&other.tokens())
  }
}

impl bluegum::Bluegum for ClientSortKey {
  fn node(&self, b: &mut bluegum::Builder) {
    b.name("ClientSortKey");
  }
}

impl bluegum::BluegumWithState<dyn crate::SourceResolver> for ClientSortKey {
  fn node_with_state(
    &self,
    b: &mut bluegum::Builder,
    resolver: &(dyn crate::SourceResolver + 'static),
  ) {
    b.name("ClientSortKey")
      .alt(crate::database::PartitionSortKey::resolve(self, resolver));
  }
}

impl crate::database::PartitionSortKey for ClientSortKey {
  fn is_prefix_of(&self, other: &Self) -> bool {
    other.tokens().starts_with(&self.tokens())
  }

  fn resolve(&self, _resolver: &dyn crate::SourceResolver) -> String {
    self.to_string()
  }
}

use crate::database::chunk::RecordWriter;

pub trait ClientWriteExt<P: Partitions> {
  fn write_client_connection(&mut self, client_id: u64, kind: ClientKind)
  where
    P::Stores: HasPartition<Clients>,
    Clients: crate::database::partitions::SortKeyOf<P>;

  fn write_client_subscription(&mut self, client_id: u64, topic: String)
  where
    P::Stores: HasPartition<Clients>,
    Clients: crate::database::partitions::SortKeyOf<P>;
}

impl<P: Partitions> ClientWriteExt<P> for RecordWriter<P> {
  fn write_client_connection(&mut self, client_id: u64, kind: ClientKind)
  where
    P::Stores: HasPartition<Clients>,
    Clients: crate::database::partitions::SortKeyOf<P>,
  {
    let sort_key = ClientSortKey::Connection { client_id };
    let record = ClientRecord::Connection { client_id, kind };

    self.insert::<Clients>(sort_key, record);
  }

  fn write_client_subscription(&mut self, client_id: u64, topic: String)
  where
    P::Stores: HasPartition<Clients>,
    Clients: crate::database::partitions::SortKeyOf<P>,
  {
    let sort_key = ClientSortKey::Subscription {
      client_id,
      topic: topic.clone(),
    };
    let record = ClientRecord::Subscription { client_id, topic };

    self.insert::<Clients>(sort_key, record);
  }
}