reifydb_core/key/

index_entry.rs

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (c) 2025 ReifyDB

use std::collections::Bound;

use reifydb_type::util::cowvec::CowVec;

use super::{EncodableKey, EncodableKeyRange, KeyKind};
use crate::{
	encoded::key::{EncodedKey, EncodedKeyRange},
	interface::catalog::{id::IndexId, primitive::PrimitiveId},
	util::encoding::keycode::{deserializer::KeyDeserializer, serializer::KeySerializer},
	value::index::{encoded::EncodedIndexKey, range::EncodedIndexKeyRange},
};

const VERSION: u8 = 1;

/// Key for storing actual index entries with the encoded index key data
#[derive(Debug, Clone, PartialEq)]
pub struct IndexEntryKey {
	pub primitive: PrimitiveId,
	pub index: IndexId,
	pub key: EncodedIndexKey,
}

impl IndexEntryKey {
	pub fn new(primitive: impl Into<PrimitiveId>, index: IndexId, key: EncodedIndexKey) -> Self {
		Self {
			primitive: primitive.into(),
			index,
			key,
		}
	}

	pub fn encoded(primitive: impl Into<PrimitiveId>, index: IndexId, key: EncodedIndexKey) -> EncodedKey {
		Self::new(primitive, index, key).encode()
	}
}

#[derive(Debug, Clone, PartialEq)]
pub struct IndexEntryKeyRange {
	pub primitive: PrimitiveId,
	pub index: IndexId,
}

impl IndexEntryKeyRange {
	fn decode_key(key: &EncodedKey) -> Option<Self> {
		let mut de = KeyDeserializer::from_bytes(key.as_slice());

		let version = de.read_u8().ok()?;
		if version != VERSION {
			return None;
		}

		let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
		if kind != Self::KIND {
			return None;
		}

		let primitive = de.read_primitive_id().ok()?;
		let index = de.read_index_id().ok()?;

		Some(IndexEntryKeyRange {
			primitive,
			index,
		})
	}
}

impl EncodableKeyRange for IndexEntryKeyRange {
	const KIND: KeyKind = KeyKind::IndexEntry;

	fn start(&self) -> Option<EncodedKey> {
		let mut serializer = KeySerializer::with_capacity(20); // 1 + 1 + 9 + 9
		serializer
			.extend_u8(VERSION)
			.extend_u8(Self::KIND as u8)
			.extend_primitive_id(self.primitive)
			.extend_index_id(self.index);
		Some(serializer.to_encoded_key())
	}

	fn end(&self) -> Option<EncodedKey> {
		let mut serializer = KeySerializer::with_capacity(20);
		serializer
			.extend_u8(VERSION)
			.extend_u8(Self::KIND as u8)
			.extend_primitive_id(self.primitive)
			.extend_index_id(self.index.prev());
		Some(serializer.to_encoded_key())
	}

	fn decode(range: &EncodedKeyRange) -> (Option<Self>, Option<Self>)
	where
		Self: Sized,
	{
		let start_key = match &range.start {
			Bound::Included(key) | Bound::Excluded(key) => Self::decode_key(key),
			Bound::Unbounded => None,
		};

		let end_key = match &range.end {
			Bound::Included(key) | Bound::Excluded(key) => Self::decode_key(key),
			Bound::Unbounded => None,
		};

		(start_key, end_key)
	}
}

impl EncodableKey for IndexEntryKey {
	const KIND: KeyKind = KeyKind::IndexEntry;

	fn encode(&self) -> EncodedKey {
		let mut serializer = KeySerializer::with_capacity(20 + self.key.len());
		serializer
			.extend_u8(VERSION)
			.extend_u8(Self::KIND as u8)
			.extend_primitive_id(self.primitive)
			.extend_index_id(self.index)
			// Append the raw index key bytes
			.extend_raw(self.key.as_slice());
		serializer.to_encoded_key()
	}

	fn decode(key: &EncodedKey) -> Option<Self> {
		let mut de = KeyDeserializer::from_bytes(key.as_slice());

		let version = de.read_u8().ok()?;
		if version != VERSION {
			return None;
		}

		let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
		if kind != Self::KIND {
			return None;
		}

		let primitive = de.read_primitive_id().ok()?;
		let index = de.read_index_id().ok()?;

		// The remaining bytes are the index key
		let remaining = de.remaining();
		if remaining > 0 {
			let remaining_bytes = de.read_raw(remaining).ok()?;
			let index_key = EncodedIndexKey(CowVec::new(remaining_bytes.to_vec()));
			Some(Self {
				primitive,
				index,
				key: index_key,
			})
		} else {
			None
		}
	}
}

