reifydb_core/value/encoded/

key.rs

// Copyright (c) reifydb.com 2025
// This file is licensed under the AGPL-3.0-or-later, see license.md file

use std::{
	collections::{
		Bound,
		Bound::{Excluded, Included, Unbounded},
	},
	ops::{Deref, RangeBounds},
};

use serde::{Deserialize, Serialize};

use crate::util::{CowVec, encoding::binary::decode_binary};

#[derive(Debug, Clone, PartialOrd, Ord, Hash, PartialEq, Eq, Serialize, Deserialize)]
pub struct EncodedKey(pub CowVec<u8>);

impl Deref for EncodedKey {
	type Target = CowVec<u8>;

	fn deref(&self) -> &Self::Target {
		&self.0
	}
}

impl EncodedKey {
	pub fn new(key: impl Into<Vec<u8>>) -> Self {
		Self(CowVec::new(key.into()))
	}

	pub fn as_bytes(&self) -> &[u8] {
		self.0.as_ref()
	}

	pub fn as_slice(&self) -> &[u8] {
		self.0.as_ref()
	}
}

#[derive(Clone, Debug)]
pub struct EncodedKeyRange {
	pub start: Bound<EncodedKey>,
	pub end: Bound<EncodedKey>,
}

impl EncodedKeyRange {
	pub fn new(start: Bound<EncodedKey>, end: Bound<EncodedKey>) -> Self {
		Self {
			start,
			end,
		}
	}

	/// Generates a key range for a key prefix, used e.g. for prefix scans.
	///
	/// The exclusive end bound is generated by adding 1 to the value of the
	/// last byte. If the last byte(s) is 0xff (so adding 1 would
	/// saturation), we instead find the latest non-0xff byte, increment
	/// that, and truncate the rest. If all bytes are 0xff, we scan to the
	/// end of the range, since there can't be other prefixes after it.
	pub fn prefix(prefix: &[u8]) -> Self {
		let start = Bound::Included(EncodedKey::new(prefix));
		let end = match prefix.iter().rposition(|&b| b != 0xff) {
			Some(i) => Bound::Excluded(EncodedKey::new(
				prefix.iter()
					.take(i)
					.copied()
					.chain(std::iter::once(prefix[i] + 1))
					.collect::<Vec<_>>(),
			)),
			None => Bound::Unbounded,
		};
		Self {
			start,
			end,
		}
	}

	pub fn with_prefix(&self, prefix: EncodedKey) -> Self {
		let start = match self.start_bound() {
			Included(key) => {
				let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
				prefixed.extend_from_slice(prefix.as_ref());
				prefixed.extend_from_slice(key.as_ref());
				Included(EncodedKey::new(prefixed))
			}
			Excluded(key) => {
				let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
				prefixed.extend_from_slice(prefix.as_ref());
				prefixed.extend_from_slice(key.as_ref());
				Excluded(EncodedKey::new(prefixed))
			}
			Unbounded => Included(prefix.clone()),
		};

		let end = match self.end_bound() {
			Included(key) => {
				let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
				prefixed.extend_from_slice(prefix.as_ref());
				prefixed.extend_from_slice(key.as_ref());
				Included(EncodedKey::new(prefixed))
			}
			Excluded(key) => {
				let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
				prefixed.extend_from_slice(prefix.as_ref());
				prefixed.extend_from_slice(key.as_ref());
				Excluded(EncodedKey::new(prefixed))
			}
			Unbounded => match prefix.as_ref().iter().rposition(|&b| b != 0xff) {
				Some(i) => {
					let mut next_prefix = prefix.as_ref()[..=i].to_vec();
					next_prefix[i] += 1;
					Excluded(EncodedKey::new(next_prefix))
				}
				None => Unbounded,
			},
		};

		EncodedKeyRange::new(start, end)
	}

	/// Constructs a key range from an optional inclusive start key to an
	/// optional inclusive end key.
	///
	/// - `start`: If provided, marks the inclusive lower bound of the range. If `None`, the range is unbounded
	///   below.
	/// - `end`: If provided, marks the inclusive upper bound of the range. If `None`, the range is unbounded above.
	///
	/// This function does not modify the input keys and assumes they are
	/// already exact keys (not prefixes). If you need to scan all keys
	/// with a given prefix, use [`EncodedKeyRange::prefix`] instead.
	///
	/// Useful for scanning between two explicit keys in a sorted key-value
	/// store.
	pub fn start_end(start: Option<EncodedKey>, end: Option<EncodedKey>) -> Self {
		let start = match start {
			Some(s) => Bound::Included(s),
			None => Bound::Unbounded,
		};

		let end = match end {
			Some(e) => Bound::Included(e),
			None => Bound::Unbounded,
		};

		Self {
			start,
			end,
		}
	}

