reifydb_core/value/index/
shape.rs1use 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); 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); }
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}