impl IndexEntryKey {
	/// Create a range for scanning all entries of a specific index
	pub fn index_range(primitive: impl Into<PrimitiveId>, index: IndexId) -> EncodedKeyRange {
		let range = IndexEntryKeyRange {
			primitive: primitive.into(),
			index,
		};
		EncodedKeyRange::new(Bound::Included(range.start().unwrap()), Bound::Excluded(range.end().unwrap()))
	}

	/// Create a range for scanning all entries of a primitive (all indexes)
	pub fn primitive_range(primitive: impl Into<PrimitiveId>) -> EncodedKeyRange {
		let primitive = primitive.into();
		let mut start_serializer = KeySerializer::with_capacity(11);
		start_serializer.extend_u8(VERSION).extend_u8(KeyKind::IndexEntry as u8).extend_primitive_id(primitive);

		let next_primitive = primitive.next();
		let mut end_serializer = KeySerializer::with_capacity(11);
		end_serializer
			.extend_u8(VERSION)
			.extend_u8(KeyKind::IndexEntry as u8)
			.extend_primitive_id(next_primitive);

		EncodedKeyRange {
			start: Bound::Included(start_serializer.to_encoded_key()),
			end: Bound::Excluded(end_serializer.to_encoded_key()),
		}
	}

	/// Create a range for scanning entries within an index with a specific
	/// key prefix
	pub fn key_prefix_range(
		primitive: impl Into<PrimitiveId>,
		index: IndexId,
		key_prefix: &[u8],
	) -> EncodedKeyRange {
		let primitive = primitive.into();
		let mut serializer = KeySerializer::with_capacity(20 + key_prefix.len());
		serializer
			.extend_u8(VERSION)
			.extend_u8(KeyKind::IndexEntry as u8)
			.extend_primitive_id(primitive)
			.extend_index_id(index)
			.extend_raw(key_prefix);
		let start = serializer.to_encoded_key();

		// For the end key, append 0xFF to get all keys with this prefix
		let mut end = start.as_slice().to_vec();
		end.push(0xFF);

		EncodedKeyRange {
			start: Bound::Included(start),
			end: Bound::Excluded(EncodedKey::new(end)),
		}
	}

	/// Create a range for entries from an EncodedIndexKeyRange
	/// This method leverages the EncodedIndexKeyRange type for cleaner
	/// range handling.
	pub fn key_range(
		primitive: impl Into<PrimitiveId>,
		index: IndexId,
		index_range: EncodedIndexKeyRange,
	) -> EncodedKeyRange {
		let primitive = primitive.into();
		// Build the prefix for this primitive and index
		let mut prefix_serializer = KeySerializer::with_capacity(20);
		prefix_serializer
			.extend_u8(VERSION)
			.extend_u8(KeyKind::IndexEntry as u8)
			.extend_primitive_id(primitive)
			.extend_index_id(index);
		let prefix = prefix_serializer.to_encoded_key().to_vec();

		// Convert bounds to include the prefix
		let start = match index_range.start {
			Bound::Included(key) => {
				let mut bytes = prefix.clone();
				bytes.extend_from_slice(key.as_slice());
				Bound::Included(EncodedKey::new(bytes))
			}
			Bound::Excluded(key) => {
				let mut bytes = prefix.clone();
				bytes.extend_from_slice(key.as_slice());
				Bound::Excluded(EncodedKey::new(bytes))
			}
			Bound::Unbounded => {
				// Start from the beginning of this index
				Bound::Included(EncodedKey::new(prefix.clone()))
			}
		};

		let end = match index_range.end {
			Bound::Included(key) => {
				let mut bytes = prefix.clone();
				bytes.extend_from_slice(key.as_slice());
				Bound::Included(EncodedKey::new(bytes))
			}
			Bound::Excluded(key) => {
				let mut bytes = prefix.clone();
				bytes.extend_from_slice(key.as_slice());
				Bound::Excluded(EncodedKey::new(bytes))
			}
			Bound::Unbounded => {
				// End at the beginning of the next index
				let mut serializer = KeySerializer::with_capacity(20);
				serializer
					.extend_u8(VERSION)
					.extend_u8(KeyKind::IndexEntry as u8)
					.extend_primitive_id(primitive)
					// Use prev() for end bound in descending order
					.extend_index_id(index.prev());
				Bound::Excluded(serializer.to_encoded_key())
			}
		};

		EncodedKeyRange {
			start,
			end,
		}
	}
}

#[cfg(test)]
pub mod tests {
	use reifydb_type::value::r#type::Type;

	use super::*;
	use crate::{sort::SortDirection, value::index::layout::EncodedIndexLayout};

