Skip to main content

reifydb_core/value/index/
layout.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2025 ReifyDB
3
4use std::{ops::Deref, sync::Arc};
5
6use reifydb_type::{error, util::cowvec::CowVec, value::r#type::Type};
7
8use crate::{
9	error::diagnostic::catalog::{index_types_directions_mismatch, index_variable_length_not_supported},
10	sort::SortDirection,
11	value::index::encoded::EncodedIndexKey,
12};
13
14#[derive(Debug, Clone)]
15pub struct EncodedIndexLayout(Arc<EncodedIndexLayoutInner>);
16
17impl Deref for EncodedIndexLayout {
18	type Target = EncodedIndexLayoutInner;
19
20	fn deref(&self) -> &Self::Target {
21		&self.0
22	}
23}
24
25impl EncodedIndexLayout {
26	pub fn new(types: &[Type], directions: &[SortDirection]) -> reifydb_type::Result<Self> {
27		if types.len() != directions.len() {
28			return Err(error!(index_types_directions_mismatch(types.len(), directions.len())));
29		}
30
31		for typ in types {
32			if matches!(typ, Type::Utf8 | Type::Blob) {
33				return Err(error!(index_variable_length_not_supported()));
34			}
35		}
36
37		Ok(Self(Arc::new(EncodedIndexLayoutInner::new(types, directions))))
38	}
39}
40
41#[derive(Debug)]
42pub struct EncodedIndexLayoutInner {
43	pub fields: Vec<IndexField>,
44	pub total_size: usize,
45	pub bitvec_size: usize,
46	pub alignment: usize,
47}
48
49#[derive(Debug)]
50pub struct IndexField {
51	pub offset: usize,
52	pub size: usize,
53	pub align: usize,
54	pub value: Type,
55	pub direction: SortDirection,
56}
57
58impl EncodedIndexLayoutInner {
59	fn new(types: &[Type], directions: &[SortDirection]) -> Self {
60		assert!(!types.is_empty());
61		assert_eq!(types.len(), directions.len());
62
63		let num_fields = types.len();
64		let bitvec_bytes = (num_fields + 7) / 8;
65
66		let mut offset = bitvec_bytes;
67		let mut fields = Vec::with_capacity(num_fields);
68		let mut max_align = 1;
69
70		for (i, value) in types.iter().enumerate() {
71			let size = value.size();
72			let align = value.alignment();
73
74			offset = align_up(offset, align);
75			fields.push(IndexField {
76				offset,
77				size,
78				align,
79				value: value.clone(),
80				direction: directions[i].clone(),
81			});
82
83			offset += size;
84			max_align = max_align.max(align);
85		}
86
87		let total_size = align_up(offset, max_align);
88
89		EncodedIndexLayoutInner {
90			fields,
91			total_size,
92			alignment: max_align,
93			bitvec_size: bitvec_bytes,
94		}
95	}
96
97	pub fn allocate_key(&self) -> EncodedIndexKey {
98		let layout = std::alloc::Layout::from_size_align(self.total_size, self.alignment).unwrap();
99		unsafe {
100			let ptr = std::alloc::alloc_zeroed(layout);
101			if ptr.is_null() {
102				std::alloc::handle_alloc_error(layout);
103			}
104			let vec = Vec::from_raw_parts(ptr, self.total_size, self.total_size);
105			EncodedIndexKey(CowVec::new(vec))
106		}
107	}
108
109	pub const fn data_offset(&self) -> usize {
110		self.bitvec_size
111	}
112
113	pub fn all_defined(&self, key: &EncodedIndexKey) -> bool {
114		let bits = self.fields.len();
115		if bits == 0 {
116			return false;
117		}
118
119		let bitvec_slice = &key[..self.bitvec_size];
120		for (i, &byte) in bitvec_slice.iter().enumerate() {
121			let bits_in_byte = if i == self.bitvec_size - 1 && bits % 8 != 0 {
122				bits % 8
123			} else {
124				8
125			};
126
127			let mask = if bits_in_byte == 8 {
128				0xFF
129			} else {
130				(1u8 << bits_in_byte) - 1
131			};
132			if (byte & mask) != mask {
133				return false;
134			}
135		}
136
137		true
138	}
139
140	pub fn value(&self, index: usize) -> Type {
141		self.fields[index].value.clone()
142	}
143
144	pub fn direction(&self, index: usize) -> &SortDirection {
145		&self.fields[index].direction
146	}
147}
148
149fn align_up(offset: usize, align: usize) -> usize {
150	(offset + align - 1) & !(align - 1)
151}
152
153#[cfg(test)]
154pub mod tests {
155	use super::*;
156	use crate::sort::SortDirection;
157
158	#[test]
159	fn test_single_field_int() {
160		let layout = EncodedIndexLayout::new(&[Type::Int4], &[SortDirection::Asc]).unwrap();
161
162		assert_eq!(layout.bitvec_size, 1);
163		assert_eq!(layout.fields.len(), 1);
164		assert_eq!(layout.fields[0].offset, 4); // aligned to 4 bytes
165		assert_eq!(layout.fields[0].value, Type::Int4);
166		assert_eq!(layout.fields[0].direction, SortDirection::Asc);
167		assert_eq!(layout.alignment, 4);
168		assert_eq!(layout.total_size, 8); // 1 bitvec + 3 padding + 4 int
169	}
170
171	#[test]
172	fn test_multiple_fields_mixed_directions() {
173		let layout = EncodedIndexLayout::new(
174			&[Type::Int4, Type::Int8, Type::Uint8],
175			&[SortDirection::Desc, SortDirection::Asc, SortDirection::Asc],
176		)
177		.unwrap();
178
179		assert_eq!(layout.bitvec_size, 1);
180		assert_eq!(layout.fields.len(), 3);
181
182		assert_eq!(layout.fields[0].value, Type::Int4);
183		assert_eq!(layout.fields[0].direction, SortDirection::Desc);
184
185		assert_eq!(layout.fields[1].value, Type::Int8);
186		assert_eq!(layout.fields[1].direction, SortDirection::Asc);
187
188		assert_eq!(layout.fields[2].value, Type::Uint8);
189		assert_eq!(layout.fields[2].direction, SortDirection::Asc);
190
191		assert_eq!(layout.alignment, 8);
192	}
193
194	#[test]
195	fn test_reject_variable_length_types() {
196		let result =
197			EncodedIndexLayout::new(&[Type::Int4, Type::Utf8], &[SortDirection::Asc, SortDirection::Asc]);
198		assert!(result.is_err());
199
200		let result = EncodedIndexLayout::new(&[Type::Blob], &[SortDirection::Desc]);
201		assert!(result.is_err());
202	}
203
204	#[test]
205	fn test_allocate_key() {
206		let layout = EncodedIndexLayout::new(
207			&[Type::Boolean, Type::Int4],
208			&[SortDirection::Asc, SortDirection::Desc],
209		)
210		.unwrap();
211
212		let key = layout.allocate_key();
213		assert_eq!(key.len(), layout.total_size);
214
215		for byte in key.as_slice() {
216			assert_eq!(*byte, 0);
217		}
218	}
219}