	/// Constructs a key range that fragments the entire keyspace.
	///
	/// This range has no lower or upper bounds, making it suitable for full
	/// scans over all keys in a sorted key-value store.
	///
	/// Equivalent to: `..` (in Rust range syntax)
	pub fn all() -> Self {
		Self {
			start: Bound::Unbounded,
			end: Bound::Unbounded,
		}
	}

	/// Parses a human-readable range string into a `KeyRange`.
	///
	/// The expected format is `<start>..[=]<end>`, where:
	/// - `<start>` is the inclusive lower bound (optional),
	/// - `..` separates the bounds,
	/// - `=` after `..` makes the upper bound inclusive,
	/// - `<end>` is the upper bound (optional).
	///
	/// Examples:
	/// - `"a..z"`       => start = Included("a"), end = Excluded("z")
	/// - `"a..=z"`      => start = Included("a"), end = Included("z")
	/// - `"..z"`        => start = Unbounded,     end = Excluded("z")
	/// - `"a.."`        => start = Included("a"), end = Unbounded
	///
	/// If parsing fails, it defaults to a degenerate range from `0xff` to
	/// `0xff` (empty).
	pub fn parse(str: &str) -> Self {
		let (mut start, mut end) = (Bound::<EncodedKey>::Unbounded, Bound::<EncodedKey>::Unbounded);

		// Find the ".." separator
		if let Some(dot_pos) = str.find("..") {
			let start_part = &str[..dot_pos];
			let end_part = &str[dot_pos + 2..];

			// Parse start bound
			if !start_part.is_empty() {
				start = Bound::Included(EncodedKey(decode_binary(start_part)));
			}

			// Parse end bound - check for inclusive marker "="
			if let Some(end_str) = end_part.strip_prefix('=') {
				// Inclusive end: "..="
				if !end_str.is_empty() {
					end = Bound::Included(EncodedKey(decode_binary(end_str)));
				}
			} else {
				// Exclusive end: ".."
				if !end_part.is_empty() {
					end = Bound::Excluded(EncodedKey(decode_binary(end_part)));
				}
			}

			Self {
				start,
				end,
			}
		} else {
			// Invalid format - return degenerate range
			Self {
				start: Bound::Included(EncodedKey::new([0xff])),
				end: Bound::Excluded(EncodedKey::new([0xff])),
			}
		}
	}
}

impl RangeBounds<EncodedKey> for EncodedKeyRange {
	fn start_bound(&self) -> Bound<&EncodedKey> {
		self.start.as_ref()
	}

	fn end_bound(&self) -> Bound<&EncodedKey> {
		self.end.as_ref()
	}
}

#[cfg(test)]
mod tests {
	use std::collections::Bound;

	use super::EncodedKey;

	macro_rules! as_key {
		($key:expr) => {{ EncodedKey::new(keycode::serialize(&$key)) }};
	}

	mod prefix {
		use std::ops::Bound;

		use crate::value::encoded::key::{
			EncodedKeyRange,
			tests::{excluded, included},
		};

		#[test]
		fn test_simple() {
			let range = EncodedKeyRange::prefix(&[0x12, 0x34]);
			assert_eq!(range.start, included(&[0x12, 0x34]));
			assert_eq!(range.end, excluded(&[0x12, 0x35]));
		}

		#[test]
		fn test_with_trailing_ff() {
			let range = EncodedKeyRange::prefix(&[0x12, 0xff]);
			assert_eq!(range.start, included(&[0x12, 0xff]));
			assert_eq!(range.end, excluded(&[0x13]));
		}

		#[test]
		fn test_with_multiple_trailing_ff() {
			let range = EncodedKeyRange::prefix(&[0x12, 0xff, 0xff]);
			assert_eq!(range.start, included(&[0x12, 0xff, 0xff]));
			assert_eq!(range.end, excluded(&[0x13]));
		}

		#[test]
		fn test_all_ff() {
			let range = EncodedKeyRange::prefix(&[0xff, 0xff]);
			assert_eq!(range.start, included(&[0xff, 0xff]));
			assert_eq!(range.end, Bound::Unbounded);
		}

		#[test]
		fn test_empty() {
			let range = EncodedKeyRange::prefix(&[]);
			assert_eq!(range.start, included(&[]));
			assert_eq!(range.end, Bound::Unbounded);
		}

		#[test]
		fn test_mid_increment() {
			let range = EncodedKeyRange::prefix(&[0x12, 0x00, 0xff]);
			assert_eq!(range.start, included(&[0x12, 0x00, 0xff]));
			assert_eq!(range.end, excluded(&[0x12, 0x01]));
		}
	}