	#[test]
	fn test_encode_decode() {
		// Create a simple index key
		let layout =
			EncodedIndexLayout::new(&[Type::Uint8, Type::Uint8], &[SortDirection::Asc, SortDirection::Asc])
				.unwrap();

		let mut index_key = layout.allocate_key();
		layout.set_u64(&mut index_key, 0, 100u64);
		layout.set_row_number(&mut index_key, 1, 1u64);

		let entry = IndexEntryKey {
			primitive: PrimitiveId::table(42),
			index: IndexId::primary(7),
			key: index_key.clone(),
		};

		let encoded = entry.encode();
		let decoded = IndexEntryKey::decode(&encoded).unwrap();

		assert_eq!(decoded.primitive, PrimitiveId::table(42));
		assert_eq!(decoded.index, IndexId::primary(7));
		assert_eq!(decoded.key.as_slice(), index_key.as_slice());
	}

	#[test]
	fn test_ordering() {
		let layout = EncodedIndexLayout::new(&[Type::Uint8], &[SortDirection::Asc]).unwrap();

		let mut key1 = layout.allocate_key();
		layout.set_u64(&mut key1, 0, 100u64);

		let mut key2 = layout.allocate_key();
		layout.set_u64(&mut key2, 0, 200u64);

		// Same primitive and index, different keys
		let entry1 = IndexEntryKey {
			primitive: PrimitiveId::table(1),
			index: IndexId::primary(1),
			key: key1,
		};

		let entry2 = IndexEntryKey {
			primitive: PrimitiveId::table(1),
			index: IndexId::primary(1),
			key: key2,
		};

		let encoded1 = entry1.encode();
		let encoded2 = entry2.encode();

		// entry1 should come before entry2 because 100 < 200
		assert!(encoded1.as_slice() < encoded2.as_slice());
	}

	#[test]
	fn test_index_range() {
		let range = IndexEntryKey::index_range(PrimitiveId::table(10), IndexId::primary(5));

		// Create entries that should be included
		let layout = EncodedIndexLayout::new(&[Type::Uint8], &[SortDirection::Asc]).unwrap();

		let mut key = layout.allocate_key();
		layout.set_u64(&mut key, 0, 50u64);

		let entry = IndexEntryKey {
			primitive: PrimitiveId::table(10),
			index: IndexId::primary(5),
			key,
		};

		let encoded = entry.encode();

		// Check that the entry falls within the range
		if let (Bound::Included(start), Bound::Excluded(end)) = (&range.start, &range.end) {
			assert!(encoded.as_slice() >= start.as_slice());
			assert!(encoded.as_slice() < end.as_slice());
		} else {
			panic!("Expected Included/Excluded bounds");
		}

		// Entry with different index should not be in range
		// Note: Due to keycode encoding, IndexId(6) will have a smaller
		// encoded value than IndexId(5) since keycode inverts bits
		// (larger numbers become smaller byte sequences)
		let entry2 = IndexEntryKey {
			primitive: PrimitiveId::table(10),
			index: IndexId::primary(6),
			key: layout.allocate_key(),
		};

		let encoded2 = entry2.encode();
		// The entry with IndexId(6) should not be within the range for
		// IndexId(5)
		if let (Bound::Included(start), Bound::Excluded(end)) = (&range.start, &range.end) {
			// encoded2 should either be < start or >= end
			assert!(encoded2.as_slice() < start.as_slice() || encoded2.as_slice() >= end.as_slice());
		}
	}

	#[test]
	fn test_key_prefix_range() {
		let layout =
			EncodedIndexLayout::new(&[Type::Uint8, Type::Uint8], &[SortDirection::Asc, SortDirection::Asc])
				.unwrap();

		let mut key = layout.allocate_key();
		layout.set_u64(&mut key, 0, 100u64);
		layout.set_row_number(&mut key, 1, 0u64); // Set to 0 to get the minimal key with this prefix

		// Use the full encoded key up to the first field as the prefix
		let prefix = &key.as_slice()[..layout.fields[1].offset]; // Include bitvec and first field
		let range = IndexEntryKey::key_prefix_range(PrimitiveId::table(1), IndexId::primary(1), prefix);

		// Now create a full key with the same prefix
		layout.set_row_number(&mut key, 1, 999u64);
		let entry = IndexEntryKey {
			primitive: PrimitiveId::table(1),
			index: IndexId::primary(1),
			key: key.clone(),
		};

		let encoded = entry.encode();

		// Should be within range
		if let (Bound::Included(start), Bound::Excluded(end)) = (&range.start, &range.end) {
			assert!(encoded.as_slice() >= start.as_slice());
			assert!(encoded.as_slice() < end.as_slice());
		}

		// Create a key with different prefix
		let mut key2 = layout.allocate_key();
		layout.set_u64(&mut key2, 0, 200u64); // Different first field
		layout.set_row_number(&mut key2, 1, 1u64);

		let entry2 = IndexEntryKey {
			primitive: PrimitiveId::table(1),
			index: IndexId::primary(1),
			key: key2,
		};

		let encoded2 = entry2.encode();

		// Should not be in range
		if let Bound::Excluded(end) = &range.end {
			assert!(encoded2.as_slice() >= end.as_slice());
		}
	}
}