reifydb_core/value/index/
layout.rs

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