	mod start_end {
		use std::ops::Bound;

		use crate::{
			EncodedKey,
			util::encoding::keycode,
			value::encoded::key::{EncodedKeyRange, tests::included},
		};

		#[test]
		fn test_start_and_end() {
			let range = EncodedKeyRange::start_end(Some(as_key!(1)), Some(as_key!(2)));
			assert_eq!(range.start, included(&as_key!(1)));
			assert_eq!(range.end, included(&as_key!(2)));
		}

		#[test]
		fn test_start_only() {
			let range = EncodedKeyRange::start_end(Some(as_key!(1)), None);
			assert_eq!(range.start, included(&as_key!(1)));
			assert_eq!(range.end, Bound::Unbounded);
		}

		#[test]
		fn test_end_only() {
			let range = EncodedKeyRange::start_end(None, Some(as_key!(2)));
			assert_eq!(range.start, Bound::Unbounded);
			assert_eq!(range.end, included(&as_key!(2)));
		}

		#[test]
		fn test_unbounded_range() {
			let range = EncodedKeyRange::start_end(None, None);
			assert_eq!(range.start, Bound::Unbounded);
			assert_eq!(range.end, Bound::Unbounded);
		}

		#[test]
		fn test_full_byte_range() {
			let range = EncodedKeyRange::start_end(Some(as_key!(0x00)), Some(as_key!(0xff)));
			assert_eq!(range.start, included(&as_key!(0x00)));
			assert_eq!(range.end, included(&as_key!(0xff)));
		}

		#[test]
		fn test_identical_bounds() {
			let range = EncodedKeyRange::start_end(Some(as_key!(0x42)), Some(as_key!(0x42)));
			assert_eq!(range.start, included(&as_key!(0x42)));
			assert_eq!(range.end, included(&as_key!(0x42)));
		}
	}

	mod all {
		use std::ops::Bound;

		use crate::value::encoded::key::EncodedKeyRange;

		#[test]
		fn test_is_unbounded() {
			let range = EncodedKeyRange::all();
			assert_eq!(range.start, Bound::Unbounded);
			assert_eq!(range.end, Bound::Unbounded);
		}
	}

	mod parse {
		use std::ops::Bound;

		use crate::value::encoded::key::{
			EncodedKey, EncodedKeyRange,
			tests::{excluded, included},
		};

		#[test]
		fn test_full_range() {
			let r = EncodedKeyRange::parse("a..z");
			assert_eq!(r.start, included(b"a"));
			assert_eq!(r.end, excluded(b"z"));
		}

		#[test]
		fn test_inclusive_end() {
			let r = EncodedKeyRange::parse("a..=z");
			assert_eq!(r.start, included(b"a"));
			assert_eq!(r.end, included(b"z"));
		}

		#[test]
		fn test_unbounded_start() {
			let r = EncodedKeyRange::parse("..z");
			assert_eq!(r.start, Bound::Unbounded);
			assert_eq!(r.end, excluded(b"z"));
		}

		#[test]
		fn test_unbounded_end() {
			let r = EncodedKeyRange::parse("a..");
			assert_eq!(r.start, included(b"a"));
			assert_eq!(r.end, Bound::Unbounded);
		}

		#[test]
		fn test_inclusive_only() {
			let r = EncodedKeyRange::parse("..=z");
			assert_eq!(r.start, Bound::Unbounded);
			assert_eq!(r.end, included(b"z"));
		}

		#[test]
		fn test_invalid_string_returns_degenerate_range() {
			let r = EncodedKeyRange::parse("not a range");
			let expected = EncodedKey::new([0xff]);
			assert_eq!(r.start, Bound::Included(expected.clone()));
			assert_eq!(r.end, Bound::Excluded(expected));
		}

		#[test]
		fn test_empty_string_returns_degenerate_range() {
			let r = EncodedKeyRange::parse("");
			let expected = EncodedKey::new([0xff]);
			assert_eq!(r.start, Bound::Included(expected.clone()));
			assert_eq!(r.end, Bound::Excluded(expected));
		}

		#[test]
		fn test_binary_encoded_values() {
			let r = EncodedKeyRange::parse("0101..=0aff");
			// decode_binary("0101") = [0x01, 0x01]
			assert_eq!(r.start, included(b"0101"));
			// decode_binary("0aff") = [0x0a, 0xff]
			assert_eq!(r.end, included(b"0aff"));
		}
	}

	fn included(key: &[u8]) -> Bound<EncodedKey> {
		Bound::Included(EncodedKey::new(key))
	}

	fn excluded(key: &[u8]) -> Bound<EncodedKey> {
		Bound::Excluded(EncodedKey::new(key))
	}
}