const ROOT: &[u8] = b"kv_sorted:";
const SCORE: u8 = b's';
const PK: u8 = b'p';
pub(super) const FIELD_LEN: usize = 8;
pub(crate) fn entries_prefix(index_name: &str) -> Vec<u8> {
let len = u32::try_from(index_name.len()).unwrap_or(u32::MAX);
let mut key = Vec::with_capacity(ROOT.len() + 4 + index_name.len() + 1);
key.extend_from_slice(ROOT);
key.extend_from_slice(&len.to_be_bytes());
key.extend_from_slice(index_name.as_bytes());
key
}
pub(super) fn score_prefix(index_name: &str) -> Vec<u8> {
let mut key = entries_prefix(index_name);
key.push(SCORE);
key
}
pub(super) fn score_key(index_name: &str, score: &[u8; FIELD_LEN], pk: &[u8]) -> Vec<u8> {
let mut key = score_prefix(index_name);
key.extend_from_slice(score);
key.extend_from_slice(pk);
key
}
pub(super) fn pk_entry_key(index_name: &str, pk: &[u8]) -> Vec<u8> {
let mut key = entries_prefix(index_name);
key.push(PK);
key.extend_from_slice(pk);
key
}
pub(crate) fn f64_to_sort_bytes(score: f64) -> [u8; FIELD_LEN] {
let bits = score.to_bits();
let key_bits = if bits >> 63 == 0 {
bits ^ (1u64 << 63)
} else {
!bits
};
key_bits.to_be_bytes()
}
pub(super) fn sort_bytes_to_f64(bytes: &[u8; FIELD_LEN]) -> f64 {
let bits = u64::from_be_bytes(*bytes);
let original = if bits >> 63 != 0 {
bits ^ (1u64 << 63)
} else {
!bits
};
f64::from_bits(original)
}
pub(super) fn complement(bytes: [u8; FIELD_LEN]) -> [u8; FIELD_LEN] {
bytes.map(|b| !b)
}
pub(super) fn encode_score(score: f64, desc: bool) -> [u8; FIELD_LEN] {
let bytes = f64_to_sort_bytes(score);
if desc { complement(bytes) } else { bytes }
}
pub(super) fn decode_score(bytes: &[u8; FIELD_LEN], desc: bool) -> f64 {
let bytes = if desc { complement(*bytes) } else { *bytes };
sort_bytes_to_f64(&bytes)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct PkEntry {
pub score: [u8; FIELD_LEN],
pub ts: Option<[u8; FIELD_LEN]>,
pub deadline: u64,
}
impl PkEntry {
pub fn encode(&self) -> Vec<u8> {
let mut out = Vec::with_capacity(3 * FIELD_LEN);
out.extend_from_slice(&self.score);
out.extend_from_slice(&self.deadline.to_le_bytes());
if let Some(ts) = &self.ts {
out.extend_from_slice(ts);
}
out
}
pub fn decode(bytes: &[u8]) -> Option<Self> {
let score: [u8; FIELD_LEN] = bytes.get(..FIELD_LEN)?.try_into().ok()?;
let deadline_bytes: [u8; FIELD_LEN] =
bytes.get(FIELD_LEN..2 * FIELD_LEN)?.try_into().ok()?;
let ts = match bytes.len() {
n if n == 2 * FIELD_LEN => None,
n if n == 3 * FIELD_LEN => Some(bytes[2 * FIELD_LEN..].try_into().ok()?),
_ => return None,
};
Some(Self {
score,
ts,
deadline: u64::from_le_bytes(deadline_bytes),
})
}
pub fn same_position(&self, other: &Self) -> bool {
self.score == other.score && self.ts == other.ts
}
pub fn score_key(&self, index_name: &str, pk: &[u8]) -> Vec<u8> {
match &self.ts {
Some(ts) => super::window::windowed_score_key(index_name, &self.score, ts, pk),
None => score_key(index_name, &self.score, pk),
}
}
pub fn score_value(&self) -> Vec<u8> {
self.deadline.to_le_bytes().to_vec()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn scores_sort_numerically_and_round_trip() {
let values = [
f64::NEG_INFINITY,
-100.0,
-0.5,
0.0,
0.5,
100.0,
f64::INFINITY,
];
for desc in [false, true] {
let encoded: Vec<_> = values.iter().map(|v| encode_score(*v, desc)).collect();
for pair in encoded.windows(2) {
if desc {
assert!(pair[0] > pair[1]);
} else {
assert!(pair[0] < pair[1]);
}
}
for (v, e) in values.iter().zip(&encoded) {
assert_eq!(decode_score(e, desc).to_bits(), v.to_bits());
}
}
}
#[test]
fn index_prefixes_do_not_nest() {
let a = entries_prefix("a");
let b = entries_prefix("a:score:x");
assert!(!b.starts_with(&a));
assert!(!a.starts_with(&b));
}
#[test]
fn a_pk_entry_round_trips_with_and_without_timestamp() {
for ts in [None, Some([7u8; FIELD_LEN])] {
let entry = PkEntry {
score: encode_score(3.5, false),
ts,
deadline: 42,
};
assert_eq!(PkEntry::decode(&entry.encode()), Some(entry));
}
assert_eq!(PkEntry::decode(&[0u8; 9]), None);
}
}