Skip to main content

reifydb_core/value/index/
shape.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{alloc, ops::Deref, sync::Arc};
5
6use reifydb_value::{Result, value::value_type::ValueType};
7
8use crate::{error::CoreError, sort::SortDirection, value::index::encoded::EncodedIndexKey};
9
10#[derive(Debug, Clone)]
11pub struct IndexShape(Arc<IndexShapeInner>);
12
13impl Deref for IndexShape {
14	type Target = IndexShapeInner;
15
16	fn deref(&self) -> &Self::Target {
17		&self.0
18	}
19}
20
21impl IndexShape {
22	pub fn new(types: &[ValueType], directions: &[SortDirection]) -> Result<Self> {
23		if types.len() != directions.len() {
24			return Err(CoreError::IndexTypesDirectionsMismatch {
25				types_len: types.len(),
26				directions_len: directions.len(),
27			}
28			.into());
29		}
30
31		for typ in types {
32			if matches!(typ, ValueType::Utf8 | ValueType::Blob | ValueType::Digest { .. }) {
33				return Err(CoreError::IndexVariableLengthNotSupported.into());
34			}
35		}
36
37		Ok(Self(Arc::new(IndexShapeInner::new(types, directions))))
38	}
39}
40
41#[derive(Debug)]
42pub struct IndexShapeInner {
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: ValueType,
55	pub direction: SortDirection,
56}
57
58impl IndexShapeInner {
59	fn new(types: &[ValueType], 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.div_ceil(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		IndexShapeInner {
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 = alloc::Layout::from_size_align(self.total_size, self.alignment).unwrap();
99		unsafe {
100			let ptr = alloc::alloc_zeroed(layout);
101			if ptr.is_null() {
102				alloc::handle_alloc_error(layout);
103			}
104			let vec = Vec::from_raw_parts(ptr, self.total_size, self.total_size);
105			EncodedIndexKey::new(vec)
106		}
107	}
108
109	pub const fn data_offset(&self) -> usize {
110		self.bitvec_size
111	}
112
113	pub fn value(&self, index: usize) -> ValueType {
114		self.fields[index].value.clone()
115	}
116
117	pub fn direction(&self, index: usize) -> &SortDirection {
118		&self.fields[index].direction
119	}
120}
121
122fn align_up(offset: usize, align: usize) -> usize {
123	(offset + align - 1) & !(align - 1)
124}
125
126#[cfg(test)]
127pub mod tests {
128	use super::*;
129	use crate::sort::SortDirection;
130
131	#[test]
132	fn test_single_field_int() {
133		let layout = IndexShape::new(&[ValueType::Int4], &[SortDirection::Asc]).unwrap();
134
135		assert_eq!(layout.bitvec_size, 1);
136		assert_eq!(layout.fields.len(), 1);
137		assert_eq!(layout.fields[0].offset, 4); // aligned to 4 bytes
138		assert_eq!(layout.fields[0].value, ValueType::Int4);
139		assert_eq!(layout.fields[0].direction, SortDirection::Asc);
140		assert_eq!(layout.alignment, 4);
141		assert_eq!(layout.total_size, 8); // 1 bitvec + 3 padding + 4 int
142	}
143
144	#[test]
145	fn test_multiple_fields_mixed_directions() {
146		let layout = IndexShape::new(
147			&[ValueType::Int4, ValueType::Int8, ValueType::Uint8],
148			&[SortDirection::Desc, SortDirection::Asc, SortDirection::Asc],
149		)
150		.unwrap();
151
152		assert_eq!(layout.bitvec_size, 1);
153		assert_eq!(layout.fields.len(), 3);
154
155		assert_eq!(layout.fields[0].value, ValueType::Int4);
156		assert_eq!(layout.fields[0].direction, SortDirection::Desc);
157
158		assert_eq!(layout.fields[1].value, ValueType::Int8);
159		assert_eq!(layout.fields[1].direction, SortDirection::Asc);
160
161		assert_eq!(layout.fields[2].value, ValueType::Uint8);
162		assert_eq!(layout.fields[2].direction, SortDirection::Asc);
163
164		assert_eq!(layout.alignment, 8);
165	}
166
167	#[test]
168	fn test_reject_variable_length_types() {
169		let result =
170			IndexShape::new(&[ValueType::Int4, ValueType::Utf8], &[SortDirection::Asc, SortDirection::Asc]);
171		assert!(result.is_err());
172
173		let result = IndexShape::new(&[ValueType::Blob], &[SortDirection::Desc]);
174		assert!(result.is_err());
175	}
176
177	#[test]
178	fn test_allocate_key() {
179		let layout = IndexShape::new(
180			&[ValueType::Boolean, ValueType::Int4],
181			&[SortDirection::Asc, SortDirection::Desc],
182		)
183		.unwrap();
184
185		let key = layout.allocate_key();
186		assert_eq!(key.len(), layout.total_size);
187
188		for byte in key.as_slice() {
189			assert_eq!(*byte, 0);
190		}
191	}
192}