Skip to main content

reifydb_core/key/
index.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::collections::Bound;
5
6use super::{EncodableKey, EncodableKeyRange, KeyKind};
7use crate::{
8	encoded::key::{EncodedKey, EncodedKeyRange},
9	interface::catalog::{
10		id::{IndexId, PrimaryKeyId},
11		shape::ShapeId,
12	},
13	util::encoding::keycode::{deserializer::KeyDeserializer, serializer::KeySerializer},
14};
15
16#[derive(Debug, Clone, PartialEq)]
17pub struct IndexKey {
18	pub shape: ShapeId,
19	pub index: IndexId,
20}
21
22impl EncodableKey for IndexKey {
23	const KIND: KeyKind = KeyKind::Index;
24
25	fn encode(&self) -> EncodedKey {
26		let mut serializer = KeySerializer::with_capacity(18);
27		serializer.extend_u8(Self::KIND as u8).extend_shape_id(self.shape).extend_u64(self.index);
28		serializer.to_encoded_key()
29	}
30
31	fn decode(key: &EncodedKey) -> Option<Self> {
32		let mut de = KeyDeserializer::from_bytes(key.as_slice());
33
34		let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
35		if kind != Self::KIND {
36			return None;
37		}
38
39		let shape = de.read_shape_id().ok()?;
40		let index_value = de.read_u64().ok()?;
41
42		Some(Self {
43			shape,
44			index: IndexId::Primary(PrimaryKeyId(index_value)),
45		})
46	}
47}
48
49#[derive(Debug, Clone, PartialEq)]
50pub struct ShapeIndexKeyRange {
51	pub shape: ShapeId,
52}
53
54impl ShapeIndexKeyRange {
55	fn decode_key(key: &EncodedKey) -> Option<Self> {
56		let mut de = KeyDeserializer::from_bytes(key.as_slice());
57
58		let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
59		if kind != Self::KIND {
60			return None;
61		}
62
63		let shape = de.read_shape_id().ok()?;
64
65		Some(ShapeIndexKeyRange {
66			shape,
67		})
68	}
69}
70
71impl EncodableKeyRange for ShapeIndexKeyRange {
72	const KIND: KeyKind = KeyKind::Index;
73
74	fn start(&self) -> Option<EncodedKey> {
75		let mut serializer = KeySerializer::with_capacity(10);
76		serializer.extend_u8(Self::KIND as u8).extend_shape_id(self.shape);
77		Some(serializer.to_encoded_key())
78	}
79
80	fn end(&self) -> Option<EncodedKey> {
81		let mut serializer = KeySerializer::with_capacity(10);
82		serializer.extend_u8(Self::KIND as u8).extend_shape_id(self.shape.prev());
83		Some(serializer.to_encoded_key())
84	}
85
86	fn decode(range: &EncodedKeyRange) -> (Option<Self>, Option<Self>)
87	where
88		Self: Sized,
89	{
90		let start_key = match &range.start {
91			Bound::Included(key) | Bound::Excluded(key) => Self::decode_key(key),
92			Bound::Unbounded => None,
93		};
94
95		let end_key = match &range.end {
96			Bound::Included(key) | Bound::Excluded(key) => Self::decode_key(key),
97			Bound::Unbounded => None,
98		};
99
100		(start_key, end_key)
101	}
102}
103
104impl IndexKey {
105	pub fn encoded(shape: impl Into<ShapeId>, index: impl Into<IndexId>) -> EncodedKey {
106		Self {
107			shape: shape.into(),
108			index: index.into(),
109		}
110		.encode()
111	}
112
113	pub fn full_scan(shape: impl Into<ShapeId>) -> EncodedKeyRange {
114		let shape = shape.into();
115		EncodedKeyRange::start_end(Some(Self::shape_start(shape)), Some(Self::shape_end(shape)))
116	}
117
118	pub fn shape_start(shape: impl Into<ShapeId>) -> EncodedKey {
119		let shape = shape.into();
120		let mut serializer = KeySerializer::with_capacity(10);
121		serializer.extend_u8(Self::KIND as u8).extend_shape_id(shape);
122		serializer.to_encoded_key()
123	}
124
125	pub fn shape_end(shape: impl Into<ShapeId>) -> EncodedKey {
126		let shape = shape.into();
127		let mut serializer = KeySerializer::with_capacity(10);
128		serializer.extend_u8(Self::KIND as u8).extend_shape_id(shape.prev());
129		serializer.to_encoded_key()
130	}
131}
132
133#[cfg(test)]
134pub mod tests {
135	use super::{EncodableKey, IndexKey};
136	use crate::interface::catalog::{id::IndexId, shape::ShapeId};
137
138	#[test]
139	fn test_encode_decode() {
140		let key = IndexKey {
141			shape: ShapeId::table(0xABCD),
142			index: IndexId::primary(0x123456789ABCDEF0u64),
143		};
144		let encoded = key.encode();
145
146		let expected: Vec<u8> =
147			vec![0xF3, 0x01, 0x3F, 0x54, 0x32, 0x00, 0xED, 0xCB, 0xA9, 0x87, 0x65, 0x43, 0x21, 0x0F];
148
149		assert_eq!(encoded.as_slice(), expected);
150
151		let key = IndexKey::decode(&encoded).unwrap();
152		assert_eq!(key.shape, 0xABCD);
153		assert_eq!(key.index, 0x123456789ABCDEF0);
154	}
155
156	#[test]
157	fn test_order_preserving() {
158		let key1 = IndexKey {
159			shape: ShapeId::table(1),
160			index: IndexId::primary(100),
161		};
162		let key2 = IndexKey {
163			shape: ShapeId::table(1),
164			index: IndexId::primary(200),
165		};
166		let key3 = IndexKey {
167			shape: ShapeId::table(2),
168			index: IndexId::primary(50),
169		};
170
171		let encoded1 = key1.encode();
172		let encoded2 = key2.encode();
173		let encoded3 = key3.encode();
174
175		assert!(encoded3 < encoded2, "ordering not preserved");
176		assert!(encoded2 < encoded1, "ordering not preserved");
177	}
178}