use core::{
cmp::Ordering,
hash::{Hash, Hasher},
};
use crate::{
bftag::BfTag,
buf::stack_heap_buf,
error::{Error, Result},
meta::kv_prefix,
};
pub const MEMBER_KEY_HEADER_SIZE: usize = 17;
pub const SCORE_KEY_HEADER_SIZE: usize = 25;
pub const ZSET_SUBKEY_STACK_CAP: usize = 128;
#[inline(always)]
const fn check_member_len(member_len: usize, header_size: usize) -> Result<usize> {
if member_len > u32::MAX as usize - header_size {
return Err(Error::KeyLengthOverflow(member_len));
}
match header_size.checked_add(member_len) {
Some(len) => Ok(len),
None => Err(Error::RecordSizeOverflow),
}
}
#[inline(always)]
pub const fn encode_order_preserving_f64(val: f64) -> [u8; 8] {
let bits = val.to_bits();
let mask = (((bits as i64) >> 63) as u64) | (1 << 63);
(bits ^ mask).to_be_bytes()
}
#[inline(always)]
pub const fn decode_order_preserving_f64(bytes: [u8; 8]) -> f64 {
let sortable = u64::from_be_bytes(bytes);
let mask = ((((!sortable) as i64) >> 63) as u64) | (1 << 63);
f64::from_bits(sortable ^ mask)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ZMemberKeyRef<'a> {
pub key_id: u64,
pub version: u64,
pub member: &'a [u8],
}
impl<'a> ZMemberKeyRef<'a> {
#[inline(always)]
pub const fn new(key_id: u64, version: u64, member: &'a [u8]) -> Self {
Self {
key_id,
version,
member,
}
}
#[inline(always)]
pub const fn from_slice(slice: &'a [u8]) -> Result<Self> {
ZSetSubKeyCodec::decode_member_key(slice)
}
#[inline(always)]
pub const fn header(&self) -> [u8; MEMBER_KEY_HEADER_SIZE] {
ZSetSubKeyCodec::encode_member_header(self.key_id, self.version)
}
#[inline(always)]
pub const fn encoded_len(&self) -> usize {
MEMBER_KEY_HEADER_SIZE + self.member.len()
}
#[inline]
pub fn write_to_slice(&self, dst: &mut [u8]) -> Result<usize> {
ZSetSubKeyCodec::encode_member_key_to_slice(self.key_id, self.version, self.member, dst)
}
#[inline]
pub fn to_vec(&self) -> Vec<u8> {
self.try_to_vec().unwrap_or_default()
}
#[inline]
pub fn try_to_vec(&self) -> Result<Vec<u8>> {
ZSetSubKeyCodec::encode_member_key(self.key_id, self.version, self.member)
}
}
#[derive(Debug, Clone, Copy)]
pub struct ZScoreKeyRef<'a> {
pub key_id: u64,
pub version: u64,
pub score: f64,
pub raw_score: [u8; 8],
pub member: &'a [u8],
}
impl<'a> ZScoreKeyRef<'a> {
#[inline(always)]
pub const fn new(key_id: u64, version: u64, score: f64, member: &'a [u8]) -> Self {
let raw_score = encode_order_preserving_f64(score);
Self {
key_id,
version,
score,
raw_score,
member,
}
}
#[inline(always)]
pub const fn from_raw(key_id: u64, version: u64, raw_score: [u8; 8], member: &'a [u8]) -> Self {
let score = decode_order_preserving_f64(raw_score);
Self {
key_id,
version,
score,
raw_score,
member,
}
}
#[inline(always)]
pub const fn from_slice(slice: &'a [u8]) -> Result<Self> {
ZSetSubKeyCodec::decode_score_key(slice)
}
#[inline(always)]
pub const fn header(&self) -> [u8; SCORE_KEY_HEADER_SIZE] {
ZSetSubKeyCodec::encode_score_header_from_raw(self.key_id, self.version, self.raw_score)
}
#[inline(always)]
pub const fn encoded_len(&self) -> usize {
SCORE_KEY_HEADER_SIZE + self.member.len()
}
#[inline]
pub fn write_to_slice(&self, dst: &mut [u8]) -> Result<usize> {
let total_len = check_member_len(self.member.len(), SCORE_KEY_HEADER_SIZE)?;
if dst.len() < total_len {
return Err(Error::BufferTooShort {
expected: total_len,
actual: dst.len(),
});
}
dst[..SCORE_KEY_HEADER_SIZE].copy_from_slice(&self.header());
dst[SCORE_KEY_HEADER_SIZE..total_len].copy_from_slice(self.member);
Ok(total_len)
}
#[inline]
pub fn to_vec(&self) -> Vec<u8> {
self.try_to_vec().unwrap_or_default()
}
#[inline]
pub fn try_to_vec(&self) -> Result<Vec<u8>> {
ZSetSubKeyCodec::encode_score_key(self.key_id, self.version, self.score, self.member)
}
}
impl<'a> PartialEq for ZScoreKeyRef<'a> {
#[inline(always)]
fn eq(&self, other: &Self) -> bool {
self.key_id == other.key_id
&& self.version == other.version
&& self.raw_score == other.raw_score
&& self.member == other.member
}
}
impl<'a> Eq for ZScoreKeyRef<'a> {}
impl<'a> PartialOrd for ZScoreKeyRef<'a> {
#[inline(always)]
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<'a> Ord for ZScoreKeyRef<'a> {
#[inline(always)]
fn cmp(&self, other: &Self) -> Ordering {
(self.key_id, self.version, &self.raw_score, self.member).cmp(&(
other.key_id,
other.version,
&other.raw_score,
other.member,
))
}
}
impl<'a> Hash for ZScoreKeyRef<'a> {
#[inline(always)]
fn hash<H: Hasher>(&self, state: &mut H) {
self.key_id.hash(state);
self.version.hash(state);
self.raw_score.hash(state);
self.member.hash(state);
}
}
stack_heap_buf!(ZSetSubKeyBuf, ZSET_SUBKEY_STACK_CAP);
impl ZSetSubKeyBuf {
#[inline]
pub fn from_member(key_id: u64, version: u64, member: &[u8]) -> Result<Self> {
ZSetSubKeyCodec::encode_member_key_buf(key_id, version, member)
}
#[inline]
pub fn from_score(key_id: u64, version: u64, score: f64, member: &[u8]) -> Result<Self> {
ZSetSubKeyCodec::encode_score_key_buf(key_id, version, score, member)
}
}
pub struct ZSetSubKeyCodec;
impl ZSetSubKeyCodec {
#[inline(always)]
pub const fn encode_member_header(key_id: u64, version: u64) -> [u8; MEMBER_KEY_HEADER_SIZE] {
kv_prefix(BfTag::ZMember.as_u8(), key_id, version)
}
#[inline(always)]
pub const fn encode_score_prefix(key_id: u64, version: u64) -> [u8; MEMBER_KEY_HEADER_SIZE] {
kv_prefix(BfTag::ZScore.as_u8(), key_id, version)
}
#[inline(always)]
pub const fn encode_score_header_from_raw(
key_id: u64,
version: u64,
raw_score: [u8; 8],
) -> [u8; SCORE_KEY_HEADER_SIZE] {
let prefix = Self::encode_score_prefix(key_id, version);
let mut out = [0u8; SCORE_KEY_HEADER_SIZE];
let mut i = 0;
while i < MEMBER_KEY_HEADER_SIZE {
out[i] = prefix[i];
i += 1;
}
let mut j = 0;
while j < 8 {
out[MEMBER_KEY_HEADER_SIZE + j] = raw_score[j];
j += 1;
}
out
}
#[inline(always)]
pub const fn encode_score_header(
key_id: u64,
version: u64,
score: f64,
) -> [u8; SCORE_KEY_HEADER_SIZE] {
Self::encode_score_header_from_raw(key_id, version, encode_order_preserving_f64(score))
}
#[inline]
pub const fn decode_member_header(slice: &[u8]) -> Result<(u64, u64)> {
if slice.len() < MEMBER_KEY_HEADER_SIZE {
return Err(Error::BufferTooShort {
expected: MEMBER_KEY_HEADER_SIZE,
actual: slice.len(),
});
}
if slice[0] != BfTag::ZMember.as_u8() {
return Err(Error::InvalidKeyTag(slice[0]));
}
let key_id = u64::from_be_bytes([
slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7], slice[8],
]);
let version = u64::from_be_bytes([
slice[9], slice[10], slice[11], slice[12], slice[13], slice[14], slice[15], slice[16],
]);
Ok((key_id, version))
}
#[inline]
pub const fn decode_score_header_raw(slice: &[u8]) -> Result<(u64, u64, [u8; 8])> {
if slice.len() < SCORE_KEY_HEADER_SIZE {
return Err(Error::BufferTooShort {
expected: SCORE_KEY_HEADER_SIZE,
actual: slice.len(),
});
}
if slice[0] != BfTag::ZScore.as_u8() {
return Err(Error::InvalidKeyTag(slice[0]));
}
let key_id = u64::from_be_bytes([
slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7], slice[8],
]);
let version = u64::from_be_bytes([
slice[9], slice[10], slice[11], slice[12], slice[13], slice[14], slice[15], slice[16],
]);
let raw_score = [
slice[17], slice[18], slice[19], slice[20], slice[21], slice[22], slice[23], slice[24],
];
Ok((key_id, version, raw_score))
}
#[inline]
pub const fn decode_score_header(slice: &[u8]) -> Result<(u64, u64, f64)> {
match Self::decode_score_header_raw(slice) {
Ok((key_id, version, raw_score)) => {
Ok((key_id, version, decode_order_preserving_f64(raw_score)))
}
Err(e) => Err(e),
}
}
#[inline]
pub fn encode_member_key_to_slice(
key_id: u64,
version: u64,
member: &[u8],
dst: &mut [u8],
) -> Result<usize> {
let total_len = check_member_len(member.len(), MEMBER_KEY_HEADER_SIZE)?;
if dst.len() < total_len {
return Err(Error::BufferTooShort {
expected: total_len,
actual: dst.len(),
});
}
let header = Self::encode_member_header(key_id, version);
dst[..MEMBER_KEY_HEADER_SIZE].copy_from_slice(&header);
dst[MEMBER_KEY_HEADER_SIZE..total_len].copy_from_slice(member);
Ok(total_len)
}
#[inline]
pub fn encode_member_key(key_id: u64, version: u64, member: &[u8]) -> Result<Vec<u8>> {
let total_len = check_member_len(member.len(), MEMBER_KEY_HEADER_SIZE)?;
let mut vec = Vec::with_capacity(total_len);
let header = Self::encode_member_header(key_id, version);
vec.extend_from_slice(&header);
vec.extend_from_slice(member);
Ok(vec)
}
#[inline]
pub fn encode_member_key_buf(key_id: u64, version: u64, member: &[u8]) -> Result<ZSetSubKeyBuf> {
let total_len = check_member_len(member.len(), MEMBER_KEY_HEADER_SIZE)?;
if total_len <= ZSET_SUBKEY_STACK_CAP {
let mut buf = [0u8; ZSET_SUBKEY_STACK_CAP];
let header = Self::encode_member_header(key_id, version);
buf[..MEMBER_KEY_HEADER_SIZE].copy_from_slice(&header);
buf[MEMBER_KEY_HEADER_SIZE..total_len].copy_from_slice(member);
Ok(ZSetSubKeyBuf::Stack(buf, total_len as u8))
} else {
let mut vec = Vec::with_capacity(total_len);
let header = Self::encode_member_header(key_id, version);
vec.extend_from_slice(&header);
vec.extend_from_slice(member);
Ok(ZSetSubKeyBuf::Heap(vec))
}
}
#[inline]
pub const fn decode_member_key<'a>(slice: &'a [u8]) -> Result<ZMemberKeyRef<'a>> {
let (key_id, version) = match Self::decode_member_header(slice) {
Ok(v) => v,
Err(e) => return Err(e),
};
let member = slice.split_at(MEMBER_KEY_HEADER_SIZE).1;
if member.len() > u32::MAX as usize - MEMBER_KEY_HEADER_SIZE {
return Err(Error::KeyLengthOverflow(member.len()));
}
Ok(ZMemberKeyRef {
key_id,
version,
member,
})
}
#[inline]
pub fn encode_score_key_to_slice(
key_id: u64,
version: u64,
score: f64,
member: &[u8],
dst: &mut [u8],
) -> Result<usize> {
let total_len = check_member_len(member.len(), SCORE_KEY_HEADER_SIZE)?;
if dst.len() < total_len {
return Err(Error::BufferTooShort {
expected: total_len,
actual: dst.len(),
});
}
let header = Self::encode_score_header(key_id, version, score);
dst[..SCORE_KEY_HEADER_SIZE].copy_from_slice(&header);
dst[SCORE_KEY_HEADER_SIZE..total_len].copy_from_slice(member);
Ok(total_len)
}
#[inline]
pub fn encode_score_key(key_id: u64, version: u64, score: f64, member: &[u8]) -> Result<Vec<u8>> {
let total_len = check_member_len(member.len(), SCORE_KEY_HEADER_SIZE)?;
let mut vec = Vec::with_capacity(total_len);
let header = Self::encode_score_header(key_id, version, score);
vec.extend_from_slice(&header);
vec.extend_from_slice(member);
Ok(vec)
}
#[inline]
pub fn encode_score_key_buf(
key_id: u64,
version: u64,
score: f64,
member: &[u8],
) -> Result<ZSetSubKeyBuf> {
let total_len = check_member_len(member.len(), SCORE_KEY_HEADER_SIZE)?;
if total_len <= ZSET_SUBKEY_STACK_CAP {
let mut buf = [0u8; ZSET_SUBKEY_STACK_CAP];
let header = Self::encode_score_header(key_id, version, score);
buf[..SCORE_KEY_HEADER_SIZE].copy_from_slice(&header);
buf[SCORE_KEY_HEADER_SIZE..total_len].copy_from_slice(member);
Ok(ZSetSubKeyBuf::Stack(buf, total_len as u8))
} else {
let mut vec = Vec::with_capacity(total_len);
let header = Self::encode_score_header(key_id, version, score);
vec.extend_from_slice(&header);
vec.extend_from_slice(member);
Ok(ZSetSubKeyBuf::Heap(vec))
}
}
#[inline]
pub const fn decode_score_key<'a>(slice: &'a [u8]) -> Result<ZScoreKeyRef<'a>> {
let (key_id, version, raw_score) = match Self::decode_score_header_raw(slice) {
Ok(v) => v,
Err(e) => return Err(e),
};
let score = decode_order_preserving_f64(raw_score);
let member = slice.split_at(SCORE_KEY_HEADER_SIZE).1;
if member.len() > u32::MAX as usize - SCORE_KEY_HEADER_SIZE {
return Err(Error::KeyLengthOverflow(member.len()));
}
Ok(ZScoreKeyRef {
key_id,
version,
score,
raw_score,
member,
})
}
}